misc changes to support optimal bufferization

PiperOrigin-RevId: 931403381

Author: j2kun

lib/Dialect/Lattigo/IR/LattigoBGVOps.td

@@ -13,7 +13,7 @@ class Lattigo_BGVOp<string mnemonic, list<Trait> traits = []> :
 // This operation cannot be marked as Pure because it represents a buffer allocation.
 // If marked as Pure, CSE will deduplicate them, causing multiple in-place operations
 // (like encode) to share the same buffer and overwrite each other.
-def Lattigo_BGVNewPlaintextOp : Lattigo_BGVOp<"new_plaintext", []> {
+def Lattigo_BGVNewPlaintextOp : Lattigo_BGVOp<"new_plaintext"> {
   let summary = "Create a new plaintext in the Lattigo BGV dialect";
   let description = [{
     This operation creates a new plaintext value in the Lattigo BGV dialect.

lib/Dialect/Lattigo/IR/LattigoCKKSOps.td

@@ -15,7 +15,7 @@ class Lattigo_CKKSOp<string mnemonic, list<Trait> traits = []> :
 // This operation cannot be marked as Pure because it represents a buffer allocation.
 // If marked as Pure, CSE will deduplicate them, causing multiple in-place operations
 // (like encode) to share the same buffer and overwrite each other.
-def Lattigo_CKKSNewPlaintextOp : Lattigo_CKKSOp<"new_plaintext", []> {
+def Lattigo_CKKSNewPlaintextOp : Lattigo_CKKSOp<"new_plaintext"> {
   let summary = "Create a new plaintext in the Lattigo CKKS dialect";
   let description = [{
     This operation creates a new plaintext value in the Lattigo CKKS dialect.
@@ -48,7 +48,7 @@ def Lattigo_CKKSNewEncoderOp : Lattigo_CKKSOp<"new_encoder", [Pure]> {
   let results = (outs Lattigo_CKKSEncoder:$encoder);
 }
 
-def Lattigo_CKKSEncodeOp : Lattigo_CKKSOp<"encode", [InPlaceOpInterface, Pure, PlaintextEncodeOpInterface]> {
+def Lattigo_CKKSEncodeOp : Lattigo_CKKSOp<"encode", [InPlaceOpInterface, PlaintextEncodeOpInterface]> {
   let summary = "Encode a plaintext value in the Lattigo CKKS dialect";
   let description = [{
     This operation encodes a plaintext value using the specified encoder in the Lattigo CKKS dialect.
@@ -68,7 +68,7 @@ def Lattigo_CKKSEncodeOp : Lattigo_CKKSOp<"encode", [InPlaceOpInterface, Pure, P
   let extraClassDeclaration = "int getInPlaceOperandIndex() { return 2; }";
 }
 
-def Lattigo_CKKSDecodeOp : Lattigo_CKKSOp<"decode", [AllTypesMatch<["value", "decoded"]>, Pure]> {
+def Lattigo_CKKSDecodeOp : Lattigo_CKKSOp<"decode", [AllTypesMatch<["value", "decoded"]>]> {
   let summary = "Decode a plaintext value in the Lattigo CKKS dialect";
   let description = [{
     This operation decodes a plaintext value using the specified encoder in the Lattigo CKKS dialect.

lib/Dialect/TensorExt/Conversions/TensorExtToTensor/TensorExtToTensor.cpp

@@ -155,11 +155,8 @@ struct ConvertRotateOp : public OpRewritePattern<RotateOp> {
     leftSliceOffsets[leftSliceOffsets.size() - 1] = dimMinusShift;
     rightSliceOffsets[rightSliceOffsets.size() - 1] = b.getIndexAttr(0);
 
-    auto empty =
-        tensor::EmptyOp::create(rewriter, op.getLoc(), tensorType.getShape(),
-                                tensorType.getElementType());
     auto insertedLeftSlice = tensor::InsertSliceOp::create(
-        b, left.getResult(), empty, leftSliceOffsets, leftSliceSizes,
+        b, left.getResult(), op.getTensor(), leftSliceOffsets, leftSliceSizes,
         allOneStrides);
     auto insertedRightSlice = tensor::InsertSliceOp::create(
         b, right.getResult(), insertedLeftSlice.getResult(), rightSliceOffsets,

lib/Kernel/BUILD

@@ -134,6 +134,7 @@ cc_test(
         ":AbstractValue",
         ":ArithmeticDag",
         ":EvalVisitor",
+        ":IRMaterializingVisitor",
         ":Kernel",
         ":KernelImplementation",
         ":RotationCountVisitor",
@@ -142,8 +143,11 @@ cc_test(
         "@heir//lib/Utils/Layout:Convolution",
         "@heir//lib/Utils/Layout:Evaluate",
         "@heir//lib/Utils/Layout:Utils",
+        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
     ],
 )
 

lib/Kernel/IRMaterializingVisitor.cpp

@@ -4,6 +4,7 @@
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
+#include <utility>
 #include <variant>
 #include <vector>
 

lib/Pipelines/BUILD

@@ -109,7 +109,6 @@ cc_library(
         "@heir//lib/Dialect/LWE/Conversions/LWEToPolynomial",
         "@heir//lib/Dialect/LWE/Transforms:AddDebugPort",
         "@heir//lib/Dialect/LWE/Transforms:ImplementTrivialEncryptionAsAddition",
-        "@heir//lib/Dialect/Lattigo/Transforms:AllocToInPlace",
         "@heir//lib/Dialect/Lattigo/Transforms:ConfigureCryptoContext",
         "@heir//lib/Dialect/Openfhe/Transforms:AllocToInPlace",
         "@heir//lib/Dialect/Openfhe/Transforms:ConfigureCryptoContext",

lib/Target/Lattigo/LattigoEmitter.cpp

@@ -34,6 +34,7 @@
 #include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"   // from @llvm-project
+#include "mlir/include/mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/SCF/IR/SCF.h"        // from @llvm-project
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Utils/StaticValueUtils.h"  // from @llvm-project
@@ -119,8 +120,9 @@ LogicalResult LattigoEmitter::translate(Operation& op) {
           .Case<memref::AllocOp, memref::LoadOp, memref::StoreOp,
                 memref::CopyOp, memref::GlobalOp, memref::GetGlobalOp,
                 memref::ExpandShapeOp, memref::CollapseShapeOp, memref::CastOp,
-                memref::SubViewOp, memref::ExtractStridedMetadataOp,
-                memref::DimOp>([&](auto op) { return printOperation(op); })
+                memref::ReinterpretCastOp, memref::SubViewOp,
+                memref::ExtractStridedMetadataOp, memref::DimOp>(
+              [&](auto op) { return printOperation(op); })
 
           // Lattigo ops
           .Case<
@@ -1197,6 +1199,17 @@ LogicalResult LattigoEmitter::printOperation(memref::CastOp op) {
   return success();
 }
 
+LogicalResult LattigoEmitter::printOperation(memref::ReinterpretCastOp op) {
+  if (!op.getType().getLayout().isIdentity()) {
+    return op.emitOpError(
+        "requires result type to have identity layout (contiguous and 0 "
+        "offset)");
+  }
+  std::string name = getName(op.getResult());
+  emitAssignment(name, getName(op.getSource()));
+  return success();
+}
+
 LogicalResult LattigoEmitter::printOperation(memref::SubViewOp op) {
   MemRefType srcType = op.getSourceType();
   MemRefType resultType = op.getType();

lib/Target/Lattigo/LattigoEmitter.h

@@ -149,6 +149,7 @@ class LattigoEmitter {
   LogicalResult printOperation(::mlir::memref::ExpandShapeOp op);
   LogicalResult printOperation(::mlir::memref::CollapseShapeOp op);
   LogicalResult printOperation(::mlir::memref::CastOp op);
+  LogicalResult printOperation(::mlir::memref::ReinterpretCastOp op);
   LogicalResult printOperation(::mlir::memref::SubViewOp op);
   LogicalResult printOperation(::mlir::memref::ExtractStridedMetadataOp op);
   LogicalResult printOperation(::mlir::memref::DimOp op);

tests/Dialect/TensorExt/Conversions/tensor_ext_to_tensor/rotate_dynamic.mlir

@@ -13,18 +13,15 @@ func.func @test_rotate_dynamic(%0: tensor<16xi32>, %shift: index) -> tensor<16xi
   // CHECK: tensor.extract_slice
   // CHECK-SAME: [%[[v2]]] [%[[v3]]] [1]
 
-  // CHECK: tensor.empty
-
-  // CHECK: tensor.insert_slice
-  // CHECK-SAME: [%[[v3]]] [%[[v2]]] [1]
-  // CHECK: tensor.insert_slice
-  // CHECK-SAME: [0] [%[[v3]]] [1]
+  // CHECK: %[[INSERTED:.*]] = tensor.insert_slice %{{.*}} into %[[arg0]][%[[v3]]] [%[[v2]]] [1]
+  // CHECK: tensor.insert_slice %{{.*}} into %[[INSERTED]][0] [%[[v3]]] [1]
   %1 = tensor_ext.rotate %0, %shift : tensor<16xi32>, index
   return %1 : tensor<16xi32>
 }
 
 
 // CHECK: @test_rotate_dynamic_multidim
+// CHECK-SAME: (%[[arg0:.*]]: tensor<3x4x16xi32>, %{{.*}}: index) -> tensor<3x4x16xi32>
 func.func @test_rotate_dynamic_multidim(%0: tensor<3x4x16xi32>, %shift: index) -> tensor<3x4x16xi32> {
   // CHECK: %[[c16:.*]] = arith.constant 16
   // CHECK: %[[v0:.*]] = arith.remsi %[[shift]], %[[c16]]
@@ -37,12 +34,8 @@ func.func @test_rotate_dynamic_multidim(%0: tensor<3x4x16xi32>, %shift: index) -
   // CHECK: tensor.extract_slice
   // CHECK-SAME: [0, 0, %[[v2]]] [3, 4, %[[v3]]] [1, 1, 1]
 
-  // CHECK: tensor.empty
-
-  // CHECK: tensor.insert_slice
-  // CHECK-SAME: [0, 0, %[[v3]]] [3, 4, %[[v2]]] [1, 1, 1]
-  // CHECK: tensor.insert_slice
-  // CHECK-SAME: [0, 0, 0] [3, 4, %[[v3]]] [1, 1, 1]
+  // CHECK: %[[INSERTED:.*]] = tensor.insert_slice %{{.*}} into %[[arg0]][0, 0, %[[v3]]] [3, 4, %[[v2]]] [1, 1, 1]
+  // CHECK: tensor.insert_slice %{{.*}} into %[[INSERTED]][0, 0, 0] [3, 4, %[[v3]]] [1, 1, 1]
   %1 = tensor_ext.rotate %0, %shift : tensor<3x4x16xi32>, index
   return %1 : tensor<3x4x16xi32>
 }

tests/Dialect/TensorExt/Conversions/tensor_ext_to_tensor/rotate_static.mlir

@@ -1,37 +1,31 @@
 // RUN: heir-opt --tensor-ext-to-tensor %s | FileCheck %s
 
 // CHECK: @test_rotate
+// CHECK-SAME: (%[[arg0:.*]]: tensor<16xi32>) -> tensor<16xi32>
 func.func @test_rotate(%0: tensor<16xi32>) -> tensor<16xi32> {
   // CHECK: tensor.extract_slice
   // CHECK-SAME: [0] [1] [1]
   // CHECK: tensor.extract_slice
   // CHECK-SAME: [1] [15] [1]
 
-  // CHECK: tensor.empty
-
-  // CHECK: tensor.insert_slice
-  // CHECK-SAME: [15] [1] [1]
-  // CHECK: tensor.insert_slice
-  // CHECK-SAME: [0] [15] [1]
+  // CHECK: %[[INSERTED:.*]] = tensor.insert_slice %{{.*}} into %[[arg0]][15] [1] [1]
+  // CHECK: tensor.insert_slice %{{.*}} into %[[INSERTED]][0] [15] [1]
   %c1 = arith.constant 1 : i32
   %1 = tensor_ext.rotate %0, %c1 : tensor<16xi32>, i32
   return %1 : tensor<16xi32>
 }
 
 
 // CHECK: @test_rotate_multidim
+// CHECK-SAME: (%[[arg0:.*]]: tensor<3x4x16xi32>) -> tensor<3x4x16xi32>
 func.func @test_rotate_multidim(%0: tensor<3x4x16xi32>) -> tensor<3x4x16xi32> {
   // CHECK: tensor.extract_slice
   // CHECK-SAME: [0, 0, 0] [3, 4, 3] [1, 1, 1]
   // CHECK: tensor.extract_slice
   // CHECK-SAME: [0, 0, 3] [3, 4, 13] [1, 1, 1]
 
-  // CHECK: tensor.empty
-
-  // CHECK: tensor.insert_slice
-  // CHECK-SAME: [0, 0, 13] [3, 4, 3] [1, 1, 1]
-  // CHECK: tensor.insert_slice
-  // CHECK-SAME: [0, 0, 0] [3, 4, 13] [1, 1, 1]
+  // CHECK: %[[INSERTED:.*]] = tensor.insert_slice %{{.*}} into %[[arg0]][0, 0, 13] [3, 4, 3] [1, 1, 1]
+  // CHECK: tensor.insert_slice %{{.*}} into %[[INSERTED]][0, 0, 0] [3, 4, 13] [1, 1, 1]
   %c3 = arith.constant 3 : i32
   %1 = tensor_ext.rotate %0, %c3 : tensor<3x4x16xi32>, i32
   return %1 : tensor<3x4x16xi32>