Source file src/simd/archsimd/_gen/specgen/specexpr/solver_test.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 specexpr
     6  
     7  import (
     8  	"fmt"
     9  	"maps"
    10  	"slices"
    11  	"strings"
    12  	"testing"
    13  )
    14  
    15  func newSolver(t *testing.T) *Solver {
    16  	var s Solver
    17  	s.SetTrace(t.Output())
    18  	return &s
    19  }
    20  
    21  func TestSolver(t *testing.T) {
    22  	t.Run("constant assignment", func(t *testing.T) {
    23  		s := newSolver(t)
    24  		v1 := Variable("v1")
    25  		s.Assign(v1, Int(10))
    26  
    27  		sol := uniqueSolution(t, s)
    28  		bCheck(t, sol, map[Variable]any{v1: Int(10)})
    29  	})
    30  
    31  	t.Run("simple dependency", func(t *testing.T) {
    32  		testBinaryOp(t, 10, &BinExpr{
    33  			Op: OpTimes,
    34  			X:  Variable("v1"),
    35  			Y:  Int(2),
    36  		}, Int(20))
    37  	})
    38  
    39  	t.Run("division", func(t *testing.T) {
    40  		testBinaryOp(t, 16, &BinExpr{
    41  			Op: OpDiv,
    42  			X:  Variable("v1"),
    43  			Y:  Int(4),
    44  		}, Int(4))
    45  	})
    46  
    47  	t.Run("cycle detection", func(t *testing.T) {
    48  		s := newSolver(t)
    49  		v1 := Variable("v1")
    50  		v2 := Variable("v2")
    51  		s.Assign(v1, Variable(v2))
    52  		s.Assign(v2, Variable(v1))
    53  
    54  		err := solverError(t, s)
    55  		if !strings.Contains(err.Error(), "cyclic requirements") {
    56  			t.Fatalf("expected cycle error, got %v", err)
    57  		}
    58  	})
    59  
    60  	t.Run("multiple assignment conflicting values", func(t *testing.T) {
    61  		s := newSolver(t)
    62  		v1 := Variable("v1")
    63  		s.Assign(v1, Int(10))
    64  		s.Assign(v1, Int(20))
    65  
    66  		err := solverError(t, s)
    67  		if !strings.Contains(err.Error(), "no solutions") {
    68  			t.Fatalf("expected no solutions error, got %v", err)
    69  		}
    70  	})
    71  
    72  	t.Run("multiple assignment same value", func(t *testing.T) {
    73  		s := newSolver(t)
    74  		v1 := Variable("v1")
    75  		s.Assign(v1, Int(10))
    76  		s.Assign(v1, Int(10))
    77  
    78  		sol := uniqueSolution(t, s)
    79  		bCheck(t, sol, map[Variable]any{v1: Int(10)})
    80  	})
    81  
    82  	t.Run("swidth times int", func(t *testing.T) {
    83  		s := newSolver(t)
    84  		v1 := Variable("v1")
    85  		v2 := Variable("v2")
    86  		s.Assign(v1, mkWidth(1, 2))
    87  		s.Assign(v2, &BinExpr{
    88  			Op: OpTimes,
    89  			X:  v1,
    90  			Y:  Int(4),
    91  		})
    92  
    93  		sol := uniqueSolution(t, s)
    94  		bCheck(t, sol, map[Variable]any{
    95  			v1: mkWidth(1, 2),
    96  			v2: mkWidth(2, 1),
    97  		})
    98  	})
    99  
   100  	t.Run("int times swidth", func(t *testing.T) {
   101  		s := newSolver(t)
   102  		v1 := Variable("v1")
   103  		v2 := Variable("v2")
   104  		s.Assign(v1, mkWidth(1, 2))
   105  		s.Assign(v2, &BinExpr{
   106  			Op: OpTimes,
   107  			X:  Int(4),
   108  			Y:  v1,
   109  		})
   110  
   111  		sol := uniqueSolution(t, s)
   112  		bCheck(t, sol, map[Variable]any{
   113  			v1: mkWidth(1, 2),
   114  			v2: mkWidth(2, 1),
   115  		})
   116  	})
   117  
   118  	t.Run("swidth div int", func(t *testing.T) {
   119  		s := newSolver(t)
   120  		v1 := Variable("v1")
   121  		v2 := Variable("v2")
   122  		s.Assign(v1, mkWidth(1, 2))
   123  		s.Assign(v2, &BinExpr{
   124  			Op: OpDiv,
   125  			X:  v1,
   126  			Y:  Int(2),
   127  		})
   128  
   129  		sol := uniqueSolution(t, s)
   130  		bCheck(t, sol, map[Variable]any{
   131  			v1: mkWidth(1, 2),
   132  			v2: mkWidth(1, 4),
   133  		})
   134  	})
   135  }
   136  
   137  func testBinaryOp(t *testing.T, v1Val Int, v2Expr Expr, expectedV2Val any) {
   138  	t.Helper()
   139  	s := newSolver(t)
   140  	v1 := Variable("v1")
   141  	v2 := Variable("v2")
   142  	s.Assign(v1, v1Val)
   143  	s.Assign(v2, v2Expr)
   144  
   145  	sol := uniqueSolution(t, s)
   146  	bCheck(t, sol, map[Variable]any{v1: v1Val, v2: expectedV2Val})
   147  }
   148  
   149  func TestComparisons(t *testing.T) {
   150  	t.Run("greater than", func(t *testing.T) {
   151  		s := newSolver(t)
   152  		v1 := Variable("v1")
   153  		v2 := Variable("v2")
   154  		s.Assign(v1, Int(20))
   155  		s.Assign(v2, Int(10))
   156  
   157  		// Add comparison v1 > v2 and assert it must be true
   158  		s.Assert(&BinExpr{Op: OpGreaterThan, X: v1, Y: v2})
   159  
   160  		uniqueSolution(t, s)
   161  	})
   162  
   163  	t.Run("less than", func(t *testing.T) {
   164  		s := newSolver(t)
   165  		v1 := Variable("v1")
   166  		v2 := Variable("v2")
   167  		s.Assign(v1, Int(10))
   168  		s.Assign(v2, Int(20))
   169  
   170  		s.Assert(&BinExpr{Op: OpLessThan, X: v1, Y: v2})
   171  
   172  		uniqueSolution(t, s)
   173  	})
   174  }
   175  
   176  func TestSolveShape(t *testing.T) {
   177  	t.Run("scalar Int32", func(t *testing.T) {
   178  		s := newSolver(t)
   179  		s.Assign("x", mustParseExpr(t, "Int32"))
   180  
   181  		sol := uniqueSolution(t, s)
   182  		bCheck(t, sol, map[Variable]any{
   183  			"x": Basic{"int", 32},
   184  		})
   185  	})
   186  
   187  	t.Run("vector Int32x4", func(t *testing.T) {
   188  		s := newSolver(t)
   189  		s.Assign("x", mustParseExpr(t, "Int32x4"))
   190  
   191  		sol := uniqueSolution(t, s)
   192  		bCheck(t, sol, map[Variable]any{
   193  			"x": Vector{Basic{"int", 32}, Int(128)},
   194  		})
   195  	})
   196  
   197  	t.Run("scalar symbolic {xB}{xN}", func(t *testing.T) {
   198  		s := newSolver(t)
   199  		s.Assign("x", mustParseExpr(t, "{xB}{xN}"))
   200  		s.Assign("xB", &Literal{"int"})
   201  		s.Assign("xN", Int(32))
   202  
   203  		sol := uniqueSolution(t, s)
   204  		bCheck(t, sol, map[Variable]any{
   205  			"x":  Basic{"int", 32},
   206  			"xB": "int",
   207  			"xN": Int(32),
   208  		})
   209  	})
   210  
   211  	vectorElem := MakeField[Vector]("Elem")
   212  	basicBase := MakeField[Basic]("Base")
   213  	basicBits := MakeField[Basic]("Bits")
   214  	vectorWidth := MakeField[Vector]("Width")
   215  	assignVector := func(s *Solver, v Variable, e Expr) {
   216  		x := s.Assign(v, e)
   217  		s.Assign(v+"B", basicBase.Apply(vectorElem.Apply(x)))
   218  		xN := s.Assign(v+"N", basicBits.Apply(vectorElem.Apply(x)))
   219  		xW := s.Assign(v+"W", vectorWidth.Apply(x))
   220  		s.Assign(v+"L", &BinExpr{Op: OpDiv, X: xW, Y: xN})
   221  	}
   222  
   223  	t.Run("derived scalable vector with lane count", func(t *testing.T) {
   224  		s := newSolver(t)
   225  		assignVector(s, "x", mustParseExpr(t, "Int32s"))
   226  		s.Assign("y", mustParseExpr(t, "{xB}{xN*2}x{xL/2}"))
   227  
   228  		sol := uniqueSolution(t, s)
   229  		bCheck(t, sol, map[Variable]any{
   230  			"y":  Vector{Basic{"int", 64}, mkWidth(1, 1)},
   231  			"xW": mkWidth(1, 1),
   232  			"xL": mkWidth(1, 32),
   233  		})
   234  	})
   235  
   236  	t.Run("derived scalable vector with width", func(t *testing.T) {
   237  		s := newSolver(t)
   238  		assignVector(s, "x", mustParseExpr(t, "Int32s"))
   239  		s.Assign("y", mustParseExpr(t, "{xB}{xN*2}w{xW}"))
   240  
   241  		sol := uniqueSolution(t, s)
   242  		bCheck(t, sol, map[Variable]any{
   243  			"y":  Vector{Basic{"int", 64}, mkWidth(1, 1)},
   244  			"xW": mkWidth(1, 1),
   245  		})
   246  	})
   247  
   248  	t.Run("width rounding", func(t *testing.T) {
   249  		s := newSolver(t)
   250  		assignVector(s, "x", mustParseExpr(t, "Int64x2"))
   251  		s.Assign("y", mustParseExpr(t, "{xB}{xN/2}x{xL}"))
   252  
   253  		sol := uniqueSolution(t, s)
   254  		bCheck(t, sol, map[Variable]any{
   255  			"y": Vector{Basic{"int", 32}, Int(128)},
   256  		})
   257  	})
   258  
   259  	t.Run("domain limits", func(t *testing.T) {
   260  		s := newSolver(t)
   261  		s.Declare("x", []any{1, 2})
   262  		s.Declare("y", []any{1, 2})
   263  		s.Assign("y", mustParseExpr(t, "x*2"))
   264  
   265  		sol := uniqueSolution(t, s)
   266  		bCheck(t, sol, map[Variable]any{
   267  			"x": Int(1), "y": Int(2),
   268  		})
   269  	})
   270  
   271  	t.Run("non-scalable width", func(t *testing.T) {
   272  		s := newSolver(t)
   273  		assignVector(s, "x", mustParseExpr(t, "Int32s"))
   274  		s.Assign("y", mustParseExpr(t, "Int16x{xL}"))
   275  
   276  		err := solverError(t, s)
   277  		if !strings.Contains(err.Error(), "invalid width") {
   278  			t.Fatalf("expected invalid width error, got: %v", err)
   279  		}
   280  	})
   281  }
   282  
   283  func TestEnumerator(t *testing.T) {
   284  	t.Run("simple enumeration", func(t *testing.T) {
   285  		s := newSolver(t)
   286  		v1 := Variable("v1")
   287  		v2 := Variable("v2")
   288  
   289  		s.Declare(v1, []any{1, 2, 3, 4})
   290  
   291  		// Assert v2 = v1 * 2
   292  		s.Assign(v2, &BinExpr{Op: OpTimes, X: v1, Y: Int(2)})
   293  
   294  		sols := allSolutions(s)
   295  
   296  		if len(sols) != 4 {
   297  			t.Errorf("expected 4 solutions, got %d", len(sols))
   298  		}
   299  
   300  		// Verify that each solution maps v2 to v1*2
   301  		for _, sol := range sols {
   302  			m := bmap(sol)
   303  			v1Val := m[v1].(Int)
   304  			v2Val := m[v2].(Int)
   305  			if v2Val != v1Val*2 {
   306  				t.Errorf("solution %v violates v2 = v1*2", m)
   307  			}
   308  		}
   309  	})
   310  
   311  	t.Run("comparison constraints enumeration", func(t *testing.T) {
   312  		s := newSolver(t)
   313  		v1 := Variable("v1")
   314  		v2 := Variable("v2")
   315  
   316  		s.Declare(v1, []any{1, 2, 3, 4, 5})
   317  
   318  		// v2 = v1 * 2
   319  		s.Assign(v2, &BinExpr{Op: OpTimes, X: v1, Y: Int(2)})
   320  
   321  		// v1 > 2
   322  		s.Assert(&BinExpr{Op: OpGreaterThan, X: v1, Y: Int(2)})
   323  
   324  		// v2 < 10
   325  		s.Assert(&BinExpr{Op: OpLessThan, X: v2, Y: Int(10)})
   326  
   327  		sols := allSolutions(s)
   328  
   329  		// Solutions should be v1=3, v1=4. So 2 solutions
   330  		if len(sols) != 2 {
   331  			t.Errorf("expected 2 solutions, got %d", len(sols))
   332  		}
   333  	})
   334  }
   335  
   336  func mustParseExpr(t *testing.T, x string) Expr {
   337  	t.Helper()
   338  	e, err := ParseExpr(x)
   339  	if err != nil {
   340  		t.Fatal(err)
   341  	}
   342  	return e
   343  }
   344  
   345  func allSolutions(s *Solver) []*Bindings {
   346  	var sols []*Bindings
   347  	for b, err := range s.Solve() {
   348  		if err != nil {
   349  			panic(err) // tests expect valid execution when enumerating
   350  		}
   351  		sols = append(sols, b)
   352  	}
   353  	return sols
   354  }
   355  
   356  func uniqueSolution(t *testing.T, s *Solver) *Bindings {
   357  	t.Helper()
   358  	var sols []*Bindings
   359  	for b, err := range s.Solve() {
   360  		if err != nil {
   361  			t.Fatalf("solve failed: %v", err)
   362  		}
   363  		sols = append(sols, b)
   364  		if len(sols) >= 20 {
   365  			// Stop before we go too deep
   366  			t.Fatalf("expected exactly one solution, got >= 20")
   367  		}
   368  	}
   369  	if len(sols) != 1 {
   370  		t.Errorf("expected exactly one solution, got %d", len(sols))
   371  		for _, sol := range sols {
   372  			t.Errorf("  %s", sol)
   373  		}
   374  		t.FailNow()
   375  	}
   376  	return sols[0]
   377  }
   378  
   379  func solverError(t *testing.T, s *Solver) error {
   380  	t.Helper()
   381  	for soln, err := range s.Solve() {
   382  		if err != nil {
   383  			return err
   384  		}
   385  		t.Fatalf("expected solver error, but got solution:\n%s", soln)
   386  	}
   387  	return fmt.Errorf("no solutions")
   388  }
   389  
   390  // bmap converts a Bindings to a map.
   391  func bmap(b *Bindings) map[Variable]any {
   392  	return maps.Collect(b.All())
   393  }
   394  
   395  // bCheck fails t if got[v] != want[v] for any keys in want.
   396  func bCheck(t *testing.T, got *Bindings, want map[Variable]any) {
   397  	t.Helper()
   398  	var keys []Variable
   399  	for k := range want {
   400  		keys = append(keys, k)
   401  	}
   402  	slices.Sort(keys)
   403  
   404  	var mismatches []string
   405  	for _, k := range keys {
   406  		wantVal := want[k]
   407  		gotVal := got.Get(k)
   408  		if gotVal != wantVal {
   409  			mismatches = append(mismatches, fmt.Sprintf("  %s: got %v (%T), want %v (%T)", k, gotVal, gotVal, wantVal, wantVal))
   410  		}
   411  	}
   412  	if len(mismatches) > 0 {
   413  		t.Fatalf("solution mismatch:\n%s", strings.Join(mismatches, "\n"))
   414  	}
   415  }
   416  

View as plain text