rewrite.go

     1  // Copyright 2026 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 midway
     6  
     7  import (
     8  	"cmd/compile/internal/base"
     9  	"cmd/compile/internal/syntax"
    10  	"cmd/compile/internal/types2"
    11  	"fmt"
    12  	"internal/buildcfg"
    13  	"internal/simd/variants"
    14  	"strings"
    15  )
    16  
    17  // "Midway" rewriting
    18  //
    19  // Go attempts to provide a package similar to the "Highway" library
    20  // for C++ (https://google.github.io/highway).  The library package is "simd"
    21  // and defines vector types with unspecified widths that are bound to particular
    22  // machine dependent types as late as program execution.  This is accomplished
    23  // by rewriting code that depends on these types into code that references
    24  // architecture-specific types, perhaps more than once, and if necessary
    25  // dynamically choosing which version to execute based on hardware attributes.
    26  //
    27  // The rewriting takes place early in the compiler, after type checking but
    28  // before conversion to "unified" IR.  To ensure that types are correctly set
    29  // on the modified version of the code, type checking information is reset and
    30  // the type checking phase is re-run.  This places some limits on the shape of
    31  // the rewrites, but it also ensures that the rewritten code is well-formed.
    32  //
    33  // Rewritten code does not reference "archsimd" types directly, but instead
    34  // references types in a "bridge" package that filters the available methods
    35  // and adds a few more.  The package used relies on a builder/compiler hack;
    36  // the compiler's type checker enforces export naming conventions, but the
    37  // build system limits visibility to unrelated "internal" packages and can be
    38  // modified to allow access in special cases (like this one).  This allows the
    39  // rewritten code to reference types, functions, and methods that are not
    40  // accessible otherwise.
    41  //
    42  // The rewrite works in phases.  The first is "analysis", to discover functions,
    43  // types, methods, and variables that depend on "simd" types.  "Depend on" means
    44  // any mention of a simd type, and for types, also includes types that have a
    45  // simd-dependent method.  Dependent functions are split into two categories;
    46  // those whose dependence includes their signature, and those that do not.
    47  // The second category forms the boundary between code that depends on simd and
    48  // code that does not.  Notice that there cannot be a boundary method, because
    49  // (by design) the receiver type is simd-dependent and thus a dependent method
    50  // also has a dependent type in its signature.
    51  //
    52  // The second phase rewrites such "boundary" functions into a "dispatch" version
    53  // and (later, third phase) "specialized" versions.  The dispatch function
    54  // will choose which specialized version to call based on which simd implementation
    55  // has been chosen, and forward parameters and results to/from that specialized version
    56  // of the function.  The dispatch version shares the same name as the original function.
    57  // Note that this applies to functions only, and not methods.
    58  
    59  // The third phase specializes dependent functions (both kinds), methods,
    60  // global variables, and types into size/emulation/feature-specific variants.
    61  // Except for methods, this is done by adding a suffix beginning with "@" to
    62  // the name.  Because "@" cannot appear in legal Go identifiers this removes
    63  // the risk of a naming overlap.  Methods are specialized, but not renamed,
    64  // because their receiver type is renamed instead.  Not changing method names
    65  // preserves interface satisfaction, for example in the case of generic interfaces.
    66  //
    67  // Non-boundary dependent function and methods are not rewritten into dispatch
    68  // functions/methods, but remain in the generated code because they must be
    69  // present in the export data so that other packages that import them will still
    70  // compile before rewriting.  Their bodies are replaced with panic(...) to allow
    71  // compilation while preventing even worse chaos in the event of a bug either in
    72  // the compiler or through ambitious use of reflection or assembly language.
    73  //
    74  
    75  /* Example rewrites
    76  
    77  // Type alias, global variable, and init function:
    78  
    79  // before:
    80  type MyInt8s = simd.Int8s
    81  func Generic[T haslen](x int) int {
    82      var v T
    83      return x + v.Len()
    84  }
    85  var VL int
    86  func init() {
    87      VL = Generic[MyInt8s](1)
    88  }
    89  // dispatch:
    90  func init() {
    91      switch simd.VectorBitSize() {
    92      case
    93          128:
    94              init@simd128()
    95              return
    96      case 256:
    97              init@simd256()
    98              return
    99      case 512:
   100              init@simd512()
   101              return
   102      default:
   103          panic("unsupported vector size")
   104      }
   105  }
   106  // specialized (128)
   107  type MyInt8s@simd128 = archsimd.Int8x16
   108  func init@simd128() {
   109          VL = Generic[MyInt8s@simd128](1)
   110  }
   111  
   112  
   113  // structure containing simd fields, and with simd methods
   114  
   115  // before
   116  // A struct dependent on SIMD
   117  type VectorC struct {
   118      Field simd.Float32s
   119  }
   120  func (v *VectorC) MethodOfSimd() bool {
   121      return false
   122  }
   123  func (v VectorC) Data() simd.Float32s {
   124      return v.Field
   125  }
   126  func (v VectorC) Foo(x VectorC) VectorC {
   127      return VectorC{Field: v.Field.Add(x.Field)}
   128  }
   129  
   130  // dispatch
   131  // technically there is none, but functions with panicking bodies
   132  // remain because code must pass type checking before rewriting.
   133  type VectorC struct {
   134      Field simd.Float32s
   135  }
   136  func (v *VectorC) MethodOfSimd() bool {
   137      panic(...)
   138  }
   139  func (v VectorC) Data() simd.Float32s {
   140      panic(...)
   141  }
   142  func (v VectorC) Foo(x VectorC) VectorC {
   143      panic(...)
   144  }
   145  
   146  // specialized (128)
   147  
   148  // A struct dependent on SIMD
   149  type VectorC@simd128 struct {
   150      Field bridge.Float32x4
   151  }
   152  func (v *VectorC@simd128) MethodOfSimd() bool {
   153      return false
   154  }
   155  func (v VectorC@simd128) Data() bridge.Float32x4 {
   156      return v.Field
   157  }
   158  func (v VectorC@simd128) Foo(x VectorC@simd128) VectorC@simd128 {
   159      return VectorC@simd128{Field: v.Field.Add(x.Field)}
   160  }
   161  
   162  */
   163  
   164  type Rewriter struct {
   165  	pkg      *types2.Package
   166  	analyzer *Analyzer
   167  	info     *types2.Info
   168  	sizes    []int
   169  }
   170  
   171  func NewRewriter(pkg *types2.Package, info *types2.Info, analyzer *Analyzer, sizes []int) *Rewriter {
   172  	return &Rewriter{
   173  		pkg:      pkg,
   174  		info:     info,
   175  		analyzer: analyzer,
   176  		sizes:    sizes,
   177  	}
   178  }
   179  
   180  func (r *Rewriter) Rewrite(files []*syntax.File) {
   181  
   182  	// First duplicate and specialize all dependent functions and variables.
   183  	for _, fileAST := range files {
   184  
   185  		var newDecls []syntax.Decl
   186  		for _, k := range r.sizes {
   187  			newDecls = r.generateForSize(fileAST, k, "", newDecls)
   188  			if v := variants.Variants[variants.Key{Arch: buildcfg.GOARCH, Size: k}]; v != nil {
   189  				newDecls = r.generateForSize(fileAST, k, v.Suffix, newDecls)
   190  			}
   191  		}
   192  
   193  		// Then replace original functions with dispatchers.
   194  		// This also edits the DeclList of fileAST.
   195  		r.generateDispatchers(fileAST)
   196  
   197  		fileAST.DeclList = append(fileAST.DeclList, newDecls...)
   198  	}
   199  }
   200  
   201  func (r *Rewriter) generateDispatchers(fileAST *syntax.File) {
   202  	var newDecls []syntax.Decl
   203  
   204  	change := false
   205  
   206  	for _, decl := range fileAST.DeclList {
   207  		switch d := decl.(type) {
   208  		case *syntax.FuncDecl:
   209  			if d.Name == nil {
   210  				newDecls = append(newDecls, d)
   211  				continue
   212  			}
   213  			obj := r.info.Defs[d.Name]
   214  			if !r.analyzer.isDependentObj[obj] || r.analyzer.inSimd {
   215  				newDecls = append(newDecls, d)
   216  				continue
   217  			}
   218  
   219  			sig, ok := obj.Type().(*types2.Signature)
   220  			if !ok {
   221  				newDecls = append(newDecls, d)
   222  				continue
   223  			}
   224  
   225  			change = true
   226  			if r.analyzer.HasDependentSignature(sig) {
   227  				if base.Debug.Simd > 0 {
   228  					base.Warn("%s: removing body of dependent-sig original function %v", d.Pos().String(), d.Name.Value)
   229  				}
   230  				d.Body = r.blockOf(d.Pos(), r.panicStmt(d.Pos(),
   231  					"unexpected call of original function rewritten to specialized SIMD"))
   232  				newDecls = append(newDecls, d)
   233  				continue
   234  			}
   235  
   236  			// Clean signature -> Replace body with dispatcher
   237  			d.Body = r.createDispatcherBody(d, sig)
   238  			newDecls = append(newDecls, d)
   239  
   240  		case *syntax.VarDecl:
   241  			// Keep var decls even if rewritten, so that pre-rewrite code parses correctly.
   242  			// TODO figure out how to deal with side-effects in initializers.
   243  			newDecls = append(newDecls, d)
   244  
   245  		case *syntax.TypeDecl:
   246  			// Keep all types; we need the untranslated copy if a method referencing it
   247  			// needs to typecheck pre-translation.
   248  			newDecls = append(newDecls, d)
   249  		default:
   250  			newDecls = append(newDecls, decl)
   251  		}
   252  	}
   253  
   254  	if !change {
   255  		return
   256  	}
   257  
   258  	fileAST.DeclList = newDecls
   259  
   260  	if !r.analyzer.inSimd {
   261  		// Inject an import to the bridge package (if not exists)
   262  		hasArchSimd := false
   263  		var simdImport *syntax.ImportDecl
   264  		p := fileAST.Pos()
   265  		for _, decl := range fileAST.DeclList {
   266  			if imp, ok := decl.(*syntax.ImportDecl); ok {
   267  				if imp.Path.Value == `"`+archFullPkg+`"` {
   268  					hasArchSimd = true
   269  					if simdImport == nil {
   270  						p = imp.Pos()
   271  					}
   272  				}
   273  				if imp.Path.Value == `"`+simdPkg+`"` {
   274  					simdImport = imp
   275  					p = imp.Pos()
   276  				}
   277  			}
   278  		}
   279  
   280  		if !hasArchSimd {
   281  			r.injectImport(fileAST, archFullPkg, p)
   282  		}
   283  
   284  		// Ensure at least one use of "simd"
   285  		// var _ = simd.VectorBitLen()
   286  		fun := &syntax.SelectorExpr{
   287  			X:   syntax.NewName(p, simdPkg), // Assume this is resolvable
   288  			Sel: syntax.NewName(p, vectorSizeFn),
   289  		}
   290  		fun.SetPos(p)
   291  		call := &syntax.CallExpr{Fun: fun}
   292  		call.SetPos(p)
   293  
   294  		name := syntax.NewName(p, "_")
   295  
   296  		varDecl := &syntax.VarDecl{NameList: []*syntax.Name{name}, Values: call}
   297  		varDecl.SetPos(p)
   298  		fileAST.DeclList = append(fileAST.DeclList, varDecl)
   299  	}
   300  }
   301  
   302  func (r *Rewriter) injectImport(fileAST *syntax.File, toImport string, simdImportPos syntax.Pos) {
   303  	importDecl := &syntax.ImportDecl{
   304  		Path: &syntax.BasicLit{Value: `"` + toImport + `"`, Kind: syntax.StringLit},
   305  	}
   306  	importDecl.Path.SetPos(simdImportPos)
   307  	importDecl.SetPos(simdImportPos)
   308  	fileAST.DeclList = append([]syntax.Decl{importDecl}, fileAST.DeclList...)
   309  }
   310  
   311  func (r *Rewriter) createDispatcherBody(d *syntax.FuncDecl, sig *types2.Signature) *syntax.BlockStmt {
   312  
   313  	// Build call arguments from the function parameters
   314  	args := func() []syntax.Expr {
   315  		var args []syntax.Expr
   316  		if d.Type.ParamList != nil {
   317  			for _, field := range d.Type.ParamList {
   318  				if field.Name != nil {
   319  					paramName := syntax.NewName(field.Pos(), field.Name.Value)
   320  					args = append(args, paramName)
   321  				}
   322  			}
   323  		}
   324  		return args
   325  	}
   326  
   327  	// Slap a pos on an expression
   328  	pe := func(e syntax.Expr) syntax.Expr {
   329  		e.SetPos(d.Pos())
   330  		return e
   331  	}
   332  	// Slap a pos on a statement
   333  	ps := func(e syntax.Stmt) syntax.Stmt {
   334  		e.SetPos(d.Pos())
   335  		return e
   336  	}
   337  
   338  	// switch ast node.
   339  	// the goal is something like (for now, till there are finer-grained choices)
   340  	// switch simd.VectorSize() {
   341  	//   case 128: if simd.Emulated() { call the specialize-for-emulation-code(args) }
   342  	//             else { call the specialize-for-128-code(args) }
   343  	//   case 256: call the specialize-for-256-code(args)
   344  	//   etc
   345  	// }
   346  	//
   347  	// the cases above deal with the usual `return call(...)` vs `call(...); return`
   348  	switchStmt := &syntax.SwitchStmt{
   349  		Tag: pe(&syntax.CallExpr{
   350  			Fun: pe(&syntax.SelectorExpr{
   351  				X:   syntax.NewName(d.Pos(), simdPkg), // Assume this is resolvable
   352  				Sel: syntax.NewName(d.Pos(), vectorSizeFn),
   353  			}),
   354  		}),
   355  		Body: []*syntax.CaseClause{},
   356  	}
   357  
   358  	var emulation syntax.Stmt
   359  
   360  	makeCallReturnStmt := func(k int, variantSuffix string) syntax.Stmt {
   361  		fnName := fmt.Sprintf("%s@simd%d%s", d.Name.Value, k, variantSuffix)
   362  		fnIdent := syntax.NewName(d.Pos(), fnName)
   363  
   364  		callExpr := pe(&syntax.CallExpr{
   365  			Fun:     pe(fnIdent),
   366  			ArgList: args(),
   367  		})
   368  
   369  		// callReturnStmt is either `return call(...)` or `call(...); return`
   370  		var callReturnStmt syntax.Stmt
   371  		if d.Type.ResultList != nil && len(d.Type.ResultList) > 0 {
   372  			callReturnStmt = &syntax.ReturnStmt{Results: callExpr}
   373  		} else {
   374  			callReturnStmt = &syntax.BlockStmt{
   375  				List: []syntax.Stmt{
   376  					ps(&syntax.ExprStmt{X: callExpr}),
   377  					ps(&syntax.ReturnStmt{}),
   378  				},
   379  				Rbrace: d.Pos(),
   380  			}
   381  		}
   382  		callReturnStmt.SetPos(d.Pos())
   383  		return callReturnStmt
   384  	}
   385  
   386  	guardCallWithCondition := func(require string, stmt syntax.Stmt) syntax.Stmt {
   387  		cond := pe(&syntax.CallExpr{
   388  			Fun: pe(&syntax.SelectorExpr{
   389  				X:   syntax.NewName(d.Pos(), simdPkg), // Assume this is resolvable
   390  				Sel: syntax.NewName(d.Pos(), require),
   391  			})})
   392  
   393  		blockStmt, ok := stmt.(*syntax.BlockStmt)
   394  		if !ok {
   395  			blockStmt = &syntax.BlockStmt{
   396  				List:   []syntax.Stmt{stmt},
   397  				Rbrace: d.Pos(),
   398  			}
   399  			blockStmt.SetPos(d.Pos())
   400  		}
   401  
   402  		guarded := ps(&syntax.IfStmt{
   403  			Cond: cond,
   404  			Then: blockStmt,
   405  		})
   406  		return guarded
   407  	}
   408  
   409  	for _, k := range r.sizes {
   410  
   411  		callReturnStmt := makeCallReturnStmt(k, "")
   412  
   413  		if k == 0 {
   414  			emulation = guardCallWithCondition(emulatedFn, callReturnStmt)
   415  			continue
   416  		}
   417  
   418  		var caseBody []syntax.Stmt
   419  		// assume that 128 is a case; when we do scalable simd, this may change.
   420  		// For now, if there is emulation, it is 128-bit (only).
   421  		if emulation != nil && k == 128 {
   422  			caseBody = append(caseBody, emulation)
   423  			emulation = nil
   424  		}
   425  
   426  		// if this architecture and size has a variant, then guard
   427  		// the regular case call with variant.require
   428  		// and then follow it with the variant case.
   429  		// `if simd.<require>() { callReturnStmt }`
   430  		if v := variants.Variants[variants.Key{Arch: buildcfg.GOARCH, Size: k}]; v != nil {
   431  			callReturnStmt = guardCallWithCondition(v.DefaultRequires, callReturnStmt)
   432  			caseBody = append(caseBody, callReturnStmt)
   433  			callReturnStmt = makeCallReturnStmt(k, v.Suffix)
   434  		}
   435  
   436  		caseBody = append(caseBody, callReturnStmt)
   437  
   438  		caseClause := &syntax.CaseClause{
   439  			Cases: pe(&syntax.BasicLit{Kind: syntax.IntLit, Value: fmt.Sprintf("%d", k)}),
   440  			Body:  caseBody,
   441  		}
   442  		caseClause.SetPos(d.Pos())
   443  		switchStmt.Body = append(switchStmt.Body, caseClause)
   444  	}
   445  
   446  	panicStmt := r.panicStmt(d.Pos(), "unsupported vector size in simd-rewritten code")
   447  	return r.blockOf(d.Pos(), switchStmt, panicStmt)
   448  }
   449  
   450  func (r *Rewriter) blockOf(p syntax.Pos, stmts ...syntax.Stmt) *syntax.BlockStmt {
   451  	for _, s := range stmts {
   452  		s.SetPos(p)
   453  	}
   454  	blockStmt := &syntax.BlockStmt{List: stmts}
   455  	blockStmt.SetPos(p)
   456  	return blockStmt
   457  }
   458  
   459  func (r *Rewriter) panicStmt(p syntax.Pos, unquotedMessage string) *syntax.ExprStmt {
   460  	pe := func(e syntax.Expr) syntax.Expr {
   461  		e.SetPos(p)
   462  		return e
   463  	}
   464  	fnName := "panic"
   465  	fnIdent := pe(syntax.NewName(p, fnName))
   466  	callExpr := pe(&syntax.CallExpr{
   467  		Fun:     fnIdent,
   468  		ArgList: []syntax.Expr{pe(&syntax.BasicLit{Value: `"` + unquotedMessage + `"`, Kind: syntax.StringLit})},
   469  	})
   470  	panicStmt := &syntax.ExprStmt{X: callExpr}
   471  	panicStmt.SetPos(p)
   472  	return panicStmt
   473  }
   474  
   475  func (r *Rewriter) generateForSize(fileAST *syntax.File, k int, variantSuffix string, newDecls []syntax.Decl) []syntax.Decl {
   476  	copier := NewDeepCopier(r.pkg, r.info, k, r.analyzer, fmt.Sprintf("@simd%d%s", k, variantSuffix), variantSuffix)
   477  	for _, decl := range fileAST.DeclList {
   478  		if r.shouldIncludeDecl(decl) {
   479  			newDecl := copier.CopyDecl(decl)
   480  			newDecls = append(newDecls, newDecl)
   481  		}
   482  	}
   483  	return newDecls
   484  }
   485  
   486  func nameToElemBitWidth(name string) int {
   487  	var width int
   488  	switch name {
   489  	case "Int8s", "Uint8s", "Mask8s":
   490  		width = 8
   491  	case "Int16s", "Uint16s", "Mask16s":
   492  		width = 16
   493  	case "Int32s", "Uint32s", "Float32s", "Mask32s":
   494  		width = 32
   495  	case "Int64s", "Uint64s", "Float64s", "Mask64s":
   496  		width = 64
   497  	}
   498  	return width
   499  }
   500  
   501  func (r *Rewriter) shouldIncludeDecl(decl syntax.Decl) bool {
   502  	// Files (and declarations) in the simd package are excluded
   503  	// from processing, except for those that whose name begins
   504  	// with "tofrom_".
   505  	if r.analyzer.inSimd {
   506  		theFile := decl.Pos().Base().Filename()
   507  
   508  		lastSlash := strings.LastIndex(theFile, simdPkg+"/")
   509  		lastBackslash := strings.LastIndex(theFile, simdPkg+"\\")
   510  
   511  		// Windows paths can be chaos, all we care, is whether the very last part
   512  		// of the path is any-path-separator + "tofrom_" + anything-else, given that
   513  		// we already know that we are in the simd package.
   514  		maxSlash := max(lastSlash, lastBackslash)
   515  		if maxSlash == -1 {
   516  			return false
   517  		}
   518  		if !strings.HasPrefix(theFile[maxSlash:], simdPkg+"/tofrom_") &&
   519  			!strings.HasPrefix(theFile[maxSlash:], simdPkg+"\\tofrom_") {
   520  			return false
   521  		}
   522  	}
   523  
   524  	switch d := decl.(type) {
   525  	case *syntax.FuncDecl:
   526  		if d.Name != nil {
   527  			return r.analyzer.isDependentObj[r.info.Defs[d.Name]]
   528  		}
   529  	case *syntax.TypeDecl:
   530  		return r.analyzer.isDependentObj[r.info.Defs[d.Name]]
   531  	case *syntax.VarDecl:
   532  		for _, name := range d.NameList {
   533  			if r.analyzer.isDependentObj[r.info.Defs[name]] {
   534  				return true
   535  			}
   536  		}
   537  	}
   538  	return false
   539  }
   540  
   541  // Generate an API matching the standalone compilation call
   542  func RewriteWrapper(pkg *types2.Package, info *types2.Info, files []*syntax.File) bool {
   543  	if !buildcfg.Experiment.SIMD {
   544  		return false
   545  	}
   546  
   547  	switch buildcfg.GOARCH {
   548  	case "wasm", "amd64", "arm64":
   549  	default:
   550  		return false
   551  	}
   552  
   553  	sizes := rewriteSizes()
   554  	if len(sizes) == 0 {
   555  		return false
   556  	}
   557  	analyzer := NewAnalyzer(pkg, info)
   558  	if !analyzer.Analyze(files) {
   559  		return false
   560  	}
   561  
   562  	CheckPositions(files, "before midway")
   563  
   564  	rewriter := NewRewriter(pkg, info, analyzer, sizes)
   565  	rewriter.Rewrite(files)
   566  
   567  	CheckPositions(files, "after midway")
   568  
   569  	return true
   570  }
   571
View as plain text