num.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  	"cmp"
     9  	"fmt"
    10  	"strings"
    11  )
    12  
    13  // Num represents a number in the solver. This is abstracted because we work
    14  // with both concrete numbers and *symbolic widths* and need to be able to mix
    15  // them. A Num is also an [Expr].
    16  type Num interface {
    17  	Expr
    18  
    19  	// ValidWidth returns true if this is a valid width: either a fixed width
    20  	// (128, 256, 512) or the scalable width VW.
    21  	ValidWidth() bool
    22  
    23  	Mul(x Num) (Num, error)
    24  	Div(x Num) (Num, error)
    25  	Compare(o Num) (int, bool)
    26  
    27  	String() string
    28  }
    29  
    30  // Int is an integer that satisfies [Num].
    31  type Int int
    32  
    33  func (w Int) ValidWidth() bool {
    34  	switch w {
    35  	case 128, 256, 512:
    36  		return true
    37  	}
    38  	return false
    39  }
    40  func (w Int) Mul(x Num) (Num, error) {
    41  	switch x := x.(type) {
    42  	case Int:
    43  		return w * x, nil
    44  	case ScalableWidth:
    45  		return mkWidth(int(w)*x.num, x.denom), nil
    46  	}
    47  	panic("unknown Num")
    48  }
    49  func (w Int) Div(x Num) (Num, error) {
    50  	switch x := x.(type) {
    51  	case Int:
    52  		if x == 0 {
    53  			return nil, fmt.Errorf("division by zero")
    54  		}
    55  		if w%x != 0 {
    56  			return nil, fmt.Errorf("inexact division %d/%d", w, x)
    57  		}
    58  		return w / x, nil
    59  	case ScalableWidth:
    60  		// ScalableWidth is implicitly multiplied by VW. We have no way to
    61  		// express the inverse of VW.
    62  		return nil, fmt.Errorf("cannot divide Int by ScalableWidth")
    63  	}
    64  	panic("unknown Num")
    65  }
    66  func (w Int) Compare(o Num) (int, bool) {
    67  	if o, ok := o.(Int); ok {
    68  		return cmp.Compare(w, o), true
    69  	}
    70  	return 0, false
    71  }
    72  func (w Int) String() string {
    73  	return fmt.Sprint(int(w))
    74  }
    75  func (w Int) eval(b *Bindings) (any, error) {
    76  	return w, nil
    77  }
    78  func (w Int) preorder(yield func(Expr) bool) bool {
    79  	return true
    80  }
    81  
    82  // An ScalableWidth represents a symbolic width relative to a fixed but unknown
    83  // scalable vector width VW. This is represented as a rational factor
    84  // VW*num/denom
    85  type ScalableWidth struct {
    86  	num, denom int
    87  }
    88  
    89  // VW returns the base ScalableWidth representing a full-width scalable vector.
    90  func VW() ScalableWidth {
    91  	return ScalableWidth{1, 1}
    92  }
    93  
    94  func gcd(a, b int) int {
    95  	for b != 0 {
    96  		a, b = b, a%b
    97  	}
    98  	if a < 0 {
    99  		return -a
   100  	}
   101  	return a
   102  }
   103  
   104  func mkWidth(num, denom int) ScalableWidth {
   105  	if denom == 0 {
   106  		panic("denominator cannot be zero")
   107  	}
   108  	g := gcd(num, denom)
   109  	num /= g
   110  	denom /= g
   111  	if denom < 0 {
   112  		num = -num
   113  		denom = -denom
   114  	}
   115  	return ScalableWidth{num, denom}
   116  }
   117  
   118  func (w ScalableWidth) ValidWidth() bool {
   119  	return w == ScalableWidth{1, 1}
   120  }
   121  
   122  func (w ScalableWidth) Mul(x Num) (Num, error) {
   123  	switch x := x.(type) {
   124  	case Int:
   125  		return mkWidth(w.num*int(x), w.denom), nil
   126  	case ScalableWidth:
   127  		return nil, fmt.Errorf("cannot multiply two scalable widths")
   128  	}
   129  	panic("unknown Num")
   130  }
   131  
   132  func (w ScalableWidth) Div(x Num) (Num, error) {
   133  	switch x := x.(type) {
   134  	case Int:
   135  		if x == 0 {
   136  			return nil, fmt.Errorf("division by zero")
   137  		}
   138  		return mkWidth(w.num, w.denom*int(x)), nil
   139  	case ScalableWidth:
   140  		a, b := w.num*x.denom, w.denom*x.num
   141  		if b == 0 {
   142  			return nil, fmt.Errorf("division by zero")
   143  		}
   144  		if a%b != 0 {
   145  			return nil, fmt.Errorf("inexact division %d/%d", w, x)
   146  		}
   147  		return Int(a / b), nil
   148  	}
   149  	panic("unknown Num")
   150  }
   151  
   152  func (w ScalableWidth) Compare(o Num) (int, bool) {
   153  	if o, ok := o.(ScalableWidth); ok {
   154  		return cmp.Compare(w.num*o.denom, o.num*w.denom), true
   155  	}
   156  	return 0, false
   157  }
   158  
   159  func (w ScalableWidth) String() string {
   160  	var buf strings.Builder
   161  	buf.WriteString("VW")
   162  	if w.num > 1 {
   163  		fmt.Fprintf(&buf, "*%d", w.num)
   164  	}
   165  	if w.denom > 1 {
   166  		fmt.Fprintf(&buf, "/%d", w.denom)
   167  	}
   168  	return buf.String()
   169  }
   170  
   171  func (w ScalableWidth) eval(b *Bindings) (any, error) {
   172  	return w, nil
   173  }
   174  
   175  func (w ScalableWidth) preorder(yield func(Expr) bool) bool {
   176  	return true
   177  }
   178
View as plain text