rewritetern.go

     1  // Copyright 2025 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package ssacompile
     6  
     7  import (
     8  	"internal/goarch"
     9  	"slices"
    10  
    11  	"cmd/compile/internal/ssa"
    12  )
    13  
    14  var truthTableValues [3]uint8 = [3]uint8{0b1111_0000, 0b1100_1100, 0b1010_1010}
    15  
    16  func (slop SIMDLogicalOP) String() string {
    17  	if slop == sloInterior {
    18  		return "leaf"
    19  	}
    20  	interior := ""
    21  	if slop&sloInterior != 0 {
    22  		interior = "+interior"
    23  	}
    24  	switch slop &^ sloInterior {
    25  	case sloAnd:
    26  		return "and" + interior
    27  	case sloXor:
    28  		return "xor" + interior
    29  	case sloOr:
    30  		return "or" + interior
    31  	case sloAndNot:
    32  		return "andNot" + interior
    33  	case sloNot:
    34  		return "not" + interior
    35  	}
    36  	return "wrong"
    37  }
    38  
    39  func rewriteTern(f *ssa.Func) {
    40  	if f.MaxCPUFeatures == ssa.CPUNone {
    41  		return
    42  	}
    43  
    44  	arch := f.Config.Ctxt.Arch.Family
    45  	// TODO there are other SIMD architectures
    46  	if arch != goarch.AMD64 {
    47  		return
    48  	}
    49  
    50  	boolExprTrees := make(map[*ssa.Value]SIMDLogicalOP)
    51  
    52  	// Find logical-expr expression trees, including leaves.
    53  	// interior nodes will be marked sloInterior,
    54  	// root nodes will not be marked sloInterior,
    55  	// leaf nodes are only marked sloInterior.
    56  	for _, b := range f.Blocks {
    57  		for _, v := range b.Values {
    58  			slo := classifyBooleanSIMD(v)
    59  			switch slo {
    60  			case sloOr,
    61  				sloAndNot,
    62  				sloXor,
    63  				sloAnd:
    64  				boolExprTrees[v.Args[1]] |= sloInterior
    65  				fallthrough
    66  			case sloNot:
    67  				boolExprTrees[v.Args[0]] |= sloInterior
    68  				boolExprTrees[v] |= slo
    69  			}
    70  		}
    71  	}
    72  
    73  	// get a canonical sorted set of roots
    74  	var roots []*ssa.Value
    75  	for v, slo := range boolExprTrees {
    76  		if f.Pass.Debug > 1 {
    77  			f.Warnl(v.Pos, "%s has SLO %v", v.LongString(), slo)
    78  		}
    79  
    80  		if slo&sloInterior == 0 && v.Block.CPUfeatures.HasFeature(ssa.CPUavx512) {
    81  			roots = append(roots, v)
    82  		}
    83  	}
    84  	slices.SortFunc(roots, func(u, v *ssa.Value) int { return int(u.ID - v.ID) }) // IDs are small enough to not care about overflow.
    85  
    86  	// This rewrite works by iterating over the root set.
    87  	// For each boolean expression, it walks the expression
    88  	// bottom up accumulating sets of variables mentioned in
    89  	// subexpressions, lazy-greedily finding the largest subexpressions
    90  	// of 3 inputs that can be rewritten to use ternary-truth-table instructions.
    91  
    92  	// rewrite recursively attempts to replace v and v's subexpressions with
    93  	// ternary-logic truth-table operations, returning a set of not more than 3
    94  	// subexpressions within v that may be combined into a parent's replacement.
    95  	// V need not have the CPU features that allow a ternary-logic operation;
    96  	// in that case, v will not be rewritten.  Replacements also require
    97  	// exactly 3 different variable inputs to a boolean expression.
    98  	//
    99  	// Given the CPU feature and 3 inputs, v is replaced in the following
   100  	// cases:
   101  	//
   102  	// 1) v is a root
   103  	// 2) u = NOT(v) and u lacks the CPU feature
   104  	// 3) u = OP(v, w) and u lacks the CPU feature
   105  	// 4) u = OP(v, w) and u has more than 3 variable inputs.	var rewrite func(v *Value) [3]*Value
   106  	var rewrite func(v *ssa.Value) [3]*ssa.Value
   107  
   108  	// computeTT returns the truth table for a boolean expression
   109  	// over the variables in vars, where vars[0] varies slowest in
   110  	// the truth table and vars[2] varies fastest.
   111  	// e.g. computeTT( "and(x, or(y, not(z)))", {x,y,z} ) returns
   112  	// (bit 0 first) 0 0 0 0 1 0 1 1 = (reversed) 1101_0000 = 0xD0
   113  	//            x: 0 0 0 0 1 1 1 1
   114  	//            y: 0 0 1 1 0 0 1 1
   115  	//            z: 0 1 0 1 0 1 0 1
   116  	var computeTT func(v *ssa.Value, vars [3]*ssa.Value) uint8
   117  
   118  	// combine two sets of variables into one, returning ok/not
   119  	// if the two sets contained 3 or fewer elements.  Combine
   120  	// ensures that the sets of Values never contain duplicates.
   121  	// (Duplicates would create less-efficient code, not incorrect code.)
   122  	combine := func(a, b [3]*ssa.Value) ([3]*ssa.Value, bool) {
   123  		var c [3]*ssa.Value
   124  		i := 0
   125  		for _, v := range a {
   126  			if v == nil {
   127  				break
   128  			}
   129  			c[i] = v
   130  			i++
   131  		}
   132  	bloop:
   133  		for _, v := range b {
   134  			if v == nil {
   135  				break
   136  			}
   137  			for _, u := range a {
   138  				if v == u {
   139  					continue bloop
   140  				}
   141  			}
   142  			if i == 3 {
   143  				return [3]*ssa.Value{}, false
   144  			}
   145  			c[i] = v
   146  			i++
   147  		}
   148  		return c, true
   149  	}
   150  
   151  	computeTT = func(v *ssa.Value, vars [3]*ssa.Value) uint8 {
   152  		i := 0
   153  		for ; i < len(vars); i++ {
   154  			if vars[i] == v {
   155  				return truthTableValues[i]
   156  			}
   157  		}
   158  		slo := boolExprTrees[v] &^ sloInterior
   159  		a := computeTT(v.Args[0], vars)
   160  		switch slo {
   161  		case sloNot:
   162  			return ^a
   163  		case sloAnd:
   164  			return a & computeTT(v.Args[1], vars)
   165  		case sloXor:
   166  			return a ^ computeTT(v.Args[1], vars)
   167  		case sloOr:
   168  			return a | computeTT(v.Args[1], vars)
   169  		case sloAndNot:
   170  			return a & ^computeTT(v.Args[1], vars)
   171  		}
   172  		panic("switch should have covered all cases, or unknown var in logical expression")
   173  	}
   174  
   175  	replace := func(a0 *ssa.Value, vars0 [3]*ssa.Value) {
   176  		imm := computeTT(a0, vars0)
   177  		op := ternOpForLogical(a0.Op)
   178  		if op == a0.Op {
   179  			if f.Pass.Debug > 0 {
   180  				f.Warnl(a0.Pos, "Skipping rewrite for %s, op=%v", a0.LongString(), op)
   181  			}
   182  			return
   183  		}
   184  		if f.Pass.Debug > 0 {
   185  			f.Warnl(a0.Pos, "Rewriting %s into %v of 0b%b %v %v %v", a0.LongString(), op, imm,
   186  				vars0[0], vars0[1], vars0[2])
   187  		}
   188  		a0.Reset(op)
   189  		a0.SetArgs3(vars0[0], vars0[1], vars0[2])
   190  		a0.AuxInt = int64(int8(imm))
   191  	}
   192  
   193  	// addOne ensures the no-duplicates addition of a single value
   194  	// to a set that is not full.  It seems possible that a shared
   195  	// subexpression in tricky combination with blocks lacking the
   196  	// AVX512 feature might permit this.
   197  	addOne := func(vars [3]*ssa.Value, v *ssa.Value) [3]*ssa.Value {
   198  		if vars[2] != nil {
   199  			panic("rewriteTern.addOne, vars[2] should be nil")
   200  		}
   201  		if v == vars[0] || v == vars[1] {
   202  			return vars
   203  		}
   204  		if vars[1] == nil {
   205  			vars[1] = v
   206  		} else {
   207  			vars[2] = v
   208  		}
   209  		return vars
   210  	}
   211  
   212  	rewrite = func(v *ssa.Value) [3]*ssa.Value {
   213  		slo := boolExprTrees[v]
   214  		if slo == sloInterior { // leaf node, i.e., a "variable"
   215  			return [3]*ssa.Value{v, nil, nil}
   216  		}
   217  		var vars [3]*ssa.Value
   218  		hasFeature := v.Block.CPUfeatures.HasFeature(ssa.CPUavx512)
   219  		if slo&sloNot == sloNot {
   220  			vars = rewrite(v.Args[0])
   221  			if !hasFeature {
   222  				if vars[2] != nil {
   223  					replace(v.Args[0], vars)
   224  					return [3]*ssa.Value{v, nil, nil}
   225  				}
   226  				return vars
   227  			}
   228  		} else {
   229  			var ok bool
   230  			a0, a1 := v.Args[0], v.Args[1]
   231  			vars0 := rewrite(a0)
   232  			vars1 := rewrite(a1)
   233  			vars, ok = combine(vars0, vars1)
   234  
   235  			if f.Pass.Debug > 1 {
   236  				f.Warnl(a0.Pos, "combine(%v, %v) -> %v, %v", vars0, vars1, vars, ok)
   237  			}
   238  
   239  			if !(ok && v.Block.CPUfeatures.HasFeature(ssa.CPUavx512)) {
   240  				// too many variables, or cannot rewrite current values.
   241  				// rewrite one or both subtrees if possible
   242  				if vars0[2] != nil && a0.Block.CPUfeatures.HasFeature(ssa.CPUavx512) {
   243  					replace(a0, vars0)
   244  				}
   245  				if vars1[2] != nil && a1.Block.CPUfeatures.HasFeature(ssa.CPUavx512) {
   246  					replace(a1, vars1)
   247  				}
   248  
   249  				// 3-element var arrays are either rewritten, or unable to be rewritten
   250  				// because of the features in effect in their block.  Either way, they
   251  				// are treated as a "new var" if 3 elements are present.
   252  
   253  				if vars0[2] == nil {
   254  					if vars1[2] == nil {
   255  						// both subtrees are 2-element and were not rewritten.
   256  						//
   257  						// TODO a clever person would look at subtrees of inputs,
   258  						// e.g. rewrite
   259  						//        ((a AND b) XOR b) XOR (d  XOR (c AND d))
   260  						// to    (((a AND b) XOR b) XOR  d) XOR (c AND d)
   261  						// to v = TERNLOG(truthtable, a, b, d) XOR (c AND d)
   262  						// and return the variable set {v, c, d}
   263  						//
   264  						// But for now, just restart with a0 and a1.
   265  						return [3]*ssa.Value{a0, a1, nil}
   266  					} else {
   267  						// a1 (maybe) rewrote, a0 has room for another var
   268  						vars = addOne(vars0, a1)
   269  					}
   270  				} else if vars1[2] == nil {
   271  					// a0 (maybe) rewrote, a1 has room for another var
   272  					vars = addOne(vars1, a0)
   273  				} else if !ok {
   274  					// both (maybe) rewrote
   275  					// a0 and a1 are different because otherwise their variable
   276  					// sets would have combined "ok".
   277  					return [3]*ssa.Value{a0, a1, nil}
   278  				}
   279  				// continue with either the vars from "ok" or the updated set of vars.
   280  			}
   281  		}
   282  		// if root and 3 vars and hasFeature, rewrite.
   283  		if slo&sloInterior == 0 && vars[2] != nil && hasFeature {
   284  			replace(v, vars)
   285  			return [3]*ssa.Value{v, nil, nil}
   286  		}
   287  		return vars
   288  	}
   289  
   290  	for _, v := range roots {
   291  		if f.Pass.Debug > 1 {
   292  			f.Warnl(v.Pos, "SLO root %s", v.LongString())
   293  		}
   294  		rewrite(v)
   295  	}
   296  }
   297
View as plain text