[Mlir-commits] [mlir] adcb988 - [mlir][tensor] Add a tiling interface registration for pack/unpack only.

[Hanhan Wang]

Author: Hanhan Wang
Date: 2023-01-26T17:53:02-08:00
New Revision: adcb98880972bec908249d091ed5634524250fa4

URL: https://github.com/llvm/llvm-project/commit/adcb98880972bec908249d091ed5634524250fa4
DIFF: https://github.com/llvm/llvm-project/commit/adcb98880972bec908249d091ed5634524250fa4.diff

LOG: [mlir][tensor] Add a tiling interface registration for pack/unpack only.

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D141977

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h
    mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h b/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h
index 364951756492d..30a5026cd68b3 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h
@@ -47,7 +47,7 @@ Operation *bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
 /// Registers external models for Tiling interface for tensor ops.
 /// Currently, it registers:
 ///
-/// * TilingInterface for `tensor.pad`.
+/// * TilingInterface for `tensor.pad`, `tensor.pack`, and `tensor.unpack`.
 ///
 /// Unfortunately, a "normal" internal registration is not possible at the
 /// moment, because of the dependency of the interface implementation for these
@@ -56,6 +56,11 @@ Operation *bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
 /// implementation is moved to a separate library.
 void registerTilingInterfaceExternalModels(mlir::DialectRegistry &registry);
 
+/// Similar to the above registeration, but it is only for `tensor.pack` and
+/// `tensor.unpack` ops.
+void registerTilingInterfaceExternalModelsForPackUnPackOps(
+    DialectRegistry &registry);
+
 } // namespace tensor
 } // namespace mlir
 

diff  --git a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
index 0ab41b6343f67..72ecf830ecfc5 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
@@ -654,3 +654,11 @@ void mlir::tensor::registerTilingInterfaceExternalModels(
     tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
   });
 }
+
+void mlir::tensor::registerTilingInterfaceExternalModelsForPackUnPackOps(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
+    tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
+    tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
+  });
+}