diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td --- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td +++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td @@ -1884,8 +1884,9 @@ def HoistRedundantTensorSubsetsOp : Op { + [DeclareOpInterfaceMethods, + TransformEachOpTrait, + TransformOpInterface]> { let description = [{ Hoists supported tensor subset extract/insert operation pairs out of immediately enclosing loop iteratively, if the following conditions diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp --- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp +++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp @@ -3208,6 +3208,13 @@ return DiagnosedSilenceableFailure::success(); } +void transform::HoistRedundantTensorSubsetsOp::getEffects( + SmallVectorImpl &effects) { + transform::onlyReadsHandle(getTarget(), effects); + transform::producesHandle(getTransformed(), effects); + transform::modifiesPayload(effects); +} + //===----------------------------------------------------------------------===// // Transform op registration //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/Linalg/hoisting.mlir b/mlir/test/Dialect/Linalg/hoisting.mlir --- a/mlir/test/Dialect/Linalg/hoisting.mlir +++ b/mlir/test/Dialect/Linalg/hoisting.mlir @@ -303,6 +303,9 @@ : (!pdl.operation) -> !pdl.operation transform.structured.hoist_redundant_tensor_subsets %0 : (!pdl.operation) -> !pdl.operation + // Test we can call the op twice without consuming the handle. + transform.structured.hoist_redundant_tensor_subsets %0 + : (!pdl.operation) -> !pdl.operation } // -----