[mlir][tensor] Add transform to make tensor.pad/empty loop-independent

[springerm]

Summary

Add a transform to make tensor.pad and tensor.empty ops independent of SCF loop IVs. Such ops can then be hoisted.

E.g.:

scf.for %iv = %lb to %ub step %step {
  %high = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (%i)[%ub]
  %p = tensor.pad %t low[5] high[%high] ...
  ...
}

Is transformed to:

%high_new = affine.apply affine_map<()[s0, s1] -> (-s0 + s1)> ()[%lb, %ub]
%p_hoistable = tensor.pad %t low[5] high[%high_new]
%dim = tensor.dim %t, %c0
%size = affine.apply affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 + 5)>(%iv)[%ub, %dim]
%slice = tensor.extract_slice %p_hoistable [0] [%size] [1]

Depends On: D146524

[springerm]

@dcaballe The same infrastructure can be used for computing upper bounds for memref allocations.

[springerm]

add op documentation

[nicolasvasilache]

mlir/include/mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h

17 ↗

(On Diff #497222)

Isn't the best practice here to just forward-declare ?
I am unclear why you reverted to a mix of forward-declaration for OpBuilder and include for AffineMap/Value/ValueRange ?

mlir/include/mlir/Dialect/Tensor/TransformOps/CMakeLists.txt

6 ↗

(On Diff #497222)

This fails to build for me with:

CMake Error at /usr/local/google/home/ntv/github/llvm-project/mlir/include/mlir/Dialect/Tensor/TransformOps/CMakeLists.txt:6 (add_mlir_doc):
  add_mlir_doc Function invoked with incorrect arguments for function named:
  add_mlir_doc

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td

27

Please explain a little more what this entails because it involves increasing the tensor size/dimensionality.

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

39

The footprint in this example can be too high: you'd want to take the ceildiv by %step.

mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp

87 ↗

(On Diff #497222)

Can we use step 1 instead here and mention that this is conservative but less precise in the context of this method?

[springerm]

mlir/include/mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h

17 ↗

(On Diff #497222)

I thought I needed the definition so that it can be used in FailureOr<...>. But it actually works without.

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

39

I tried to do that but FlatAffineConstraints was unable to compute an upper bound in that case. It's unclear why this is happening, I'm looking into this. Maybe the system of inequalities is getting too complex (with various "semi-affine exprs"...). It looks like FlatAffineConstraints must be extended.

[springerm]

address comments

[springerm]

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

39

This is indeed due to a shortcoming in FlatAffineValueConstraints:

// TODO: Whenever there are local variables in the dependence
// constraints, we'll conservatively over-approximate, since we don't
// always explicitly compute them above (in the while loop).

It is worth fixing this? Will probably take a while to understand and rewrite a 200 LOC function.

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

39

Yes, this is the same rationale as using step 1 below.

Can we update the doc to make this explicit?
I.e. divide by %step and mention that in the case of a symbol, we over-approximate with setting step to 1 to circumvent the limitation of the analysis.

mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp

86 ↗

(On Diff #500698)

if we use step 1, this is not an optional anymore and there is a bit of simplification in your code down the line

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td

23

Can we rename this transform to transform.loop_privatize or something similar and have it take the op to privatize and the number of enclosing loops above which we want to hoist ?

This should be kept in sync and compose with the recently landed hoist_pad.
We can later evolve the syntax from num_loops to something better, informed by how we want these things to compose.

mlir/lib/Dialect/Tensor/Transforms/LoopHoisting.cpp

21 ↗

(On Diff #500698)

The filename is misleading here, this is not performing hoisting but privatization that will later enable hoisting.
Can we move this functionality to a LoopPrivatization.cpp file ?

As a followup, we should integrate the usage of privatization its usage into HoistPadding.cpp.

We also now have SubsetHoisting.cpp for mechanical parts related to actual hoisting of loop-independent quantities that will also come in handy..

[springerm]

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td

23

This does not privatize the tensor/op though. It simply changes the size of the tensor.

Added the number of loops to the op.

mlir/lib/Dialect/Tensor/Transforms/LoopHoisting.cpp

21 ↗

(On Diff #500698)

How about LoopTransforms.cpp? This transformation does not do any privatization. (Assuming "privatization" = "making a private copy of a tensor for each loop iteration".)

[springerm]

address comments

[springerm]

update

[springerm]

reimplement with ValueBoundsOpInterface

[nicolasvasilache]

Does this update change the way this is supposed to interact with HoistPadding ?
Or in other words, how do you see this interacting with HoistPadding ?

[springerm]

In D143910#4212287, @nicolasvasilache wrote:

Does this update change the way this is supposed to interact with HoistPadding ?
Or in other words, how do you see this interacting with HoistPadding ?

HoistPadding does not require this functionality, it just clones the entire loop nest. So no interaction with HoistPadding.

This revision is already a month old and was never landed, but is useful for Diego as an example how to make ops (memref.alloca in his example) hoistable. So I reimplemented it so that it takes advantage of ValueBoundsOpInterface.

[springerm]

rebase

[dcaballe]

Looking at the ValueBounds part only for now.

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

77

I can't follow from the description... This is converting an Affine-based value bound into an Arith-based value bound?

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

76

typos

Could you please elaborate a bit more on what "independent of the values in independencies" mean?

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

394

Something like this in the header doc is what I was hoping for :)

[springerm]

address comments

[springerm]

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

77

The name was confusing. I renamed the function and added some more documentation.

[springerm]

rebase

[dcaballe]

Thanks!

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

39

What happened with this? Was it addressed?

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

67

nit: ub to var

mlir/lib/Dialect/Tensor/Transforms/IndependenceTransforms.cpp

80

nit: ub to var

[springerm]

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

39

There is a TODO in SCF/IR/ValueBoundsOpInterfaceImpl.cpp. We can't do any better at the moment.

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

 //===- Transforms.h - Transforms Entrypoints --------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines a set of transforms specific for the AffineOps
 // dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H
 #define MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H
 
+#include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 class AffineApplyOp;
+class AffineMap;
 class Location;
 class OpBuilder;
 class OpFoldResult;
 class RewritePatternSet;
 class RewriterBase;
 class Value;
 
 namespace presburger {

 /// By default, lower/equal bounds are closed and upper bounds are open. If
 /// `closedUB` is set to "true", upper bounds are also closed.
 FailureOr<OpFoldResult>
 reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
                 Value value, std::optional<int64_t> dim,
                 function_ref<bool(Value, std::optional<int64_t>)> stopCondition,
                 bool closedUB = false);
 
+/// Reify an already computed bound with Affine dialect ops.
+OpFoldResult
+reifyValueBound(OpBuilder &b, Location loc, AffineMap boundMap,
+                ArrayRef<std::pair<Value, std::optional<int64_t>>> mapOperands);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Dialect/Tensor/CMakeLists.txt

 add_subdirectory(IR)
 add_subdirectory(Transforms)
+add_subdirectory(TransformOps)

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h

+//===- TensorTransformOps.h - Tensor transformation ops ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TENSOR_TRANSFORMOPS_TENSORTRANSFORMOPS_H
+#define MLIR_DIALECT_TENSOR_TRANSFORMOPS_TENSORTRANSFORMOPS_H
+
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
+#include "mlir/IR/OpImplementation.h"
+
+namespace mlir {
+class DialectRegistry;
+
+namespace tensor {
+class PadOp;
+
+void registerTransformDialectExtension(DialectRegistry &registry);
+} // namespace tensor
+} // namespace mlir
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h.inc"
+
+#endif // MLIR_DIALECT_TENSOR_TRANSFORMOPS_TENSORTRANSFORMOPS_H

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td

+//===- TensorTransformOps.td - Tensor transformation ops ---*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TENSOR_TRANSFORM_OPS
+#define TENSOR_TRANSFORM_OPS
+
+include "mlir/Dialect/PDL/IR/PDLTypes.td"
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+include "mlir/Dialect/Transform/IR/TransformTypes.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpBase.td"
+
+def Transform_TensorPadOp : Transform_ConcreteOpType<"tensor.pad">;
+
+def MakeLoopIndependentOp
+    : Op<Transform_Dialect, "tensor.make_loop_independent",
+         [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+          TransformOpInterface, TransformEachOpTrait]> {
+  let description = [{
+    Rewrite the targeted ops such that their index-typed operands no longer
+    depend on any loop induction variable of the `num_loop` enclosing `scf.for`
+    loops. I.e., compute an upper bound that is independent of any such loop IV
+    for every tensor dimension. The transformed op could then be hoisted from
+    the `num_loop` enclosing loops. To preserve the original semantics, place a
+    `tensor.extract_slice` inside the loop.
+
+    Currently supported operations are:
+    - tensor.empty: Replaced with a new tensor.empty with upper bound sizes,
+      followed by a tensor.extract_slice.
+    - tensor.pad: Replaced by an upper bound padding, followed by a
+      tensor.extract_slice.
+
+    #### Return modes
+
+    This operation fails if at least one induction variable could not be
+    eliminated. In case the targeted op is already independent of induction
+    variables, this transform succeeds and returns the unmodified target op.
+
+    Otherwise, the returned handle points to a subset of the produced ops:
+    - tensor.empty: The returned handle points to the tensor.extract_slice op.
+    - tensor.pad: The returned handle points to the tensor.extract_slice op.
+
+    This transform op consumes the target handle and produces a result handle.
+  }];
+
+  let arguments = (ins PDL_Operation:$target, I64Attr:$num_loops);
+  let results = (outs PDL_Operation:$transformed);
+  let assemblyFormat = "$target attr-dict";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
+#endif // TENSOR_TRANSFORM_OPS

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 
 struct TilingResult;
 
 namespace tensor {
 
+/// Build a new tensor::PadOp with low/high padding that is independent of all
+/// given independencies. If the op is already independent of all
+/// independencies, the same PadOp result is returned.
+///
+/// Failure indicates the no suitable upper bound for low/high padding could be
+/// found.
+///
+/// Example:
+/// scf.for %iv = %lb to %ub step %step {
+///   %high = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (%i)[%ub]
+///   %p = tensor.pad %t low[5] high[%high] ...
+///   ...
+/// }
+///
+/// The function builds IR such as:
+/// %high_new = affine.apply affine_map<()[s0, s1] -> (-s0 + s1)> ()[%lb, %ub]
+/// %p_hoistable = tensor.pad %t low[5] high[%high_new]
+/// %dim = tensor.dim %t, %c0
+/// %size = affine.apply affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 + 5)>
+///     (%iv)[%ub, %dim]
+/// %slice = tensor.extract_slice %p_hoistable [0] [%size] [1]
+///
+/// The slice is returned.
+FailureOr<Value> buildIndependentOp(OpBuilder &b, tensor::PadOp padOp,
+                                    ValueRange independencies);
+
+/// Build a new tensor::EmptyOp who's dynamic sizes are independent of all
+/// given independencies. If the op is already independent of all
+/// independencies, the same EmptyOp result is returned.
+///
+/// Failure indicates the no suitable upper bound for the dynamic sizes could be
+/// found.
+FailureOr<Value> buildIndependentOp(OpBuilder &b, tensor::EmptyOp emptyOp,
+                                    ValueRange independencies);
+
 /// Populates `patterns` with patterns to wrap a tensor.pad op with an scf.if op
 /// to separate the cases where we don't need padding (all pad sizes are
 /// actually zeros) and where we indeed need padding.
 void populateSplitPaddingPatterns(RewritePatternSet &patterns,
                                   PatternBenefit baseBenefit = 1);
 
 /// Pattern to swap an `tensor.extract_slice` with its producer when the
 /// producer implements the `TilingInterface`. The pattern itself does not

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
 #include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Dialect.h"
inline void registerAllDialects(DialectRegistry &registry) {
 
   // Register all dialect extensions.
   affine::registerTransformDialectExtension(registry);
   bufferization::registerTransformDialectExtension(registry);
   gpu::registerTransformDialectExtension(registry);
   linalg::registerTransformDialectExtension(registry);
   memref::registerTransformDialectExtension(registry);
   scf::registerTransformDialectExtension(registry);
+  tensor::registerTransformDialectExtension(registry);
   vector::registerTransformDialectExtension(registry);
 
   // Register all external models.
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   arith::registerValueBoundsOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

public:
   /// Compute a bound in terms of the values/dimensions in `dependencies`.
   static LogicalResult computeBound(AffineMap &resultMap,
                                     ValueDimList &mapOperands,
                                     presburger::BoundType type, Value value,
                                     std::optional<int64_t> dim,
                                     ValueDimList dependencies,
                                     bool closedUB = false);
 
+  /// Compute a bound in that is independent of the of the values in
+  /// `independencies`.
+  static LogicalResult
+  computeIndependentBound(AffineMap &resultMap, ValueDimList &mapOperands,
+                          presburger::BoundType type, Value value,
+                          std::optional<int64_t> dim, ValueRange independencies,
+                          bool closedUB = false);
+
   /// Compute a constant bound for the given index-typed value or shape
   /// dimension size.
   ///
   /// `dim` must be `nullopt` if and only if `value` is index-typed. This
   /// function traverses the backward slice of the given value in a
   /// worklist-driven manner until `stopCondition` evaluates to "true". The
   /// constraint set is populated according to `ValueBoundsOpInterface` for each
   /// visited value. (No constraints are added for values for which the stop

mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp

FailureOr<OpFoldResult> mlir::reifyValueBound(
     bool closedUB) {
   // Compute bound.
   AffineMap boundMap;
   ValueDimList mapOperands;
   if (failed(ValueBoundsConstraintSet::computeBound(
           boundMap, mapOperands, type, value, dim, stopCondition, closedUB)))
     return failure();
 
+  // Reify bound.
+  return reifyValueBound(b, loc, boundMap, mapOperands);
+}
+
+OpFoldResult mlir::reifyValueBound(
+    OpBuilder &b, Location loc, AffineMap boundMap,
+    ArrayRef<std::pair<Value, std::optional<int64_t>>> mapOperands) {
   // Materialize tensor.dim/memref.dim ops.
   SmallVector<Value> operands;
   for (auto valueDim : mapOperands) {
     Value value = valueDim.first;
     std::optional<int64_t> dim = valueDim.second;
 
     if (!dim.has_value()) {
       // This is an index-typed value.

mlir/lib/Dialect/Tensor/CMakeLists.txt

 add_subdirectory(IR)
 add_subdirectory(Transforms)
+add_subdirectory(TransformOps)
 add_subdirectory(Utils)

mlir/lib/Dialect/Tensor/TransformOps/CMakeLists.txt

+add_mlir_dialect_library(MLIRTensorTransformOps
+  TensorTransformOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/TransformOps
+
+  DEPENDS
+  MLIRTensorTransformOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRAffineDialect
+  MLIRIR
+  MLIRPDLDialect
+  MLIRSCFDialect
+  MLIRTensorTransforms
+  MLIRTransformDialect
+)

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

+//===- TensorTransformOps.cpp - Implementation of tensor transform ops ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// MakeLoopIndependentOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::MakeLoopIndependentOp::applyToOne(
+    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  // Gather IVs.
+  SmallVector<Value> ivs;
+  Operation *nextOp = target;
+  for (uint64_t i = 0; i < getNumLoops(); ++i) {
+    nextOp = nextOp->getParentOfType<scf::ForOp>();
+    if (!nextOp) {
+      DiagnosedSilenceableFailure diag = emitSilenceableError()
+                                         << "could not find " << i
+                                         << "-th enclosing loop";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
+    ivs.push_back(cast<scf::ForOp>(nextOp).getInductionVar());
+  }
+
+  // Rewrite IR.
+  IRRewriter rewriter(target->getContext());
+  FailureOr<Value> replacement = failure();
+  if (auto padOp = dyn_cast<tensor::PadOp>(target)) {
+    replacement = tensor::buildIndependentOp(rewriter, padOp, ivs);
+  } else if (auto emptyOp = dyn_cast<tensor::EmptyOp>(target)) {
+    replacement = tensor::buildIndependentOp(rewriter, emptyOp, ivs);
+  } else {
+    DiagnosedSilenceableFailure diag = emitSilenceableError()
+                                       << "unsupported target p[]";
+    diag.attachNote(target->getLoc()) << "target op";
+    return diag;
+  }
+  if (failed(replacement)) {
+    DiagnosedSilenceableFailure diag =
+        emitSilenceableError() << "could not make target op loop-independent";
+    diag.attachNote(target->getLoc()) << "target op";
+    return diag;
+  }
+  rewriter.replaceOp(target, *replacement);
+  results.push_back(replacement->getDefiningOp());
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
+// Transform op registration
+//===----------------------------------------------------------------------===//
+
+namespace {
+class TensorTransformDialectExtension
+    : public transform::TransformDialectExtension<
+          TensorTransformDialectExtension> {
+public:
+  using Base::Base;
+
+  void init() {
+    declareGeneratedDialect<AffineDialect>();
+    declareGeneratedDialect<tensor::TensorDialect>();
+
+    registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.cpp.inc"
+        >();
+  }
+};
+} // namespace
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.cpp.inc"
+
+void mlir::tensor::registerTransformDialectExtension(
+    DialectRegistry &registry) {
+  registry.addExtensions<TensorTransformDialectExtension>();
+}

mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRTensorTransforms
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 
   DEPENDS
   MLIRTensorTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
+  MLIRAffineTransforms
   MLIRAffineUtils
   MLIRArithDialect
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRIR
   MLIRLinalgDialect
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFDialect
   MLIRTensorDialect
   MLIRTilingInterface
   MLIRTransforms
+  MLIRValueBoundsOpInterface
 )

mlir/lib/Dialect/Tensor/Transforms/IndependenceTransforms.cpp

+//===- IndependenceTransforms.cpp - Make ops independent of values --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+/// Make the given OpFoldResult independent of all independencies.
+static FailureOr<OpFoldResult> makeIndependent(OpBuilder &b, Location loc,
+                                               OpFoldResult ofr,
+                                               ValueRange independencies) {
+  if (ofr.is<Attribute>())
+    return ofr;
+  Value value = ofr.get<Value>();
+  AffineMap boundMap;
+  ValueDimList mapOperands;
+  if (failed(ValueBoundsConstraintSet::computeIndependentBound(
+          boundMap, mapOperands, presburger::BoundType::UB, value,
+          /*dim=*/std::nullopt, independencies, /*closedUB=*/true)))
+    return failure();
+  return reifyValueBound(b, loc, boundMap, mapOperands);
+};
+
+FailureOr<Value> tensor::buildIndependentOp(OpBuilder &b, tensor::PadOp padOp,
+                                            ValueRange independencies) {
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(padOp);
+  Location loc = padOp.getLoc();
+
+  // Non-constant padding not supported.
+  Value constantPadding = padOp.getConstantPaddingValue();
+  if (!constantPadding)
+    return failure();
+
+  SmallVector<OpFoldResult> newMixedLow, newMixedHigh;
+  for (OpFoldResult ofr : padOp.getMixedLowPad()) {
+    auto ub = makeIndependent(b, loc, ofr, independencies);
+    if (failed(ub))
+      return failure();
+    newMixedLow.push_back(*ub);
+  }
+  for (OpFoldResult ofr : padOp.getMixedHighPad()) {
+    auto ub = makeIndependent(b, loc, ofr, independencies);
+    if (failed(ub))
+      return failure();
+    newMixedHigh.push_back(*ub);
+  }
+
+  // Return existing tensor::PadOp if nothing has changed.
+  if (llvm::equal(padOp.getMixedLowPad(), newMixedLow) &&
+      llvm::equal(padOp.getMixedHighPad(), newMixedHigh))
+    return padOp.getResult();
+
+  // Create a new tensor::PadOp.
+  auto newPadOp = b.create<PadOp>(
+      loc, padOp.getResultType(), padOp.getSource(), newMixedLow, newMixedHigh,
+      constantPadding, padOp.getNofold(), /*attrs=*/ArrayRef<NamedAttribute>{});
+
+  // Create a tensor::ExtractSliceOp.
+  // Reify the result sizes of the old tensor::PadOp.
+  ReifiedRankedShapedTypeDims reifiedSizes;
+  ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
+      dyn_cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation());
+  if (failed(reifyShapedTypeInterface.reifyResultShapes(b, reifiedSizes)))
+    return failure();
+  SmallVector<OpFoldResult> offsets, sizes, strides;
+  for (int64_t i = 0; i < padOp.getResultType().getRank(); ++i) {
+    // offset = ub(low_padding) - low_padding
+    OpFoldResult prevLow = padOp.getMixedLowPad()[i];
+    if (prevLow.is<Attribute>()) {
+      offsets.push_back(b.getIndexAttr(0));
+    } else {
+      offsets.push_back(
+          b.create<AffineApplyOp>(
+               loc, b.getAffineDimExpr(0) - b.getAffineDimExpr(1),
+               std::initializer_list<Value>{newMixedLow[i].get<Value>(),
+                                            prevLow.get<Value>()})
+              .getResult());
+    }
+    // size = reified result size
+    if (!padOp.getResultType().isDynamicDim(i)) {
+      sizes.push_back(b.getIndexAttr(padOp.getResultType().getDimSize(i)));
+    } else {
+      sizes.push_back(reifiedSizes[0][i]);
+    }
+    // stride = 1
+    strides.push_back(b.getIndexAttr(1));
+  }
+
+  return b.create<ExtractSliceOp>(loc, newPadOp, offsets, sizes, strides)
+      .getResult();
+}
+
+FailureOr<Value> tensor::buildIndependentOp(OpBuilder &b,
+                                            tensor::EmptyOp emptyOp,
+                                            ValueRange independencies) {
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(emptyOp);
+  Location loc = emptyOp.getLoc();
+
+  SmallVector<OpFoldResult> newSizes;
+  for (OpFoldResult ofr : emptyOp.getMixedSizes()) {
+    auto ub = makeIndependent(b, loc, ofr, independencies);
+    if (failed(ub))
+      return failure();
+    newSizes.push_back(*ub);
+  }
+
+  // Return existing tensor::EmptyOp if nothing has changed.
+  if (llvm::equal(emptyOp.getMixedSizes(), newSizes))
+    return emptyOp.getResult();
+
+  // Create a new tensor::EmptyOp.
+  Value newEmptyOp =
+      b.create<EmptyOp>(loc, newSizes, emptyOp.getType().getElementType());
+
+  // Create a tensor::ExtractSliceOp.
+  SmallVector<OpFoldResult> offsets(newSizes.size(), b.getIndexAttr(0));
+  SmallVector<OpFoldResult> strides(newSizes.size(), b.getIndexAttr(1));
+  return b
+      .create<ExtractSliceOp>(loc, newEmptyOp, offsets, emptyOp.getMixedSizes(),
+                              strides)
+      .getResult();
+}

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

LogicalResult ValueBoundsConstraintSet::computeBound(
   return computeBound(
       resultMap, mapOperands, type, value, dim,
       [&](Value v, std::optional<int64_t> d) {
         return llvm::is_contained(dependencies, std::make_pair(v, d));
       },
       closedUB);
 }
 
+LogicalResult ValueBoundsConstraintSet::computeIndependentBound(
+    AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
+    Value value, std::optional<int64_t> dim, ValueRange independencies,
+    bool closedUB) {
+  // Return "true" if the given value is independent of all values in
+  // `independencies`. I.e., neither the value itself nor any value in the
+  // backward slice (reverse use-def chain) is contained in `independencies`.
+  auto isIndependent = [&](Value v) {
+    SmallVector<Value> worklist;
+    DenseSet<Value> visited;
+    worklist.push_back(v);
+    while (!worklist.empty()) {
+      Value next = worklist.pop_back_val();
+      if (visited.contains(next))
+        continue;
+      visited.insert(next);
+      if (llvm::is_contained(independencies, next))
+        return false;
+      // TODO: DominanceInfo could be used to stop the traversal early.
+      Operation *op = next.getDefiningOp();
+      if (!op)
+        continue;
+      worklist.append(op->getOperands().begin(), op->getOperands().end());
+    }
+    return true;
+  };
+
+  // Reify bounds in terms of any independent values.
+  return computeBound(
+      resultMap, mapOperands, type, value, dim,
+      [&](Value v, std::optional<int64_t> d) { return isIndependent(v); },
+      closedUB);
+}
+
 FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
     presburger::BoundType type, Value value, std::optional<int64_t> dim,
     StopConditionFn stopCondition, bool closedUB) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
 #endif // NDEBUG
 
   // Process the backward slice of `value` (i.e., reverse use-def chain) until

mlir/test/Dialect/Tensor/transform-op-make-loop-independent.mlir

+// RUN: mlir-opt %s -allow-unregistered-dialect \
+// RUN:     -test-transform-dialect-interpreter -canonicalize \
+// RUN:     -split-input-file -verify-diagnostics | FileCheck %s
+
+// This is a test case where "high" padding depends on the IV.
+
+//       CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 - s1)>
+//       CHECK: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 + 5)>
+// CHECK-LABEL: func @make_pad_loop_independent_1(
+//  CHECK-SAME:     %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index,
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
+func.func @make_pad_loop_independent_1(%lb: index, %ub: index, %step: index,
+                                       %t: tensor<?xf32>, %f: f32) {
+  // CHECK: scf.for %[[iv:.*]] = %[[lb]] to %[[ub]]
+  scf.for %i = %lb to %ub step %step {
+    // CHECK: %[[high:.*]] = affine.apply #[[$map]]()[%[[ub]], %[[lb]]]
+    // CHECK: %[[padded:.*]] = tensor.pad %[[t]] low[5] high[%[[high]]]
+    // CHECK: %[[dim:.*]] = tensor.dim %[[t]]
+    // CHECK: %[[size:.*]] = affine.apply #[[$map1]](%[[iv]])[%[[ub]], %[[dim]]]
+    // CHECK: %[[replacement:.*]] = tensor.extract_slice %[[padded]][0] [%[[size]]] [1]
+    %high = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (%i)[%ub]
+    %p = tensor.pad %t low[5] high[%high] {
+    ^bb0(%arg1: index):
+      tensor.yield %f : f32
+    } : tensor<?xf32> to tensor<?xf32>
+    // CHECK: "dummy.some_use"(%[[replacement]])
+    "dummy.some_use"(%p) : (tensor<?xf32>) -> ()
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.tensor.make_loop_independent %0 {num_loops = 1}
+}
+
+// -----
+
+// This is a test case where "low" padding depends on the IV.
+
+//       CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 - s1)>
+//       CHECK: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 + 5)>
+//       CHECK: #[[$map2:.*]] = affine_map<(d0)[s0] -> (d0 - s0)>
+// CHECK-LABEL: func @make_pad_loop_independent_1(
+//  CHECK-SAME:     %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index,
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
+func.func @make_pad_loop_independent_1(%lb: index, %ub: index, %step: index,
+                                       %t: tensor<?xf32>, %f: f32) {
+  // CHECK: scf.for %[[iv:.*]] = %[[lb]] to %[[ub]]
+  scf.for %i = %lb to %ub step %step {
+    // CHECK: %[[low:.*]] = affine.apply #[[$map]]()[%[[ub]], %[[lb]]]
+    // CHECK: %[[padded:.*]] = tensor.pad %[[t]] low[%[[low]]] high[5]
+    // CHECK: %[[dim:.*]] = tensor.dim %[[t]]
+    // CHECK: %[[size:.*]] = affine.apply #[[$map1]](%[[iv]])[%[[ub]], %[[dim]]]
+    // CHECK: %[[offset:.*]] = affine.apply #[[$map2]](%[[iv]])[%[[lb]]]
+    // CHECK: %[[replacement:.*]] = tensor.extract_slice %[[padded]][%[[offset]]] [%[[size]]] [1]
+    %low = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (%i)[%ub]
+    %p = tensor.pad %t low[%low] high[5] {
+    ^bb0(%arg1: index):
+      tensor.yield %f : f32
+    } : tensor<?xf32> to tensor<?xf32>
+    // CHECK: "dummy.some_use"(%[[replacement]])
+    "dummy.some_use"(%p) : (tensor<?xf32>) -> ()
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.tensor.make_loop_independent %0 {num_loops = 1}
+}
+
+// -----
+
+//       CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 * 2 - 2)>
+// CHECK-LABEL: func @two_loops(
+func.func @two_loops(%lb: index, %ub: index, %step: index,
+                     %t: tensor<?xf32>, %f: f32) {
+  scf.for %i = %lb to %ub step %step {
+    scf.for %j = %lb to %ub step %step {
+      // CHECK: affine.apply #map()[%{{.*}}]
+      %low = affine.apply affine_map<(d0, d1)[] -> (d0 + d1)> (%i, %j)[]
+      %p = tensor.pad %t low[%low] high[5] {
+      ^bb0(%arg1: index):
+        tensor.yield %f : f32
+      } : tensor<?xf32> to tensor<?xf32>
+      "dummy.some_use"(%p) : (tensor<?xf32>) -> ()
+    }
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.tensor.make_loop_independent %0 {num_loops = 2}
+}
+
+// -----
+
+func.func @not_enough_loops(%lb: index, %ub: index, %step: index,
+                            %t: tensor<?xf32>, %f: f32) {
+  scf.for %i = %lb to %ub step %step {
+    scf.for %j = %lb to %ub step %step {
+      %low = affine.apply affine_map<(d0, d1)[] -> (d0 + d1)> (%i, %j)[]
+      // expected-note@below {{target op}}
+      %p = tensor.pad %t low[%low] high[5] {
+      ^bb0(%arg1: index):
+        tensor.yield %f : f32
+      } : tensor<?xf32> to tensor<?xf32>
+      "dummy.some_use"(%p) : (tensor<?xf32>) -> ()
+    }
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  // expected-error@below {{could not find 2-th enclosing loop}}
+  %1 = transform.tensor.make_loop_independent %0 {num_loops = 3}
+}
+
+// -----
+
+// CHECK: #[[$map:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>
+// CHECK: #[[$map1:.*]] = affine_map<()[s0, s1] -> (s0 - s1)>
+// CHECK-LABEL: func @make_empty_loop_independent(
+//  CHECK-SAME:     %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)
+func.func @make_empty_loop_independent(%lb: index, %ub: index, %step: index) {
+  // CHECK: scf.for %[[iv:.*]] = %[[lb]] to %[[ub]]
+  scf.for %i = %lb to %ub step %step {
+    // CHECK: %[[slice_sz:.*]] = affine.apply #[[$map]](%[[iv]])[%[[ub]]]
+    // CHECK: %[[empty_sz:.*]] = affine.apply #[[$map1]]()[%[[ub]], %[[lb]]]
+    // CHECK: %[[empty:.*]] = tensor.empty(%[[empty_sz]]) : tensor<?xf32>
+    // CHECK: %[[replacement:.*]] = tensor.extract_slice %[[empty]][0] [%[[slice_sz]]] [1]
+    %sz = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (%i)[%ub]
+    %empty = tensor.empty(%sz) : tensor<?xf32>
+    // CHECK: "dummy.some_use"(%[[replacement]])
+    "dummy.some_use"(%empty) : (tensor<?xf32>) -> ()
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.empty"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.tensor.make_loop_independent %0 {num_loops = 1}
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

hdrs = [
         "include/mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h",
         "include/mlir/Dialect/Tensor/Transforms/Passes.h",
         "include/mlir/Dialect/Tensor/Transforms/TransformUtils.h",
         "include/mlir/Dialect/Tensor/Transforms/Transforms.h",
     ],
     includes = ["include"],
     deps = [
         ":AffineDialect",
+        ":AffineTransforms",
         ":AffineUtils",
         ":ArithDialect",
         ":ArithUtils",
         ":BufferizationDialect",
         ":BufferizationTransforms",
         ":DialectUtils",
         ":FuncDialect",
         ":IR",
         ":LinalgDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":TensorDialect",
         ":TensorPassIncGen",
         ":TilingInterface",
         ":Transforms",
+        ":ValueBoundsOpInterface",
+        "//llvm:Support",
+    ],
+)
+
+td_library(
+    name = "TensorTransformOpsTdFiles",
+    srcs = [
+        "include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td",
+    ],
+    includes = ["include"],
+    deps = [
+        ":PDLDialect",
+        ":TransformDialectTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "TensorTransformOpsIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-op-decls"],
+            "include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td",
+    deps = [
+        ":TensorTransformOpsTdFiles",
+    ],
+)
+
+cc_library(
+    name = "TensorTransformOps",
+    srcs = glob(["lib/Dialect/Tensor/TransformOps/*.cpp"]),
+    hdrs = glob(["include/mlir/Dialect/Tensor/TransformOps/*.h"]),
+    includes = ["include"],
+    deps = [
+        ":AffineDialect",
+        ":IR",
+        ":PDLDialect",
+        ":SCFDialect",
+        ":TensorDialect",
+        ":TensorTransformOpsIncGen",
+        ":TensorTransforms",
+        ":TransformDialect",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "Rewrite",
     srcs = glob([
         "lib/Rewrite/*.cpp",
deps = [
         ":ShapeTransforms",
         ":ShapeTransformsPassIncGen",
         ":SparseTensorDialect",
         ":SparseTensorPipelines",
         ":SparseTensorTransforms",
         ":TensorDialect",
         ":TensorInferTypeOpInterfaceImpl",
         ":TensorTilingInterfaceImpl",
+        ":TensorTransformOps",
         ":TensorTransforms",
         ":TosaDialect",
         ":TosaToLinalg",
         ":TransformDialect",
         ":TransformDialectTransforms",
         ":Transforms",
         ":TransformsPassIncGen",
         ":VectorDialect",