Skip to content

Commit b86945d

Browse files
LegNeatoLegNeato
committed
Add TypeDomain tracking to egraph to prevent cross-type rewrites
Adds a TypeDomain datatype (BoolType, IntType, FloatType) and an ExprType function that tracks the type domain of each expression in the e-graph. This prevents rewrite rules from crossing type boundaries (e.g., converting an int-domain Gamma to a bool-domain LogAnd). Rules that could cross type domains are now guarded with ExprType checks: cleanup.egg Select/Gamma-to-bool rules require BoolType, rvsdg.egg Gamma-to-LogAnd/LogOr rules require BoolType operands, and logical.egg short-circuit rules require BoolType results.
1 parent 24c4c1a commit b86945d

4 files changed

Lines changed: 246 additions & 12 deletions

File tree

‎rust/spirv-tools-opt/src/rules/cleanup.egg‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -71,17 +71,18 @@
7171
((union e cond)))
7272
(rule ((= e (Select cond (BoolConst 0) (BoolConst 1))))
7373
((union e (LogNot cond))))
74-
(rule ((= e (Select cond (Const 1) (Const 0))))
74+
; Guard Const variants with ExprType to prevent int→bool domain crossing
75+
(rule ((= e (Select cond (Const 1) (Const 0))) (= (ExprType e) (BoolType)))
7576
((union e cond)))
76-
(rule ((= e (Select cond (Const 0) (Const 1))))
77+
(rule ((= e (Select cond (Const 0) (Const 1))) (= (ExprType e) (BoolType)))
7778
((union e (LogNot cond))))
7879
(rule ((= e (Gamma cond (BoolConst 1) (BoolConst 0))))
7980
((union e cond)))
8081
(rule ((= e (Gamma cond (BoolConst 0) (BoolConst 1))))
8182
((union e (LogNot cond))))
82-
(rule ((= e (Gamma cond (Const 1) (Const 0))))
83+
(rule ((= e (Gamma cond (Const 1) (Const 0))) (= (ExprType e) (BoolType)))
8384
((union e cond)))
84-
(rule ((= e (Gamma cond (Const 0) (Const 1))))
85+
(rule ((= e (Gamma cond (Const 0) (Const 1))) (= (ExprType e) (BoolType)))
8586
((union e (LogNot cond))))
8687

8788
; =============================================================================

‎rust/spirv-tools-opt/src/rules/datatypes.egg‎

Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -775,6 +775,234 @@
775775
(BitCast Expr)
776776
)
777777

778+
; =============================================================================
779+
; ExprType: Type Class Tracking
780+
; =============================================================================
781+
; Tracks the SPIR-V type domain for each expression in the e-graph.
782+
; :merge old preserves the first-set type when e-classes merge across domains.
783+
; This enables type-aware rewrite guards and correct type inference during extraction.
784+
; Sym nodes are seeded from Rust code; constructors get types from rules below.
785+
786+
(datatype TypeDomain (BoolType) (IntType) (FloatType))
787+
788+
(function ExprType (Expr) TypeDomain :merge old)
789+
790+
; --- Boolean-producing operations ---
791+
; Integer comparisons
792+
(rule ((= e (Eq a b))) ((set (ExprType e) (BoolType))))
793+
(rule ((= e (Ne a b))) ((set (ExprType e) (BoolType))))
794+
(rule ((= e (SLt a b))) ((set (ExprType e) (BoolType))))
795+
(rule ((= e (SLe a b))) ((set (ExprType e) (BoolType))))
796+
(rule ((= e (SGt a b))) ((set (ExprType e) (BoolType))))
797+
(rule ((= e (SGe a b))) ((set (ExprType e) (BoolType))))
798+
(rule ((= e (ULt a b))) ((set (ExprType e) (BoolType))))
799+
(rule ((= e (ULe a b))) ((set (ExprType e) (BoolType))))
800+
(rule ((= e (UGt a b))) ((set (ExprType e) (BoolType))))
801+
(rule ((= e (UGe a b))) ((set (ExprType e) (BoolType))))
802+
; Float comparisons (ordered)
803+
(rule ((= e (FOrdEq a b))) ((set (ExprType e) (BoolType))))
804+
(rule ((= e (FOrdNe a b))) ((set (ExprType e) (BoolType))))
805+
(rule ((= e (FOrdLt a b))) ((set (ExprType e) (BoolType))))
806+
(rule ((= e (FOrdLe a b))) ((set (ExprType e) (BoolType))))
807+
(rule ((= e (FOrdGt a b))) ((set (ExprType e) (BoolType))))
808+
(rule ((= e (FOrdGe a b))) ((set (ExprType e) (BoolType))))
809+
; Float comparisons (unordered)
810+
(rule ((= e (FUnordEq a b))) ((set (ExprType e) (BoolType))))
811+
(rule ((= e (FUnordNe a b))) ((set (ExprType e) (BoolType))))
812+
(rule ((= e (FUnordLt a b))) ((set (ExprType e) (BoolType))))
813+
(rule ((= e (FUnordLe a b))) ((set (ExprType e) (BoolType))))
814+
(rule ((= e (FUnordGt a b))) ((set (ExprType e) (BoolType))))
815+
(rule ((= e (FUnordGe a b))) ((set (ExprType e) (BoolType))))
816+
; Additional float comparisons
817+
(rule ((= e (FEq a b))) ((set (ExprType e) (BoolType))))
818+
(rule ((= e (FNe a b))) ((set (ExprType e) (BoolType))))
819+
(rule ((= e (FLt a b))) ((set (ExprType e) (BoolType))))
820+
(rule ((= e (FLe a b))) ((set (ExprType e) (BoolType))))
821+
(rule ((= e (FGt a b))) ((set (ExprType e) (BoolType))))
822+
(rule ((= e (FGe a b))) ((set (ExprType e) (BoolType))))
823+
; Logical operations
824+
(rule ((= e (LogNot a))) ((set (ExprType e) (BoolType))))
825+
(rule ((= e (LogAnd a b))) ((set (ExprType e) (BoolType))))
826+
(rule ((= e (LogOr a b))) ((set (ExprType e) (BoolType))))
827+
(rule ((= e (LogEq a b))) ((set (ExprType e) (BoolType))))
828+
(rule ((= e (LogNe a b))) ((set (ExprType e) (BoolType))))
829+
; FP predicates
830+
(rule ((= e (IsNan a))) ((set (ExprType e) (BoolType))))
831+
(rule ((= e (IsInf a))) ((set (ExprType e) (BoolType))))
832+
; Boolean constant
833+
(rule ((= e (BoolConst v))) ((set (ExprType e) (BoolType))))
834+
; Subgroup boolean operations
835+
(rule ((= e (GroupElect))) ((set (ExprType e) (BoolType))))
836+
(rule ((= e (GroupAll a))) ((set (ExprType e) (BoolType))))
837+
(rule ((= e (GroupAny a))) ((set (ExprType e) (BoolType))))
838+
(rule ((= e (GroupAllEqual a))) ((set (ExprType e) (BoolType))))
839+
; Vector boolean predicates
840+
(rule ((= e (Any a))) ((set (ExprType e) (BoolType))))
841+
(rule ((= e (All a))) ((set (ExprType e) (BoolType))))
842+
843+
; --- Integer-producing operations ---
844+
; Arithmetic
845+
(rule ((= e (Add a b))) ((set (ExprType e) (IntType))))
846+
(rule ((= e (Sub a b))) ((set (ExprType e) (IntType))))
847+
(rule ((= e (Mul a b))) ((set (ExprType e) (IntType))))
848+
(rule ((= e (Neg a))) ((set (ExprType e) (IntType))))
849+
(rule ((= e (SDiv a b))) ((set (ExprType e) (IntType))))
850+
(rule ((= e (UDiv a b))) ((set (ExprType e) (IntType))))
851+
(rule ((= e (SRem a b))) ((set (ExprType e) (IntType))))
852+
(rule ((= e (SMod a b))) ((set (ExprType e) (IntType))))
853+
(rule ((= e (UMod a b))) ((set (ExprType e) (IntType))))
854+
; Bitwise
855+
(rule ((= e (Shl a b))) ((set (ExprType e) (IntType))))
856+
(rule ((= e (ShrS a b))) ((set (ExprType e) (IntType))))
857+
(rule ((= e (ShrU a b))) ((set (ExprType e) (IntType))))
858+
(rule ((= e (BitAnd a b))) ((set (ExprType e) (IntType))))
859+
(rule ((= e (BitOr a b))) ((set (ExprType e) (IntType))))
860+
(rule ((= e (BitXor a b))) ((set (ExprType e) (IntType))))
861+
(rule ((= e (BitNot a))) ((set (ExprType e) (IntType))))
862+
(rule ((= e (BitReverse a))) ((set (ExprType e) (IntType))))
863+
(rule ((= e (BitCount a))) ((set (ExprType e) (IntType))))
864+
(rule ((= e (RotL a b))) ((set (ExprType e) (IntType))))
865+
(rule ((= e (RotR a b))) ((set (ExprType e) (IntType))))
866+
; Integer min/max
867+
(rule ((= e (SMin a b))) ((set (ExprType e) (IntType))))
868+
(rule ((= e (SMax a b))) ((set (ExprType e) (IntType))))
869+
(rule ((= e (UMin a b))) ((set (ExprType e) (IntType))))
870+
(rule ((= e (UMax a b))) ((set (ExprType e) (IntType))))
871+
; Integer GLSL operations
872+
(rule ((= e (SAbs a))) ((set (ExprType e) (IntType))))
873+
(rule ((= e (Sign a))) ((set (ExprType e) (IntType))))
874+
(rule ((= e (FindILsb a))) ((set (ExprType e) (IntType))))
875+
(rule ((= e (FindSMsb a))) ((set (ExprType e) (IntType))))
876+
(rule ((= e (FindUMsb a))) ((set (ExprType e) (IntType))))
877+
; Float-to-int conversions
878+
(rule ((= e (ConvertFToS a))) ((set (ExprType e) (IntType))))
879+
(rule ((= e (ConvertFToU a))) ((set (ExprType e) (IntType))))
880+
; Integer width conversions
881+
(rule ((= e (SConvert a))) ((set (ExprType e) (IntType))))
882+
(rule ((= e (UConvert a))) ((set (ExprType e) (IntType))))
883+
; Integer constants
884+
(rule ((= e (Const v))) ((set (ExprType e) (IntType))))
885+
(rule ((= e (Const64 v))) ((set (ExprType e) (IntType))))
886+
; Packing (int result)
887+
(rule ((= e (PackSnorm4x8 a))) ((set (ExprType e) (IntType))))
888+
(rule ((= e (PackUnorm4x8 a))) ((set (ExprType e) (IntType))))
889+
(rule ((= e (PackSnorm2x16 a))) ((set (ExprType e) (IntType))))
890+
(rule ((= e (PackUnorm2x16 a))) ((set (ExprType e) (IntType))))
891+
(rule ((= e (PackHalf2x16 a))) ((set (ExprType e) (IntType))))
892+
(rule ((= e (PackDouble2x32 a))) ((set (ExprType e) (IntType))))
893+
; Subgroup integer reductions
894+
(rule ((= e (GroupIAdd a))) ((set (ExprType e) (IntType))))
895+
(rule ((= e (GroupIMul a))) ((set (ExprType e) (IntType))))
896+
(rule ((= e (GroupSMin a))) ((set (ExprType e) (IntType))))
897+
(rule ((= e (GroupUMin a))) ((set (ExprType e) (IntType))))
898+
(rule ((= e (GroupSMax a))) ((set (ExprType e) (IntType))))
899+
(rule ((= e (GroupUMax a))) ((set (ExprType e) (IntType))))
900+
(rule ((= e (GroupBitAnd a))) ((set (ExprType e) (IntType))))
901+
(rule ((= e (GroupBitOr a))) ((set (ExprType e) (IntType))))
902+
(rule ((= e (GroupBitXor a))) ((set (ExprType e) (IntType))))
903+
904+
; --- Float-producing operations ---
905+
; Arithmetic
906+
(rule ((= e (FAdd a b))) ((set (ExprType e) (FloatType))))
907+
(rule ((= e (FSub a b))) ((set (ExprType e) (FloatType))))
908+
(rule ((= e (FMul a b))) ((set (ExprType e) (FloatType))))
909+
(rule ((= e (FDiv a b))) ((set (ExprType e) (FloatType))))
910+
(rule ((= e (FNeg a))) ((set (ExprType e) (FloatType))))
911+
(rule ((= e (FRem a b))) ((set (ExprType e) (FloatType))))
912+
(rule ((= e (FMod a b))) ((set (ExprType e) (FloatType))))
913+
(rule ((= e (FMin a b))) ((set (ExprType e) (FloatType))))
914+
(rule ((= e (FMax a b))) ((set (ExprType e) (FloatType))))
915+
(rule ((= e (FAbs a))) ((set (ExprType e) (FloatType))))
916+
(rule ((= e (FFloor a))) ((set (ExprType e) (FloatType))))
917+
(rule ((= e (FCeil a))) ((set (ExprType e) (FloatType))))
918+
(rule ((= e (FRound a))) ((set (ExprType e) (FloatType))))
919+
(rule ((= e (FTrunc a))) ((set (ExprType e) (FloatType))))
920+
; Int-to-float conversions
921+
(rule ((= e (ConvertSToF a))) ((set (ExprType e) (FloatType))))
922+
(rule ((= e (ConvertUToF a))) ((set (ExprType e) (FloatType))))
923+
(rule ((= e (FConvert a))) ((set (ExprType e) (FloatType))))
924+
; Trigonometric
925+
(rule ((= e (Sin a))) ((set (ExprType e) (FloatType))))
926+
(rule ((= e (Cos a))) ((set (ExprType e) (FloatType))))
927+
(rule ((= e (Tan a))) ((set (ExprType e) (FloatType))))
928+
(rule ((= e (Asin a))) ((set (ExprType e) (FloatType))))
929+
(rule ((= e (Acos a))) ((set (ExprType e) (FloatType))))
930+
(rule ((= e (Atan a))) ((set (ExprType e) (FloatType))))
931+
(rule ((= e (Atan2 a b))) ((set (ExprType e) (FloatType))))
932+
(rule ((= e (Sinh a))) ((set (ExprType e) (FloatType))))
933+
(rule ((= e (Cosh a))) ((set (ExprType e) (FloatType))))
934+
(rule ((= e (Tanh a))) ((set (ExprType e) (FloatType))))
935+
(rule ((= e (Asinh a))) ((set (ExprType e) (FloatType))))
936+
(rule ((= e (Acosh a))) ((set (ExprType e) (FloatType))))
937+
(rule ((= e (Atanh a))) ((set (ExprType e) (FloatType))))
938+
; Exponential / logarithmic
939+
(rule ((= e (Sqrt a))) ((set (ExprType e) (FloatType))))
940+
(rule ((= e (InverseSqrt a))) ((set (ExprType e) (FloatType))))
941+
(rule ((= e (Exp a))) ((set (ExprType e) (FloatType))))
942+
(rule ((= e (Exp2 a))) ((set (ExprType e) (FloatType))))
943+
(rule ((= e (Log a))) ((set (ExprType e) (FloatType))))
944+
(rule ((= e (Log2 a))) ((set (ExprType e) (FloatType))))
945+
(rule ((= e (Pow a b))) ((set (ExprType e) (FloatType))))
946+
; GLSL math
947+
(rule ((= e (Fma a b c))) ((set (ExprType e) (FloatType))))
948+
(rule ((= e (Fract a))) ((set (ExprType e) (FloatType))))
949+
(rule ((= e (Modf a))) ((set (ExprType e) (FloatType))))
950+
(rule ((= e (Ldexp a b))) ((set (ExprType e) (FloatType))))
951+
(rule ((= e (Frexp a))) ((set (ExprType e) (FloatType))))
952+
(rule ((= e (Step a b))) ((set (ExprType e) (FloatType))))
953+
(rule ((= e (FSign a))) ((set (ExprType e) (FloatType))))
954+
(rule ((= e (Radians a))) ((set (ExprType e) (FloatType))))
955+
(rule ((= e (Degrees a))) ((set (ExprType e) (FloatType))))
956+
(rule ((= e (FMix a b c))) ((set (ExprType e) (FloatType))))
957+
(rule ((= e (SmoothStep a b c))) ((set (ExprType e) (FloatType))))
958+
(rule ((= e (FClamp a b c))) ((set (ExprType e) (FloatType))))
959+
(rule ((= e (NMin a b))) ((set (ExprType e) (FloatType))))
960+
(rule ((= e (NMax a b))) ((set (ExprType e) (FloatType))))
961+
(rule ((= e (NClamp a b c))) ((set (ExprType e) (FloatType))))
962+
; Geometry
963+
(rule ((= e (Length a))) ((set (ExprType e) (FloatType))))
964+
(rule ((= e (Distance a b))) ((set (ExprType e) (FloatType))))
965+
(rule ((= e (Dot a b))) ((set (ExprType e) (FloatType))))
966+
; Derivatives
967+
(rule ((= e (DPdx a))) ((set (ExprType e) (FloatType))))
968+
(rule ((= e (DPdy a))) ((set (ExprType e) (FloatType))))
969+
(rule ((= e (Fwidth a))) ((set (ExprType e) (FloatType))))
970+
(rule ((= e (DPdxFine a))) ((set (ExprType e) (FloatType))))
971+
(rule ((= e (DPdyFine a))) ((set (ExprType e) (FloatType))))
972+
(rule ((= e (FwidthFine a))) ((set (ExprType e) (FloatType))))
973+
(rule ((= e (DPdxCoarse a))) ((set (ExprType e) (FloatType))))
974+
(rule ((= e (DPdyCoarse a))) ((set (ExprType e) (FloatType))))
975+
(rule ((= e (FwidthCoarse a))) ((set (ExprType e) (FloatType))))
976+
; Unpacking (float result)
977+
(rule ((= e (UnpackSnorm2x16 a))) ((set (ExprType e) (FloatType))))
978+
(rule ((= e (UnpackUnorm2x16 a))) ((set (ExprType e) (FloatType))))
979+
(rule ((= e (UnpackHalf2x16 a))) ((set (ExprType e) (FloatType))))
980+
(rule ((= e (UnpackSnorm4x8 a))) ((set (ExprType e) (FloatType))))
981+
(rule ((= e (UnpackUnorm4x8 a))) ((set (ExprType e) (FloatType))))
982+
; Float constant
983+
(rule ((= e (FConst v))) ((set (ExprType e) (FloatType))))
984+
; Subgroup float reductions
985+
(rule ((= e (GroupFAdd a))) ((set (ExprType e) (FloatType))))
986+
(rule ((= e (GroupFMul a))) ((set (ExprType e) (FloatType))))
987+
(rule ((= e (GroupFMin a))) ((set (ExprType e) (FloatType))))
988+
(rule ((= e (GroupFMax a))) ((set (ExprType e) (FloatType))))
989+
; Quantization
990+
(rule ((= e (QuantizeToF16 a))) ((set (ExprType e) (FloatType))))
991+
992+
; --- Type-preserving operations (inherit ExprType from operands) ---
993+
; Gamma/Select/If inherit from branches
994+
(rule ((= e (Gamma c t f)) (= ty (ExprType t))) ((set (ExprType e) ty)))
995+
(rule ((= e (Gamma c t f)) (= ty (ExprType f))) ((set (ExprType e) ty)))
996+
(rule ((= e (Select c t f)) (= ty (ExprType t))) ((set (ExprType e) ty)))
997+
(rule ((= e (Select c t f)) (= ty (ExprType f))) ((set (ExprType e) ty)))
998+
(rule ((= e (If c t f)) (= ty (ExprType t))) ((set (ExprType e) ty)))
999+
(rule ((= e (If c t f)) (= ty (ExprType f))) ((set (ExprType e) ty)))
1000+
; CopyObject inherits from operand
1001+
(rule ((= e (CopyObject x)) (= ty (ExprType x))) ((set (ExprType e) ty)))
1002+
; Loop constructs inherit from body/init
1003+
(rule ((= e (LoopInvariant x)) (= ty (ExprType x))) ((set (ExprType e) ty)))
1004+
(rule ((= e (Theta c body init)) (= ty (ExprType init))) ((set (ExprType e) ty)))
1005+
7781006
; =============================================================================
7791007
; RVSDG Region Types
7801008
; =============================================================================

‎rust/spirv-tools-opt/src/rules/logical.egg‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -563,21 +563,24 @@
563563
; But: select(c, c && a, false) = c && a (proved by truth table)
564564
(rule ((= e (Gamma c (LogAnd c a) (BoolConst 0))))
565565
((union e (LogAnd c a))))
566-
(rule ((= e (Gamma c (LogAnd c a) (Const 0))))
566+
; Guard Const variant with ExprType to prevent int→bool domain crossing
567+
(rule ((= e (Gamma c (LogAnd c a) (Const 0))) (= (ExprType e) (BoolType)))
567568
((union e (LogAnd c a))))
568569

569570
; select(c, c || a, true) = true (c true means c||a=true; c false means result is true)
570571
(rule ((= e (Gamma c (LogOr c a) (BoolConst 1))))
571572
((union e (BoolConst 1))))
572-
(rule ((= e (Gamma c (LogOr c a) (Const 1))))
573+
; Guard Const variant with ExprType to prevent int→bool domain crossing
574+
(rule ((= e (Gamma c (LogOr c a) (Const 1))) (= (ExprType e) (BoolType)))
573575
((union e (BoolConst 1))))
574576

575577
; select(c, !c, x) = select(c, false, x) = !c && x via existing rules
576578
; But also: select(c, !c, true) = !c || c = true? No wait: c?!c:true = c?false:true = !c
577579
; Actually if c is true, !c is false; if c is false, we get true. So = !c
578580
(rule ((= e (Gamma c (LogNot c) (BoolConst 1))))
579581
((union e (LogNot c))))
580-
(rule ((= e (Gamma c (LogNot c) (Const 1))))
582+
; Guard Const variant with ExprType to prevent int→bool domain crossing
583+
(rule ((= e (Gamma c (LogNot c) (Const 1))) (= (ExprType e) (BoolType)))
581584
((union e (LogNot c))))
582585

583586
; select(c, true, c) = true (c?true:c = if c then true else c, but c is false so false = c)

‎rust/spirv-tools-opt/src/rules/rvsdg.egg‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -259,13 +259,14 @@
259259
; IMPORTANT: Use BoolConst (not Const) to prevent type confusion.
260260
; Integer (Const 0) must NOT be converted to LogAnd since LogAnd requires boolean operands.
261261
; ONE-DIRECTIONAL: Convert Gamma to LogAnd/LogOr, never reverse (prevents explosion)
262-
(rule ((= e (Gamma c x (BoolConst 0))))
262+
; Guard with ExprType to prevent converting non-bool x into LogAnd/LogOr operand
263+
(rule ((= e (Gamma c x (BoolConst 0))) (= (ExprType x) (BoolType)))
263264
((union e (LogAnd c x))))
264-
(rule ((= e (Gamma c (BoolConst 1) x)))
265+
(rule ((= e (Gamma c (BoolConst 1) x)) (= (ExprType x) (BoolType)))
265266
((union e (LogOr c x))))
266-
(rule ((= e (Gamma c (BoolConst 0) x)))
267+
(rule ((= e (Gamma c (BoolConst 0) x)) (= (ExprType x) (BoolType)))
267268
((union e (LogAnd (LogNot c) x))))
268-
(rule ((= e (Gamma c x (BoolConst 1))))
269+
(rule ((= e (Gamma c x (BoolConst 1))) (= (ExprType x) (BoolType)))
269270
((union e (LogOr (LogNot c) x))))
270271

271272
; Gamma with boolean constants 0/1 and variable
@@ -620,9 +621,10 @@
620621
; =============================================================================
621622

622623
; select(c, c && x, false) = c && x (boolean contexts)
624+
; Guard Const variant with ExprType to prevent int→bool domain crossing
623625
(rule ((= e (Gamma c (LogAnd c x) (BoolConst 0))))
624626
((union e (LogAnd c x))))
625-
(rule ((= e (Gamma c (LogAnd c x) (Const 0))))
627+
(rule ((= e (Gamma c (LogAnd c x) (Const 0))) (= (ExprType e) (BoolType)))
626628
((union e (LogAnd c x))))
627629

628630
; select(c, c || x, true) = true (when c is true, c||x is true; when false, true)

0 commit comments

Comments
 (0)