Add cheddar_gpu neg op

Adds pointwise modular negation (cheddar_gpu.neg) and the unary op
base class to the cheddar_gpu dialect.

PiperOrigin-RevId: 1015875353

Author: j2kun

lib/Dialect/CheddarGPU/IR/CheddarGPUConstraints.td

@@ -29,6 +29,11 @@ def CheddarGPU_PolyBuffer : CheddarGPU_StaticMemRefOfRank<[I32, I64], 2,
 class CheddarGPU_FormsMatch<list<string> a>
     : AllMatchSameOperatorTrait<a, "$_self", "form">;
 
+class CheddarGPU_FormsAre<list<string> attrs, EnumCase f>
+    : PredOpTrait<
+        "{" # !interleave(attrs, ", ") # "} are '" # f.str # "'",
+        And<!foreach(a, attrs, CPred<"$" # a # " == ::mlir::heir::cheddar_gpu::Form::" # f.symbol>)>>;
+
 //===----------------------------------------------------------------------===//
 // Attribute Constraints
 //===----------------------------------------------------------------------===//

lib/Dialect/CheddarGPU/IR/CheddarGPUOps.td

@@ -22,34 +22,73 @@ class CheddarGPU_Op<string mnemonic, list<Trait> traits = []>
   }];
 }
 
+defvar CheddarGPU_UnaryForms = (ins
+  CheddarGPU_FormAttr:$in_form,
+  CheddarGPU_FormAttr:$out_form
+);
+
 defvar CheddarGPU_BinaryForms = (ins
   CheddarGPU_FormAttr:$lhs_form,
   CheddarGPU_FormAttr:$rhs_form,
   CheddarGPU_FormAttr:$out_form
 );
 
+defvar CheddarGPU_UnaryLimbRanges = (ins
+  OptionalAttr<CheddarGPU_LimbRangeAttr>:$in_limb_range,
+  OptionalAttr<CheddarGPU_LimbRangeAttr>:$out_limb_range
+);
+
 defvar CheddarGPU_BinaryLimbRanges = (ins
   OptionalAttr<CheddarGPU_LimbRangeAttr>:$lhs_limb_range,
   OptionalAttr<CheddarGPU_LimbRangeAttr>:$rhs_limb_range,
   OptionalAttr<CheddarGPU_LimbRangeAttr>:$out_limb_range
 );
 
 defvar CheddarGPU_BinaryBuilders = [
-  OpBuilder<(ins "Value":$out, "Value":$lhs, "Value":$rhs,
-                 "FormAttr":$lhs_form, "FormAttr":$rhs_form, "FormAttr":$out_form,
-                 CArg<"DenseI64ArrayAttr", "{}">:$lhs_limb_range,
-                 CArg<"DenseI64ArrayAttr", "{}">:$rhs_limb_range,
-                 CArg<"DenseI64ArrayAttr", "{}">:$out_limb_range), [{
+  OpBuilder<(ins "::mlir::Value":$out,
+                 "::mlir::Value":$lhs,
+                 "::mlir::Value":$rhs,
+                 "::mlir::heir::cheddar_gpu::FormAttr":$lhs_form,
+                 "::mlir::heir::cheddar_gpu::FormAttr":$rhs_form,
+                 "::mlir::heir::cheddar_gpu::FormAttr":$out_form,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$lhs_limb_range,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$rhs_limb_range,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$out_limb_range), [{
       build($_builder, $_state, out, lhs, rhs, lhs_form, rhs_form, out_form,
-            lhs_limb_range, rhs_limb_range, out_limb_range, /*num_aux=*/IntegerAttr{});
+            lhs_limb_range, rhs_limb_range, out_limb_range, /*num_aux=*/::mlir::IntegerAttr{});
   }]>,
-  OpBuilder<(ins "Value":$out, "Value":$lhs, "Value":$rhs,
-                 "Form":$lhs_form, "Form":$rhs_form, "Form":$out_form,
-                 CArg<"DenseI64ArrayAttr", "{}">:$lhs_limb_range,
-                 CArg<"DenseI64ArrayAttr", "{}">:$rhs_limb_range,
-                 CArg<"DenseI64ArrayAttr", "{}">:$out_limb_range), [{
+  OpBuilder<(ins "::mlir::Value":$out,
+                 "::mlir::Value":$lhs,
+                 "::mlir::Value":$rhs,
+                 "::mlir::heir::cheddar_gpu::Form":$lhs_form,
+                 "::mlir::heir::cheddar_gpu::Form":$rhs_form,
+                 "::mlir::heir::cheddar_gpu::Form":$out_form,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$lhs_limb_range,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$rhs_limb_range,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$out_limb_range), [{
       build($_builder, $_state, out, lhs, rhs, lhs_form, rhs_form, out_form,
-            lhs_limb_range, rhs_limb_range, out_limb_range, /*num_aux=*/IntegerAttr{});
+            lhs_limb_range, rhs_limb_range, out_limb_range, /*num_aux=*/::mlir::IntegerAttr{});
+  }]>
+];
+
+defvar CheddarGPU_UnaryBuilders = [
+  OpBuilder<(ins "::mlir::Value":$out,
+                 "::mlir::Value":$in,
+                 "::mlir::heir::cheddar_gpu::FormAttr":$in_form,
+                 "::mlir::heir::cheddar_gpu::FormAttr":$out_form,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$in_limb_range,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$out_limb_range), [{
+      build($_builder, $_state, out, in, in_form, out_form,
+            in_limb_range, out_limb_range, /*num_aux=*/::mlir::IntegerAttr{});
+  }]>,
+  OpBuilder<(ins "::mlir::Value":$out,
+                 "::mlir::Value":$in,
+                 "::mlir::heir::cheddar_gpu::Form":$in_form,
+                 "::mlir::heir::cheddar_gpu::Form":$out_form,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$in_limb_range,
+                 CArg<"::mlir::DenseI64ArrayAttr", "{}">:$out_limb_range), [{
+      build($_builder, $_state, out, in, in_form, out_form,
+            in_limb_range, out_limb_range, /*num_aux=*/::mlir::IntegerAttr{});
   }]>
 ];
 
@@ -81,6 +120,30 @@ class CheddarGPU_BinaryOpBase<string mnemonic, list<Trait> traits = []>
   }];
 }
 
+// Unary element-wise operations on single polynomial buffers (neg).
+class CheddarGPU_UnaryOpBase<string mnemonic, dag extraAttrs = (ins),
+                             list<OpBuilder> opBuilders = [], list<Trait> traits = []>
+    : CheddarGPU_Op<mnemonic, !listconcat([
+        AllTypesMatch<["out", "in"]>,
+        CheddarGPU_LimbRangesWithin
+      ], traits)> {
+  let arguments = !con(
+    (ins
+      Arg<CheddarGPU_PolyBuffer, "destination buffer", [MemWrite]>:$out,
+      Arg<CheddarGPU_PolyBuffer, "input polynomial buffer", [MemRead]>:$in
+    ),
+    CheddarGPU_UnaryForms,
+    extraAttrs,
+    CheddarGPU_UnaryLimbRanges,
+    (ins OptionalAttr<CheddarGPU_NonNegI64Attr>:$num_aux)
+  );
+
+  let builders = opBuilders;
+  let assemblyFormat = [{
+    $out `,` $in prop-dict attr-dict `:` type($out)
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Binary and Unary Element-Wise Operations
 //===----------------------------------------------------------------------===//
@@ -99,4 +162,46 @@ def CheddarGPU_AddOp : CheddarGPU_BinaryOpBase<"add", [
   }];
 }
 
+def CheddarGPU_SubOp : CheddarGPU_BinaryOpBase<"sub", [
+    CheddarGPU_FormsMatch<["lhs_form", "rhs_form", "out_form"]>]> {
+  let summary = "Pointwise modular subtraction of two polynomial buffers";
+  let description = [{
+    Mirrors Cheddar's `kernel::Sub` and `ElementWiseHandler::Sub` in `ElementWise.cu`.
+    Computes `out = lhs - rhs mod Q` limb-wise, where the limbs of `Q` are defined
+    in `CheddarGPU_ParamsAttr`.
+
+    Form contract:
+    - Modular subtraction is valid in both `coeff` and `eval` domain.
+    - Both operands must have matching form (`lhs_form == rhs_form == out_form`).
+  }];
+}
+
+def CheddarGPU_NegOp : CheddarGPU_UnaryOpBase<"neg", (ins), CheddarGPU_UnaryBuilders, [
+    CheddarGPU_FormsMatch<["in_form", "out_form"]>]> {
+  let summary = "Pointwise modular negation of a polynomial buffer";
+  let description = [{
+    Mirrors Cheddar's `kernel::Neg` and `ElementWiseHandler::Neg` in `ElementWise.cu`.
+    Computes `out = -in mod Q` limb-wise, where the limbs of `Q` are defined
+    in `CheddarGPU_ParamsAttr`.
+
+    Form contract:
+    - Valid in both `coeff` and `eval` domain.
+    - Preserves representation form (`in_form == out_form`).
+  }];
+}
+
+def CheddarGPU_MulOp : CheddarGPU_BinaryOpBase<"mul", [
+    CheddarGPU_FormsAre<["lhs_form", "rhs_form", "out_form"], CheddarGPU_Form_Eval>]> {
+  let summary = "Pointwise Montgomery modular multiplication";
+  let description = [{
+    Mirrors Cheddar's `kernel::Mult` and `ElementWiseHandler::Mult` in `ElementWise.cu`.
+    Computes the Montgomery form multiplication of `lhs` and `rhs` mod `Q`.
+
+    Form contract:
+    - Polynomial multiplication in evaluation domain is pointwise Montgomery multiplication:
+      `(a * R) * (b * R) * R^-1 = (a * b) * R mod Q`.
+    - Operands must be in `eval` form (`lhs_form == rhs_form == out_form == eval`).
+  }];
+}
+
 #endif  // LIB_DIALECT_CHEDDARGPU_IR_CHEDDARGPU_OPS_TD_

lib/Dialect/CheddarGPU/IR/CheddarGPUTraits.cpp

@@ -73,6 +73,7 @@ LogicalResult verifyLimbRangesWithin(Operation* op) {
 
   if (failed(checkRange(0, "out_limb_range"))) return failure();
   if (failed(checkRange(1, "lhs_limb_range"))) return failure();
+  if (failed(checkRange(1, "in_limb_range"))) return failure();
   if (failed(checkRange(2, "rhs_limb_range"))) return failure();
 
   if (auto numAuxAttr = op->getAttrOfType<IntegerAttr>("num_aux")) {

tests/Dialect/CheddarGPU/IR/binary_ops.mlir

@@ -31,4 +31,26 @@ module attributes {
     return
   }
 
+  // CHECK: @test_sub
+  func.func @test_sub(%out: memref<4x65536xi32>, %lhs: memref<4x65536xi32>, %rhs: memref<4x65536xi32>) {
+    // CHECK: cheddar_gpu.sub %{{.*}}, %{{.*}}, %{{.*}} <{{.*}}lhs_form = eval{{.*}}> : memref<4x65536xi32>
+    cheddar_gpu.sub %out, %lhs, %rhs <{
+      lhs_form = #cheddar_gpu.form<eval>,
+      rhs_form = #cheddar_gpu.form<eval>,
+      out_form = #cheddar_gpu.form<eval>
+    }> : memref<4x65536xi32>
+    return
+  }
+
+  // CHECK: @test_mul
+  func.func @test_mul(%out: memref<4x65536xi32>, %lhs: memref<4x65536xi32>, %rhs: memref<4x65536xi32>) {
+    // CHECK: cheddar_gpu.mul %{{.*}}, %{{.*}}, %{{.*}} <{{.*}}lhs_form = eval{{.*}}> : memref<4x65536xi32>
+    cheddar_gpu.mul %out, %lhs, %rhs <{
+      lhs_form = #cheddar_gpu.form<eval>,
+      rhs_form = #cheddar_gpu.form<eval>,
+      out_form = #cheddar_gpu.form<eval>
+    }> : memref<4x65536xi32>
+    return
+  }
+
 }

tests/Dialect/CheddarGPU/IR/binary_verifier.mlir

@@ -71,3 +71,15 @@ func.func @test_num_aux_oob(%out: memref<4x65536xi32>, %lhs: memref<4x65536xi32>
   }> : memref<4x65536xi32>
   return
 }
+// -----
+// Mul with coeff operand (requires eval)
+
+func.func @test_mul_with_coeff(%out: memref<4x65536xi32>, %lhs: memref<4x65536xi32>, %rhs: memref<4x65536xi32>) {
+  // expected-error @below {{failed to verify that {lhs_form, rhs_form, out_form} are 'eval'}}
+  cheddar_gpu.mul %out, %lhs, %rhs <{
+    lhs_form = #cheddar_gpu.form<coeff>,
+    rhs_form = #cheddar_gpu.form<eval>,
+    out_form = #cheddar_gpu.form<eval>
+  }> : memref<4x65536xi32>
+  return
+}

tests/Dialect/CheddarGPU/IR/unary_ops.mlir

@@ -0,0 +1,19 @@
+// RUN: heir-opt --verify-roundtrip %s | FileCheck %s
+
+// CHECK: module attributes
+// CHECK-SAME: cheddar_gpu.params = #cheddar_gpu.params<word = 32, log_n = 16, ter = [1000], main = [2000, 3000], aux = [4000], level_config = {{.*}}, alpha = 1, dnum = 0>
+module attributes {
+  cheddar_gpu.params = #cheddar_gpu.params<word = 32, log_n = 16, ter = [1000], main = [2000, 3000], aux = [4000], level_config = [[[2 : i32, 1 : i32]]], alpha = 1, dnum = 0>
+} {
+
+  // CHECK: @test_neg
+  func.func @test_neg(%out: memref<4x65536xi32>, %in: memref<4x65536xi32>) {
+    // CHECK: cheddar_gpu.neg %{{.*}}, %{{.*}} <{{.*}}in_form = eval{{.*}}> : memref<4x65536xi32>
+    cheddar_gpu.neg %out, %in <{
+      in_form = #cheddar_gpu.form<eval>,
+      out_form = #cheddar_gpu.form<eval>
+    }> : memref<4x65536xi32>
+    return
+  }
+
+}