diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h b/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h --- a/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h @@ -47,7 +47,7 @@ /// 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 @@ /// implementation is moved to a separate library. void registerTilingInterfaceExternalModels(mlir::DialectRegistry ®istry); +/// Similar to the above registeration, but it is only for `tensor.pack` and +/// `tensor.unpack` ops. +void registerTilingInterfaceExternalModelsForPackUnPackOps( + DialectRegistry ®istry); + } // namespace tensor } // namespace mlir diff --git a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp --- a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp @@ -654,3 +654,11 @@ tensor::UnPackOp::attachInterface(*ctx); }); } + +void mlir::tensor::registerTilingInterfaceExternalModelsForPackUnPackOps( + DialectRegistry ®istry) { + registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) { + tensor::PackOp::attachInterface(*ctx); + tensor::UnPackOp::attachInterface(*ctx); + }); +}