From adcb98880972bec908249d091ed5634524250fa4 Mon Sep 17 00:00:00 2001
From: Hanhan Wang <hanchung@google.com>
Date: Tue, 17 Jan 2023 16:44:32 -0800
Subject: [PATCH] [mlir][tensor] Add a tiling interface registration for
 pack/unpack only.

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D141977
---
 mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h | 7 ++++++-
 mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp        | 8 ++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h b/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h
index 3649517..30a5026 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 0ab41b6..72ecf83 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);
+  });
+}
-- 
2.7.4