Dispatch vector arithmetic to Vec* constructors in the egraph

SPIR-V OpFAdd/OpIAdd/etc. work on both scalars and vectors, but the
egglog typed sorts (FloatExpr, IntExpr) are scalar-only. When an
operand was vector-typed, the sort mismatch caused egglog failures.

Following the existing Select/VecSelect and FMix/VecFMix pattern,
dispatch to Vec* constructors (Expr sort) for vector operations while
keeping typed sorts for scalars. This keeps everything in the egraph
for one big saturation pass.

Add VecFRem, VecFMod, VecSDiv, VecUDiv, VecSRem, VecSMod, VecUMod
constructors with DCE, VecSize, LICM rules, and reconstruction.

Author: LegNeato

rust/spirv-tools-opt/src/direct/context.rs

@@ -217,15 +217,15 @@ impl EgglogContext {
             Op::ConstantTrue => "(BoolConst 1)".to_string(),
             Op::ConstantFalse => "(BoolConst 0)".to_string(),
 
-            Op::IAdd => self.binary_op("Add", inst)?,
-            Op::ISub => self.binary_op("Sub", inst)?,
-            Op::IMul => self.binary_op("Mul", inst)?,
-            Op::SDiv => self.binary_op("SDiv", inst)?,
-            Op::UDiv => self.binary_op("UDiv", inst)?,
-            Op::SRem => self.binary_op("SRem", inst)?,
-            Op::SMod => self.binary_op("SMod", inst)?,
-            Op::UMod => self.binary_op("UMod", inst)?,
-            Op::SNegate => self.unary_op("Neg", inst)?,
+            Op::IAdd => self.typed_binary_op("Add", "VecAdd", inst)?,
+            Op::ISub => self.typed_binary_op("Sub", "VecSub", inst)?,
+            Op::IMul => self.typed_binary_op("Mul", "VecMul", inst)?,
+            Op::SDiv => self.typed_binary_op("SDiv", "VecSDiv", inst)?,
+            Op::UDiv => self.typed_binary_op("UDiv", "VecUDiv", inst)?,
+            Op::SRem => self.typed_binary_op("SRem", "VecSRem", inst)?,
+            Op::SMod => self.typed_binary_op("SMod", "VecSMod", inst)?,
+            Op::UMod => self.typed_binary_op("UMod", "VecUMod", inst)?,
+            Op::SNegate => self.typed_unary_op("Neg", "VecNeg", inst)?,
             Op::ShiftLeftLogical => self.binary_op("Shl", inst)?,
             Op::ShiftRightLogical => self.binary_op("ShrU", inst)?,
             Op::ShiftRightArithmetic => self.binary_op("ShrS", inst)?,
@@ -249,14 +249,14 @@ impl EgglogContext {
             Op::LogicalOr => self.binary_op("LogOr", inst)?,
             Op::LogicalEqual => self.binary_op("LogEq", inst)?,
             Op::LogicalNotEqual => self.binary_op("LogNe", inst)?,
-            // Floating-point operations
-            Op::FAdd => self.binary_op("FAdd", inst)?,
-            Op::FSub => self.binary_op("FSub", inst)?,
-            Op::FMul => self.binary_op("FMul", inst)?,
-            Op::FDiv => self.binary_op("FDiv", inst)?,
-            Op::FRem => self.binary_op("FRem", inst)?,
-            Op::FMod => self.binary_op("FMod", inst)?,
-            Op::FNegate => self.unary_op("FNeg", inst)?,
+            // Floating-point operations (scalar or vector)
+            Op::FAdd => self.typed_binary_op("FAdd", "VecFAdd", inst)?,
+            Op::FSub => self.typed_binary_op("FSub", "VecFSub", inst)?,
+            Op::FMul => self.typed_binary_op("FMul", "VecFMul", inst)?,
+            Op::FDiv => self.typed_binary_op("FDiv", "VecFDiv", inst)?,
+            Op::FRem => self.typed_binary_op("FRem", "VecFRem", inst)?,
+            Op::FMod => self.typed_binary_op("FMod", "VecFMod", inst)?,
+            Op::FNegate => self.typed_unary_op("FNeg", "VecFNeg", inst)?,
             // Floating-point comparisons (ordered)
             Op::FOrdEqual => self.binary_op("FOrdEq", inst)?,
             Op::FOrdNotEqual => self.binary_op("FOrdNe", inst)?,
@@ -930,6 +930,46 @@ impl EgglogContext {
         Some(format!("({} {})", op, operand))
     }
 
+    /// Type-dispatched binary op: uses `scalar_op` (typed sort) when the result
+    /// type is a scalar, and `vec_op` (Expr sort) when it is a vector/other type.
+    /// This prevents sort mismatches when SPIR-V opcodes like OpFAdd/OpIAdd
+    /// operate on vectors whose operands are in the generic Expr sort.
+    fn typed_binary_op(
+        &mut self,
+        scalar_op: &str,
+        vec_op: &str,
+        inst: &Instruction,
+    ) -> Option<String> {
+        let is_scalar = inst
+            .result_type
+            .map(|ty| self.type_class_of_type(ty) != TypeClass::Other)
+            .unwrap_or(false);
+        if is_scalar {
+            self.binary_op(scalar_op, inst)
+        } else {
+            self.binary_op(vec_op, inst)
+        }
+    }
+
+    /// Type-dispatched unary op: uses `scalar_op` (typed sort) when the result
+    /// type is a scalar, and `vec_op` (Expr sort) when it is a vector/other type.
+    fn typed_unary_op(
+        &mut self,
+        scalar_op: &str,
+        vec_op: &str,
+        inst: &Instruction,
+    ) -> Option<String> {
+        let is_scalar = inst
+            .result_type
+            .map(|ty| self.type_class_of_type(ty) != TypeClass::Other)
+            .unwrap_or(false);
+        if is_scalar {
+            self.unary_op(scalar_op, inst)
+        } else {
+            self.unary_op(vec_op, inst)
+        }
+    }
+
     /// Convert GLSL.std.450 extended instruction to egglog term.
     fn extended_instruction_to_term(&mut self, inst: &Instruction) -> Option<String> {
         // ExtInst operands: %set %instruction operands...

rust/spirv-tools-opt/src/direct/mod.rs

@@ -2454,19 +2454,24 @@ fn instruction_has_valid_types(
     let op = inst.class.opcode;
 
     // Check result type: if the opcode requires a specific type class,
-    // the result type MUST match (no TypeClass::Other escape).
+    // the result type MUST match OR be TypeClass::Other (vectors/matrices).
+    // SPIR-V arithmetic ops (FAdd, IAdd, etc.) work on both scalars and vectors.
+    // TypeClass::Other means the type is a vector/matrix/struct which is valid
+    // for these component-wise operations.
     if let (Some(required), Some(result_type)) = (required_result_type_class(op), inst.result_type)
     {
         let actual = type_classes
             .get(&result_type)
             .copied()
             .unwrap_or(TypeClass::Other);
-        if actual != required {
+        if actual != required && actual != TypeClass::Other {
             return false;
         }
     }
 
     // Check operand types for comparisons
+    // Again, TypeClass::Other (vector operands) is always acceptable since
+    // SPIR-V comparison ops work component-wise on vectors.
     if let Some(required_op_class) = required_operand_type_class(op) {
         for operand in &inst.operands {
             if let Some(operand_id) = operand.id_ref_any() {
@@ -2475,7 +2480,7 @@ fn instruction_has_valid_types(
                         .get(&operand_type)
                         .copied()
                         .unwrap_or(TypeClass::Other);
-                    if actual != required_op_class {
+                    if actual != required_op_class && actual != TypeClass::Other {
                         return false;
                     }
                 }

rust/spirv-tools-opt/src/direct/parse/vector.rs

@@ -5,7 +5,8 @@ use rspirv::spirv::{Op, Word};
 use std::collections::HashMap;
 
 use super::util::{
-    parse_binary_args, parse_expr_list, parse_ternary_args, resolve_term_to_id, split_terms,
+    parse_binary_args, parse_expr_list, parse_ternary_args, parse_unary_arg, resolve_term_to_id,
+    split_terms,
 };
 
 /// Matrix binary operations.
@@ -21,6 +22,36 @@ const MATRIX_BINARY_OPS: &[(&str, Op)] = &[
 /// Matrix unary operations.
 const MATRIX_UNARY_OPS: &[(&str, Op)] = &[("Transpose", Op::Transpose)];
 
+/// Vector arithmetic binary operations (component-wise).
+/// These are the Expr-sort constructors for ops that work on both scalars and vectors.
+const VEC_ARITHMETIC_BINARY_OPS: &[(&str, Op)] = &[
+    // Integer vector arithmetic
+    ("VecAdd", Op::IAdd),
+    ("VecSub", Op::ISub),
+    ("VecMul", Op::IMul),
+    ("VecDiv", Op::SDiv),
+    ("VecSDiv", Op::SDiv),
+    ("VecUDiv", Op::UDiv),
+    ("VecSRem", Op::SRem),
+    ("VecSMod", Op::SMod),
+    ("VecUMod", Op::UMod),
+    // Float vector arithmetic
+    ("VecFAdd", Op::FAdd),
+    ("VecFSub", Op::FSub),
+    ("VecFMul", Op::FMul),
+    ("VecFDiv", Op::FDiv),
+    ("VecFRem", Op::FRem),
+    ("VecFMod", Op::FMod),
+    // Scalar-vector multiply
+    ("VecTimesScalar", Op::VectorTimesScalar),
+];
+
+/// Vector arithmetic unary operations (component-wise).
+const VEC_ARITHMETIC_UNARY_OPS: &[(&str, Op)] = &[
+    ("VecNeg", Op::SNegate),
+    ("VecFNeg", Op::FNegate),
+];
+
 /// Try to parse a vector or composite operation.
 pub fn try_parse_vector(
     term: &str,
@@ -63,6 +94,39 @@ pub fn try_parse_vector(
         }
     }
 
+    // Vector arithmetic binary operations (component-wise)
+    for (name, opcode) in VEC_ARITHMETIC_BINARY_OPS {
+        let prefix = format!("({} ", name);
+        if let Some(rest) = term.strip_prefix(&prefix) {
+            if let Some((lhs, rhs)) = parse_binary_args(rest, id_map) {
+                return Some(Instruction::new(
+                    *opcode,
+                    Some(result_type),
+                    Some(result_id),
+                    vec![
+                        rspirv::dr::Operand::IdRef(lhs),
+                        rspirv::dr::Operand::IdRef(rhs),
+                    ],
+                ));
+            }
+        }
+    }
+
+    // Vector arithmetic unary operations (component-wise)
+    for (name, opcode) in VEC_ARITHMETIC_UNARY_OPS {
+        let prefix = format!("({} ", name);
+        if let Some(rest) = term.strip_prefix(&prefix) {
+            if let Some(operand) = parse_unary_arg(rest, id_map) {
+                return Some(Instruction::new(
+                    *opcode,
+                    Some(result_type),
+                    Some(result_id),
+                    vec![rspirv::dr::Operand::IdRef(operand)],
+                ));
+            }
+        }
+    }
+
     // VectorInsertDynamic (ternary: vector, component, index)
     if let Some(rest) = term.strip_prefix("(VectorInsertDynamic ") {
         if let Some((vec, component, idx)) = parse_ternary_args(rest, id_map) {

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

@@ -393,6 +393,13 @@
 (constructor VecFMul (Expr Expr) Expr)
 (constructor VecFDiv (Expr Expr) Expr)
 (constructor VecFNeg (Expr) Expr)
+(constructor VecFRem (Expr Expr) Expr)
+(constructor VecFMod (Expr Expr) Expr)
+(constructor VecSDiv (Expr Expr) Expr)
+(constructor VecUDiv (Expr Expr) Expr)
+(constructor VecSRem (Expr Expr) Expr)
+(constructor VecSMod (Expr Expr) Expr)
+(constructor VecUMod (Expr Expr) Expr)
 (constructor VecTimesScalar (Expr Expr) Expr)
 
 ; Vector-level GLSL ops (vector variants of scalar FloatExpr constructors)
@@ -1103,6 +1110,13 @@
 (rule ((Live e) (= e (VecFMul a b))) ((Live a) (Live b)))
 (rule ((Live e) (= e (VecFDiv a b))) ((Live a) (Live b)))
 (rule ((Live e) (= e (VecFNeg a))) ((Live a)))
+(rule ((Live e) (= e (VecFRem a b))) ((Live a) (Live b)))
+(rule ((Live e) (= e (VecFMod a b))) ((Live a) (Live b)))
+(rule ((Live e) (= e (VecSDiv a b))) ((Live a) (Live b)))
+(rule ((Live e) (= e (VecUDiv a b))) ((Live a) (Live b)))
+(rule ((Live e) (= e (VecSRem a b))) ((Live a) (Live b)))
+(rule ((Live e) (= e (VecSMod a b))) ((Live a) (Live b)))
+(rule ((Live e) (= e (VecUMod a b))) ((Live a) (Live b)))
 (rule ((Live e) (= e (VecTimesScalar v s))) ((Live v) (Live s)))
 (rule ((Live e) (= e (VecFMix a b c))) ((Live a) (Live b) (Live c)))
 (rule ((Live e) (= e (VecSelect mask t f))) ((Live mask) (Live t) (Live f)))

rust/spirv-tools-opt/src/rules/vector.egg

@@ -619,6 +619,14 @@
 (rule ((= e (VecFSub (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (VecFSub a b)))))
 (rule ((= e (VecFMul (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (VecFMul a b)))))
 (rule ((= e (VecFNeg (LoopInvariant a)))) ((union e (LoopInvariant (VecFNeg a)))))
+(rule ((= e (VecFRem (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (VecFRem a b)))))
+(rule ((= e (VecFMod (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (VecFMod a b)))))
+; Note: VecDiv and VecFDiv loop invariant rules are in licm.egg
+(rule ((= e (VecSDiv (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (VecSDiv a b)))))
+(rule ((= e (VecUDiv (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (VecUDiv a b)))))
+(rule ((= e (VecSRem (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (VecSRem a b)))))
+(rule ((= e (VecSMod (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (VecSMod a b)))))
+(rule ((= e (VecUMod (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (VecUMod a b)))))
 ; Note: Dot returns FloatExpr, use LoopInvariantF (FloatExpr -> FloatExpr)
 (rule ((= e (Dot (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariantF (Dot a b)))))
 (rule ((= e (Cross (LoopInvariant a) (LoopInvariant b)))) ((union e (LoopInvariant (Cross a b)))))
@@ -649,6 +657,14 @@
 (rule ((= e (VecFSub a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
 (rule ((= e (VecFMul a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
 (rule ((= e (VecFNeg a)) (= sz (VecSize a))) ((set (VecSize e) sz)))
+(rule ((= e (VecFRem a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
+(rule ((= e (VecFMod a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
+(rule ((= e (VecDiv a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
+(rule ((= e (VecSDiv a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
+(rule ((= e (VecUDiv a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
+(rule ((= e (VecSRem a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
+(rule ((= e (VecSMod a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
+(rule ((= e (VecUMod a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
 (rule ((= e (VecTimesScalar a _)) (= sz (VecSize a))) ((set (VecSize e) sz)))
 (rule ((= e (VecFMix a _ _)) (= sz (VecSize a))) ((set (VecSize e) sz)))