[Mlir-commits] [mlir] [mlir][linalg] Implement bufferization for UnPackOp. (PR #182837)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Han-Chung Wang (hanhanW)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/182837.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp (+39) 
- (modified) mlir/test/Dialect/Linalg/bufferize.mlir (+16) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
index aed4fbf12bd43..ca5ee62e5ffeb 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -231,6 +231,44 @@ struct PackOpInterface
     return success();
   }
 };
+
+struct UnPackOpInterface
+    : public DstBufferizableOpInterfaceExternalModel<UnPackOpInterface,
+                                                     linalg::UnPackOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    auto unPackOp = cast<linalg::UnPackOp>(op);
+    return !unPackOp.isDpsInit(&opOperand);
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options,
+                          BufferizationState &state) const {
+    auto unPackOp = cast<linalg::UnPackOp>(op);
+    assert(!unPackOp.hasPureBufferSemantics() && "expected op with tensors");
+    if (!unPackOp.hasPureTensorSemantics())
+      return unPackOp.emitError()
+             << "mixed tensor/buffer semantic op not supported yet";
+    FailureOr<Value> sourceBuffer =
+        getBuffer(rewriter, unPackOp.getSource(), options, state);
+    if (failed(sourceBuffer))
+      return failure();
+    FailureOr<Value> destBuffer =
+        getBuffer(rewriter, unPackOp.getDest(), options, state);
+    if (failed(destBuffer))
+      return failure();
+
+    SmallVector<Value> operands;
+    operands.push_back(*sourceBuffer);
+    operands.push_back(*destBuffer);
+    llvm::append_range(operands, unPackOp.getInnerTiles());
+
+    linalg::UnPackOp::create(rewriter, unPackOp.getLoc(), TypeRange{}, operands,
+                             op->getAttrs());
+    replaceOpWithBufferizedValues(rewriter, op, *destBuffer);
+    return success();
+  }
+};
 } // namespace
 
 void mlir::linalg::registerBufferizableOpInterfaceExternalModels(
@@ -247,5 +285,6 @@ void mlir::linalg::registerBufferizableOpInterfaceExternalModels(
 
     SoftmaxOp::attachInterface<SoftmaxOpInterface>(*ctx);
     PackOp::attachInterface<PackOpInterface>(*ctx);
+    UnPackOp::attachInterface<UnPackOpInterface>(*ctx);
   });
 }
diff --git a/mlir/test/Dialect/Linalg/bufferize.mlir b/mlir/test/Dialect/Linalg/bufferize.mlir
index 6729cc4b76c4d..7b1af139f0f3e 100644
--- a/mlir/test/Dialect/Linalg/bufferize.mlir
+++ b/mlir/test/Dialect/Linalg/bufferize.mlir
@@ -226,3 +226,19 @@ func.func @bufferize_pack(%arg0: tensor<200x127x256xf32>, %arg1: tensor<256x64x2
   return %0 : tensor<256x64x200x2xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func @bufferize_unpack(
+// CHECK-SAME:   %[[SRC:.*]]: tensor<256x64x200x2xf32>, %[[DST:.*]]: tensor<200x128x256xf32>) -> tensor<200x128x256xf32> {
+// CHECK-DAG:     %[[SRC_BUF:.*]] = bufferization.to_buffer %[[SRC]] : tensor<256x64x200x2xf32> to memref<256x64x200x2xf32>
+// CHECK-DAG:     %[[DST_BUF:.*]] = memref.alloc() {{.*}} : memref<200x128x256xf32>
+// CHECK-NOT:     memref.copy
+// CHECK:         linalg.unpack %[[SRC_BUF]] outer_dims_perm = [2, 1, 0] inner_dims_pos = [1] inner_tiles = [2] into %[[DST_BUF]] : memref<256x64x200x2xf32> -> memref<200x128x256xf32>
+// CHECK:         %[[RESULT:.*]] = bufferization.to_tensor %[[DST_BUF]] : memref<200x128x256xf32> to tensor<200x128x256xf32>
+// CHECK:         return %[[RESULT]] : tensor<200x128x256xf32>
+func.func @bufferize_unpack(%arg0: tensor<256x64x200x2xf32>, %arg1: tensor<200x128x256xf32>) -> tensor<200x128x256xf32> {
+  %0 = linalg.unpack %arg0 outer_dims_perm = [2, 1, 0]
+      inner_dims_pos = [1] inner_tiles = [2] into %arg1
+      : tensor<256x64x200x2xf32> -> tensor<200x128x256xf32>
+  return %0 : tensor<200x128x256xf32>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/182837