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: 1019376409

Author: j2kun

lib/Dialect/CheddarGPU/IR/CheddarGPUOps.td

@@ -22,12 +22,22 @@ 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,
@@ -61,6 +71,27 @@ defvar CheddarGPU_BinaryBuilders = [
   }]>
 ];
 
+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{});
+  }]>
+];
+
 //===----------------------------------------------------------------------===//
 // Operation Families
 //===----------------------------------------------------------------------===//
@@ -89,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
 //===----------------------------------------------------------------------===//
@@ -121,6 +176,20 @@ def CheddarGPU_SubOp : CheddarGPU_BinaryOpBase<"sub", [
   }];
 }
 
+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";

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/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
+  }
+
+}