add modular add/sub/mul ops to arith_ext

lib/Dialect/ArithExt/IR/ArithExtOps.td

@@ -14,6 +14,35 @@ class ArithExt_Op<string mnemonic, list<Trait> traits = [Pure]> :
   let cppNamespace = "::mlir::heir::arith_ext";
 }
 
+class ArithExt_BinaryOp<string mnemonic, list<Trait> traits = []> :
+    ArithExt_Op<mnemonic, traits # [SameOperandsAndResultType, Pure, ElementwiseMappable]>,
+    Arguments<(ins SignlessIntegerLike:$lhs, SignlessIntegerLike:$rhs, I64Attr:$modulus)>,
+    Results<(outs SignlessIntegerLike:$output)> {
+  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` qualified(type($output))";
+}
+
+def ArithExt_AddOp : ArithExt_BinaryOp<"add", [Commutative]> {
+  let summary = "modular addition operation";
+  let description = [{
+    Computes addition modulo a statically known modulus $q$.
+  }];
+}
+
+def ArithExt_SubOp : ArithExt_BinaryOp<"sub"> {
+  let summary = "modular subtraction operation";
+  let description = [{
+    Computes subtraction modulo a statically known modulus $q$.
+  }];
+}
+
+def ArithExt_MulOp : ArithExt_BinaryOp<"mul", [Commutative]> {
+  let summary = "modular multiplication operation";
+  let description = [{
+    Computes multiplication modulo a statically known modulus $q$.
+  }];
+}
+
+
 def ArithExt_BarrettReduceOp : ArithExt_Op<"barrett_reduce", [SameOperandsAndResultType]> {
   let summary = "Compute the first step of the Barrett reduction.";
   let description = [{

tests/arith_ext/syntax.mlir

@@ -3,10 +3,28 @@
 // CHECK-LABEL: @test_arith_syntax
 func.func @test_arith_syntax() {
   %zero = arith.constant 1 : i10
+  %c5 = arith.constant 5 : i10
+  %c6 = arith.constant 6 : i10
   %cmod = arith.constant 17 : i10
   %c_vec = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi10>
+  %c_vec2 = arith.constant dense<[4,3,2,1]> : tensor<4xi10>
   %cmod_vec = arith.constant dense<17> : tensor<4xi10>
 
+  // CHECK: arith_ext.add
+  // CHECK: arith_ext.add
+  %add = arith_ext.add %c5, %c6 { modulus = 17 } : i10
+  %add_vec = arith_ext.add %c_vec, %c_vec2 { modulus = 17 } : tensor<4xi10>
+
+  // CHECK: arith_ext.sub
+  // CHECK: arith_ext.sub
+  %sub = arith_ext.sub %c5, %c6 { modulus = 17 } : i10
+  %sub_vec = arith_ext.sub %c_vec, %c_vec2 { modulus = 17 } : tensor<4xi10>
+
+  // CHECK: arith_ext.mul
+  // CHECK: arith_ext.mul
+  %mul = arith_ext.mul %c5, %c6 { modulus = 17 } : i10
+  %mul_vec = arith_ext.mul %c_vec, %c_vec2 { modulus = 17 } : tensor<4xi10>
+
   // CHECK: arith_ext.barrett_reduce
   // CHECK: arith_ext.barrett_reduce
   %barrett = arith_ext.barrett_reduce %zero { modulus = 17 } : i10