dialects: (bufferization) add materialize_in_destination (#3301)

Currently we need only the command with tensor arguments but we should extend it to memrefs as well at some point.
xdslproject · Oct 15, 2024 · 810d653 · 810d653
1 parent 877b70c
commit 810d653
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/tests/filecheck/mlir-conversion/with-mlir/dialects/bufferization/ops.mlir b/tests/filecheck/mlir-conversion/with-mlir/dialects/bufferization/ops.mlir
@@ -5,11 +5,18 @@ module{
     %m = "test.op"() : () -> memref<30x20x10xf32>
     %m_t = bufferization.to_tensor %m restrict writable : memref<30x20x10xf32>
     %t_m = bufferization.to_memref %m_t read_only : memref<30x20x10xf32>
+
+    %tensor1 = "test.op"() : () -> tensor<2x2xf64>
+    %tensor2 = "test.op"() : () -> tensor<2x2xf64>
+    %b = bufferization.materialize_in_destination %tensor1 in %tensor2 : (tensor<2x2xf64>, tensor<2x2xf64>) -> tensor<2x2xf64>
 }
 
 // CHECK:       builtin.module {
 // CHECK-NEXT:    %0 = "bufferization.alloc_tensor"() <{"operandSegmentSizes" = array<i32: 0, 0, 0>}> : () -> tensor<10x20x30xf64>
 // CHECK-NEXT:    %1 = "test.op"() : () -> memref<30x20x10xf32>
 // CHECK-NEXT:    %2 = bufferization.to_tensor %1 restrict writable : memref<30x20x10xf32>
 // CHECK-NEXT:    %3 = bufferization.to_memref %2 read_only : memref<30x20x10xf32>
+// CHECK-NEXT:    %4 = "test.op"() : () -> tensor<2x2xf64>
+// CHECK-NEXT:    %5 = "test.op"() : () -> tensor<2x2xf64>
+// CHECK-NEXT:    %6 = bufferization.materialize_in_destination %4 in %5 : (tensor<2x2xf64>, tensor<2x2xf64>) -> tensor<2x2xf64>
 // CHECK-NEXT:  }
diff --git a/xdsl/dialects/bufferization.py b/xdsl/dialects/bufferization.py
@@ -168,12 +168,38 @@ class ToMemrefOp(IRDLOperation):
     assembly_format = "$tensor (`read_only` $read_only^)?  `:` attr-dict type($memref)"
 
 
+@irdl_op_definition
+class MaterializeInDestination(IRDLOperation):
+    name = "bufferization.materialize_in_destination"
+
+    source = operand_def(
+        TensorMemrefInferenceConstraint(
+            "T", AnyTensorTypeConstr | AnyUnrankedTensorTypeConstr
+        )
+    )
+    dest = operand_def(
+        TensorMemrefInferenceConstraint(
+            "T", AnyTensorTypeConstr | AnyUnrankedTensorTypeConstr
+        )
+    )
+    result = result_def(
+        TensorMemrefInferenceConstraint(
+            "T", AnyTensorTypeConstr | AnyUnrankedTensorTypeConstr
+        )
+    )
+    restrict = opt_prop_def(UnitAttr)
+    writable = opt_prop_def(UnitAttr)
+
+    assembly_format = "$source `in` (`restrict` $restrict^)? (`writable` $writable^)? $dest attr-dict `:` `(` type($source) `,` type($dest) `)` `->` type($result)"
+
+
 Bufferization = Dialect(
     "bufferization",
     [
         AllocTensorOp,
         ToTensorOp,
         ToMemrefOp,
+        MaterializeInDestination,
     ],
     [],
 )