[mlir][memref] Add transform to make alloca ops loop-independent

[springerm]

Summary

Add a helper function that makes dynamic sizes of memref.alloca ops independent of a given set of values. This functionality can be used to make dynamic allocations hoistable from loops.

Depends On: D143910

[springerm]

add test case

[dcaballe]

mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h

129

can't parse this but maybe it's just me: who's -> whose?

138

same

176

Awesome!

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

111

ub to var

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

104

same

114

We can hoist this to a utility function as it looks like we may have to add support for more ops. We should also set/restore the insertion point per user (iteration) to prevent accidental uses of a wrong insertion point.

150

Any way we could minimize the risk of creating invalid code here?

[springerm]

address comments

mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h

129

ah yes i always mix those two up

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

104

Here we actually want unrealizedConversions.size() in the check because new ops are getting added to unrealizedConversions in the loop body.

150

We will need an interface for this eventually. This interface can be used for more than just this transform here. E.g., I think it would also be useful for canonicalization in general: it could be used to canonicalize away/simplify layout maps. There, we have the same problem, that some of the IR has to be re-created in the worst case.

(I'm not aware of any ops where we would generate invalid IR at the moment.)

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

OpBuilder<(ins "MemRefType":$memrefType, "ValueRange":$dynamicSizes,
       $_state.addOperands(dynamicSizes);
       $_state.addOperands(symbolOperands);
       $_state.addAttribute(getOperandSegmentSizeAttr(),
           $_builder.getDenseI32ArrayAttr({
               static_cast<int32_t>(dynamicSizes.size()),
               static_cast<int32_t>(symbolOperands.size())}));
       if (alignment)
         $_state.addAttribute(getAlignmentAttrStrName(), alignment);
-    }]>];
+    }]>,
+    OpBuilder<(ins "ArrayRef<OpFoldResult>":$sizes, "Type":$elementType,
+                   CArg<"Attribute", "{}">:$memorySpace), [{
+      SmallVector<int64_t> staticShape;
+      SmallVector<Value> dynamicSizes;
+      dispatchIndexOpFoldResults(sizes, dynamicSizes, staticShape);
+      MemRefLayoutAttrInterface layout;
+      MemRefType memrefType = MemRefType::get(staticShape, elementType, layout,
+                                              memorySpace);
+      return build($_builder, $_state, memrefType, dynamicSizes);
+    }]>
+  ];
 
   let extraClassDeclaration = [{
     static StringRef getAlignmentAttrStrName() { return "alignment"; }
 
     MemRefType getType() { return getResult().getType().cast<MemRefType>(); }
 
     SmallVector<OpFoldResult> getMixedSizes() {
       SmallVector<OpFoldResult> result;

mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

def MemRefExtractAddressComputationsOp :
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &transformResults,
         ::mlir::transform::TransformState &state);
   }];
 }
+
+def MemRefMakeLoopIndependentOp
+    : Op<Transform_Dialect, "memref.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
+    `memref.subview` inside the loop.
+
+    Currently supported operations are:
+    - memref.alloca: Replaced with a new memref.alloca with upper bound sizes,
+      followed by a memref.subview.
+
+    #### 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:
+    - memref.alloca: The returned handle points to the memref.subview 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 // MEMREF_TRANSFORM_OPS

mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MEMREF_TRANSFORMS_TRANSFORMS_H
 #define MLIR_DIALECT_MEMREF_TRANSFORMS_TRANSFORMS_H
 
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
+class OpBuilder;
 class RewritePatternSet;
 class RewriterBase;
+class Value;
+class ValueRange;
 
 namespace arith {
 class WideIntEmulationConverter;
 } // namespace arith
 
 namespace memref {
 class AllocOp;
+class AllocaOp;
+
 //===----------------------------------------------------------------------===//
 // Patterns
 //===----------------------------------------------------------------------===//
 
 /// Collects a set of patterns to rewrite ops within the memref dialect.
 void populateExpandOpsPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns for folding memref aliasing ops into consumer load/store

 ///
 /// Will be rewritten in:
 /// ```mlir
 /// %new_base = memref.subview %base[%off0,...][1,...][1,...]
 /// memref.load %new_base[%c0,...]
 /// ```
 void populateExtractAddressComputationsPatterns(RewritePatternSet &patterns);
 
+/// Build a new memref::AllocaOp whose dynamic sizes are independent of all
+/// given independencies. If the op is already independent of all
+/// independencies, the same AllocaOp result is returned.
+///
+/// Failure indicates the no suitable upper bound for the dynamic sizes could be
+/// found.
+FailureOr<Value> buildIndependentOp(OpBuilder &b, AllocaOp allocaOp,
+                                    ValueRange independencies);
+
+/// Build a new memref::AllocaOp whose dynamic sizes are independent of all
+/// given independencies. If the op is already independent of all
+/// independencies, the same AllocaOp result is returned.
+///
+/// The original AllocaOp is replaced with the new one, wrapped in a SubviewOp.
+/// The result type of the replacement is different from the original allocation
+/// type: it has the same shape, but a different layout map. This function
+/// updates all users that do not have a memref result or memref region block
+/// argument, and some frequently used memref dialect ops (such as
+/// memref.subview). It does not update other uses such as the init_arg of an
+/// scf.for op. Such uses are wrapped in unrealized_conversion_cast.
+///
+/// Failure indicates the no suitable upper bound for the dynamic sizes could be
+/// found.
+///
+/// Example (make independent of %iv):
+/// ```
+/// scf.for %iv = %c0 to %sz step %c1 {
+///   %0 = memref.alloca(%iv) : memref<?xf32>
+///   %1 = memref.subview %0[0][5][1] : ...
+///   linalg.generic outs(%1 : ...) ...
+///   %2 = scf.for ... iter_arg(%arg0 = %0) ...
+///   ...
+/// }
+/// ```
+///
+/// The above IR is rewritten to:
+///
+/// ```
+/// scf.for %iv = %c0 to %sz step %c1 {
+///   %0 = memref.alloca(%sz - 1) : memref<?xf32>
+///   %0_subview = memref.subview %0[0][%iv][1]
+///       : memref<?xf32> to memref<?xf32, #map>
+///   %1 = memref.subview %0_subview[0][5][1] : ...
+///   linalg.generic outs(%1 : ...) ...
+///   %cast = unrealized_conversion_cast %0_subview
+///       : memref<?xf32, #map> to memref<?xf32>
+///   %2 = scf.for ... iter_arg(%arg0 = %cast) ...
+///  ...
+/// }
+/// ```
+FailureOr<Value> replaceWithIndependentOp(RewriterBase &rewriter,
+                                          memref::AllocaOp allocaOp,
+                                          ValueRange independencies);
+
 } // namespace memref
 } // namespace mlir
 
 #endif

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

 //===- MemRefTransformOps.cpp - Implementation of Memref 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/MemRef/TransformOps/MemRefTransformOps.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
transform::MemRefExtractAddressComputationsOp::applyToOne(
   if (failed(applyPatternsAndFoldGreedily(target, std::move(patterns))))
     return emitDefaultDefiniteFailure(target);
 
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
+// MemRefMakeLoopIndependentOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::MemRefMakeLoopIndependentOp::applyToOne(
+    Operation *target, transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  // Gather IVs.
+  SmallVector<Value> ivs;
+  Operation *nextOp = target;
+  for (uint64_t i = 0, e = getNumLoops(); i < e; ++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 allocaOp = dyn_cast<memref::AllocaOp>(target)) {
+    replacement = memref::replaceWithIndependentOp(rewriter, allocaOp, ivs);
+  } else {
+    DiagnosedSilenceableFailure diag = emitSilenceableError()
+                                       << "unsupported target op";
+    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;
+  }
+  results.push_back(replacement->getDefiningOp());
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 class MemRefTransformDialectExtension
     : public transform::TransformDialectExtension<
           MemRefTransformDialectExtension> {
 public:

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

 add_mlir_dialect_library(MLIRMemRefTransforms
   BufferizableOpInterfaceImpl.cpp
   ComposeSubView.cpp
   ExpandOps.cpp
   ExpandStridedMetadata.cpp
   EmulateWideInt.cpp
   ExtractAddressComputations.cpp
   FoldMemRefAliasOps.cpp
+  IndependenceTransforms.cpp
   MultiBuffer.cpp
   NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
   RuntimeOpVerification.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
+  MLIRAffineTransforms
   MLIRAffineUtils
   MLIRArithDialect
   MLIRArithTransforms
   MLIRBufferizationDialect
   MLIRDialectUtils
   MLIRFuncDialect
   MLIRGPUOps
   MLIRInferTypeOpInterface
   MLIRLoopLikeInterface
   MLIRMemRefDialect
   MLIRNVGPUDialect
   MLIRPass
   MLIRTensorDialect
   MLIRTransforms
+  MLIRValueBoundsOpInterface
   MLIRVectorDialect
 )
 

mlir/lib/Dialect/MemRef/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/MemRef/Transforms/Transforms.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Transforms/Transforms.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+using namespace mlir;
+using namespace mlir::memref;
+
+/// 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 affine::materializeComputedBound(b, loc, boundMap, mapOperands);
+}
+
+FailureOr<Value> memref::buildIndependentOp(OpBuilder &b,
+                                            memref::AllocaOp allocaOp,
+                                            ValueRange independencies) {
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(allocaOp);
+  Location loc = allocaOp.getLoc();
+
+  SmallVector<OpFoldResult> newSizes;
+  for (OpFoldResult ofr : allocaOp.getMixedSizes()) {
+    auto ub = makeIndependent(b, loc, ofr, independencies);
+    if (failed(ub))
+      return failure();
+    newSizes.push_back(*ub);
+  }
+
+  // Return existing memref::AllocaOp if nothing has changed.
+  if (llvm::equal(allocaOp.getMixedSizes(), newSizes))
+    return allocaOp.getResult();
+
+  // Create a new memref::AllocaOp.
+  Value newAllocaOp =
+      b.create<AllocaOp>(loc, newSizes, allocaOp.getType().getElementType());
+
+  // Create a memref::SubViewOp.
+  SmallVector<OpFoldResult> offsets(newSizes.size(), b.getIndexAttr(0));
+  SmallVector<OpFoldResult> strides(newSizes.size(), b.getIndexAttr(1));
+  return b
+      .create<SubViewOp>(loc, newAllocaOp, offsets, allocaOp.getMixedSizes(),
+                         strides)
+      .getResult();
+}
+
+/// Push down an UnrealizedConversionCastOp past a SubViewOp.
+static UnrealizedConversionCastOp
+propagateSubViewOp(RewriterBase &rewriter,
+                   UnrealizedConversionCastOp conversionOp, SubViewOp op) {
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(op);
+  auto newResultType =
+      SubViewOp::inferRankReducedResultType(
+          op.getType().getShape(), op.getSourceType(), op.getMixedOffsets(),
+          op.getMixedSizes(), op.getMixedStrides())
+          .cast<MemRefType>();
+  Value newSubview = rewriter.create<SubViewOp>(
+      op.getLoc(), newResultType, conversionOp.getOperand(0),
+      op.getMixedOffsets(), op.getMixedSizes(), op.getMixedStrides());
+  auto newConversionOp = rewriter.create<UnrealizedConversionCastOp>(
+      op.getLoc(), op.getType(), newSubview);
+  rewriter.replaceAllUsesWith(op.getResult(), newConversionOp->getResult(0));
+  return newConversionOp;
+}
+
+/// Given an original op and a new, modified op with the same number of results,
+/// whose memref return types may differ, replace all uses of the original op
+/// with the new op and propagate the new memref types through the IR.
+///
+/// Example:
+/// %from = memref.alloca(%sz) : memref<?xf32>
+/// %to = memref.subview ... : ... to memref<?xf32, strided<[1], offset: ?>>
+/// memref.store %cst, %from[%c0] : memref<?xf32>
+///
+/// In the above example, all uses of %from are replaced with %to. This can be
+/// done directly for ops such as memref.store. For ops that have memref results
+/// (e.g., memref.subview), the result type may depend on the operand type, so
+/// we cannot just replace all uses. There is special handling for common memref
+/// ops. For all other ops, unrealized_conversion_cast is inserted.
+static void replaceAndPropagateMemRefType(RewriterBase &rewriter,
+                                          Operation *from, Operation *to) {
+  assert(from->getNumResults() == to->getNumResults() &&
+         "expected same number of results");
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPointAfter(to);
+
+  // Wrap new results in unrealized_conversion_cast and replace all uses of the
+  // original op.
+  SmallVector<UnrealizedConversionCastOp> unrealizedConversions;
+  for (const auto &it :
+       llvm::enumerate(llvm::zip(from->getResults(), to->getResults()))) {
+    unrealizedConversions.push_back(rewriter.create<UnrealizedConversionCastOp>(
+        to->getLoc(), std::get<0>(it.value()).getType(),
+        std::get<1>(it.value())));
+    rewriter.replaceAllUsesWith(from->getResult(it.index()),
+                                unrealizedConversions.back()->getResult(0));
+  }
+
+  // Push unrealized_conversion_cast ops further down in the IR. I.e., try to
+  // wrap results instead of operands in a cast.
+  for (int i = 0; i < static_cast<int>(unrealizedConversions.size()); ++i) {
+    UnrealizedConversionCastOp conversion = unrealizedConversions[i];
+    assert(conversion->getNumOperands() == 1 &&
+           conversion->getNumResults() == 1 &&
+           "expected single operand and single result");
+    SmallVector<Operation *> users = llvm::to_vector(conversion->getUsers());
+    for (Operation *user : users) {
+      // Handle common memref dialect ops that produce new memrefs and must
+      // be recreated with the new result type.
+      if (auto subviewOp = dyn_cast<SubViewOp>(user)) {
+        unrealizedConversions.push_back(
+            propagateSubViewOp(rewriter, conversion, subviewOp));
+        continue;
+      }
+
+      // TODO: Other memref ops such as memref.collapse_shape/expand_shape
+      // should also be handled here.
+
+      // Skip any ops that produce MemRef result or have MemRef region block
+      // arguments. These may need special handling (e.g., scf.for).
+      if (llvm::any_of(user->getResultTypes(),
+                       [](Type t) { return isa<MemRefType>(t); }))
+        continue;
+      if (llvm::any_of(user->getRegions(), [](Region &r) {
+            return llvm::any_of(r.getArguments(), [](BlockArgument bbArg) {
+              return isa<MemRefType>(bbArg.getType());
+            });
+          }))
+        continue;
+
+      // For all other ops, we assume that we can directly replace the operand.
+      // This may have to be revised in the future; e.g., there may be ops that
+      // do not support non-identity layout maps.
+      for (OpOperand &operand : user->getOpOperands()) {
+        if (auto castOp =
+                operand.get().getDefiningOp<UnrealizedConversionCastOp>()) {
+          rewriter.updateRootInPlace(
+              user, [&]() { operand.set(conversion->getOperand(0)); });
+        }
+      }
+    }
+  }
+
+  // Erase all unrealized_conversion_cast ops without uses.
+  for (auto op : unrealizedConversions)
+    if (op->getUses().empty())
+      rewriter.eraseOp(op);
+}
+
+FailureOr<Value> memref::replaceWithIndependentOp(RewriterBase &rewriter,
+                                                  memref::AllocaOp allocaOp,
+                                                  ValueRange independencies) {
+  auto replacement =
+      memref::buildIndependentOp(rewriter, allocaOp, independencies);
+  if (failed(replacement))
+    return failure();
+  replaceAndPropagateMemRefType(rewriter, allocaOp,
+                                replacement->getDefiningOp());
+  return replacement;
+}

mlir/test/Dialect/MemRef/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
+
+// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 - 1)>
+// CHECK-LABEL: func @make_alloca_loop_independent(
+//  CHECK-SAME:     %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)
+func.func @make_alloca_loop_independent(%lb: index, %ub: index, %step: index) {
+  %cst = arith.constant 5.5 : f32
+  %c0 = arith.constant 0 : index
+  // CHECK: scf.for %[[iv:.*]] = %[[lb]] to %[[ub]]
+  scf.for %i = %lb to %ub step %step {
+    // CHECK: %[[sz:.*]] = affine.apply #[[$map]]()[%[[ub]]]
+    // CHECK: %[[alloca:.*]] = memref.alloca(%[[sz]])
+    // CHECK: %[[subview:.*]] = memref.subview %[[alloca]][0] [%[[iv]]] [1] : memref<?xf32> to memref<?xf32, strided<[1]>>
+    // CHECK: %[[cast:.*]] = builtin.unrealized_conversion_cast %[[subview]] : memref<?xf32, strided<[1]>> to memref<?xf32>
+    %alloc = memref.alloca(%i) : memref<?xf32>
+
+    // memref.subview has special handling.
+    // CHECK: %[[subview2:.*]] = memref.subview %[[subview]][1] [5] [1] : memref<?xf32, strided<[1]>> to memref<5xf32, strided<[1], offset: 1>>
+    %view = memref.subview %alloc[1][5][1] : memref<?xf32> to memref<5xf32, strided<[1], offset: 1>>
+
+    // This op takes a memref but does not produce one. The new alloc is used
+    // directly.
+    // CHECK: "test.some_use"(%[[subview2]])
+    "test.some_use"(%view) : (memref<5xf32, strided<[1], offset: 1>>) -> ()
+
+    // This op produces a memref, so the new alloc cannot be used directly.
+    // It is wrapped in a unrealized_conversion_cast.
+    // CHECK: "test.another_use"(%[[cast]]) : (memref<?xf32>) -> memref<?xf32>
+    "test.another_use"(%alloc) : (memref<?xf32>) -> (memref<?xf32>)
+
+    // CHECK: memref.store %{{.*}}, %[[subview]]
+    memref.store %cst, %alloc[%c0] : memref<?xf32>
+  }
+  return
+}
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["memref.alloca"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.memref.make_loop_independent %0 {num_loops = 1}
+}
+
+// -----
+
+// CHECK: #[[$map:.*]] = affine_map<(d0) -> (-d0 + 128)>
+// CHECK-LABEL: func @make_alloca_loop_independent_static(
+func.func @make_alloca_loop_independent_static(%step: index) {
+  %cst = arith.constant 5.5 : f32
+  %c0 = arith.constant 0 : index
+  %ub = arith.constant 128 : index
+  // CHECK: scf.for %[[iv:.*]] =
+  scf.for %i = %c0 to %ub step %step {
+    // CHECK: %[[sz:.*]] = affine.apply #[[$map]](%[[iv]])
+    %sz = affine.apply affine_map<(d0)[s0] -> (-d0 + s0)>(%i)[%ub]
+
+    // CHECK: %[[alloca:.*]] = memref.alloca() : memref<128xf32>
+    // CHECK: %[[subview:.*]] = memref.subview %[[alloca]][0] [%[[sz]]] [1] : memref<128xf32> to memref<?xf32, strided<[1]>>
+    %alloc = memref.alloca(%sz) : memref<?xf32>
+
+    // CHECK: memref.store %{{.*}}, %[[subview]]
+    memref.store %cst, %alloc[%c0] : memref<?xf32>
+
+    // CHECK: vector.print %[[sz]]
+    %dim = memref.dim %alloc, %c0 : memref<?xf32>
+    vector.print %dim : index
+  }
+  return
+}
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["memref.alloca"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.memref.make_loop_independent %0 {num_loops = 1}
+}

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

srcs = glob(
             "lib/Dialect/MemRef/Transforms/*.cpp",
             "lib/Dialect/MemRef/Transforms/*.h",
         ],
     ),
     hdrs = glob(["include/mlir/Dialect/MemRef/Transforms/*.h"]),
     includes = ["include"],
     deps = [
         ":AffineDialect",
+        ":AffineTransforms",
         ":AffineUtils",
         ":ArithDialect",
         ":ArithTransforms",
         ":ArithUtils",
         ":BufferizationDialect",
         ":ControlFlowDialect",
         ":DialectUtils",
         ":FuncDialect",
         ":GPUDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":MemRefPassIncGen",
         ":NVGPUDialect",
         ":Pass",
         ":RuntimeVerifiableOpInterface",
         ":Support",
         ":TensorDialect",
         ":Transforms",
+        ":ValueBoundsOpInterface",
         ":VectorDialect",
         "//llvm:Support",
     ],
 )
 
 td_library(
     name = "MemRefTransformOpsTdFiles",
     srcs = [
deps = [
         ":ArithDialect",
         ":IR",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":MemRefTransformOpsIncGen",
         ":MemRefTransforms",
         ":NVGPUDialect",
         ":PDLDialect",
+        ":SCFDialect",
         ":TransformDialect",
         ":TransformUtils",
         ":VectorDialect",
         "//llvm:Support",
     ],
 )
 
 ##---------------------------------------------------------------------------##