[mlir][memref] Add transformation to do loop multi-buffering

[ThomasRaoux]

Summary

This transformation is useful to break dependency between consecutive loop iterations by increasing the size of a temporary buffer.
This is usually combined with heavy software pipelining.
Started a discourse thread in case we need this needs further discussions: https://discourse.llvm.org/t/loop-double-buffering-multi-buffering/59979

[ThomasRaoux]

clang-format

[bondhugula]

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

69–71

This should in general have nothing to do with an scf::for op -- you aren't transforming the loop in any way. Just use LoopLikeOpInterface.

100–105

Avoid generating lower-level arithmetic ops. Instead, just generate a modulo using affine.apply.

[ThomasRaoux]

Use affine.apply for arithmetic

[ThomasRaoux]

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

69–71

I need to access the single induction variable as well as the lower bound and step. This is not possible with LoopLikeInterface as it cannot be applied to any loops.

100–105

Good point, I didn't realize this was a valid affine map. Changed it.

[bondhugula]

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

77

Please document the return value.

77

This is a utility and should be returning LogicalResult instead of bool I think.

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

65

Please document the return value.

66

You don't need a triple doc comment here. Just the more detailed version on the declaration will suffice.

69–71

This should be pretty straightforward: you can just have the op interface method returning the IV (say getInductionVar()) return None for multi-dim loops (like for scf.parallel) and have this method fail in such a case just like you already do at line 87. It's going to be as simple and strictly more general. For the other two things you mention (lower bound and step), that's straightforward as well but the interface method will have to materialize some trivial IR (for eg. when it's a constant as in the case of affine.for):

// For IV.
Optional<Value> getSingleInductionVar();

1. lower bound

Optional<Value> getOrCreateLowerBound(OpBuilder &b);
Just return the lower bound SSA value (or materialize the step in case of affine.for)

2. step

Optional<Value> getOrCreateStep(OpBuilder &b);

With the above, it would be strictly more general than restricting it to scf::ForOp and as simple. And in your case of interest here, you'll get output IR identical to what you have in the test cases below.

mlir/test/Dialect/MemRef/multibuffering.mlir

67–68

You don't need the %{{.*}} = part.

[bondhugula]

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

77

The utilities are almost always named imperatively instead of a gerund. multiBuffering -> multiBuffer?

[ThomasRaoux]

Address review feedback

[ThomasRaoux]

I switched to using loop interface and added a test for affine.for op. Please take another look.

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

69–71

I added the interface and moved the code to use it. This means it now works on affine.for op. I added a test for it.

[bondhugula]

The comments that I had have been addressed. Dismissing change request.

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

23 ↗

(On Diff #409160)

Define using dyn_cast to avoid a double check in assert mode.

[bondhugula]

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

29 ↗

(On Diff #409160)

You are also collecting and erasing ops here.

36–53 ↗

(On Diff #409160)

Invert the conditional?

if (!....) {
   push_back
   continue;
}

...

70 ↗

(On Diff #409160)

Unchecked use of dyn_cast result. This should just be a single cast?

130–131 ↗

(On Diff #409160)

strides: append -> assign (since this is empty at this stage).

[ThomasRaoux]

Address review comments

[ThomasRaoux]

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

70 ↗

(On Diff #409160)

OpFoldResult/PointerUnion don't have cast they only have dyn_cast, I see it used like that in the rest of the code

[mravishankar]

LGTM!

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

89 ↗

(On Diff #409160)

Nit: if this user doesnt dominate...

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

 namespace tensor {
 class TensorDialect;
 } // namespace tensor
 namespace vector {
 class VectorDialect;
 } // namespace vector
 
 namespace memref {
-
+class AllocOp;
 //===----------------------------------------------------------------------===//
 // Patterns
 //===----------------------------------------------------------------------===//
 
 /// Collects a set of patterns to rewrite ops within the memref dialect.
 void populateExpandOpsPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns for folding memref.subview ops into consumer load/store ops
 /// into `patterns`.
 void populateFoldSubViewOpPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns that resolve `memref.dim` operations with values that are
 /// defined by operations that implement the
 /// `ReifyRankedShapeTypeShapeOpInterface`, in terms of shapes of its input
 /// operands.
 void populateResolveRankedShapeTypeResultDimsPatterns(
     RewritePatternSet &patterns);
 
 /// Appends patterns that resolve `memref.dim` operations with values that are
 /// defined by operations that implement the `InferShapedTypeOpInterface`, in
 /// terms of shapes of its input operands.
 void populateResolveShapedTypeResultDimsPatterns(RewritePatternSet &patterns);
 
+/// Transformation to do multi-buffering/array expansion to remove dependencies
+/// on the temporary allocation between consecutive loop iterations.
+/// This converts:
+/// ```
+/// %0 = memref.alloc() : memref<4x128xf32>
+/// scf.for %iv = %c1 to %c1024 step %c3 {
+///   memref.copy %1, %0 : memref<4x128xf32> to memref<4x128xf32>
+///   "some_use"(%0) : (memref<4x128xf32>) -> ()
+/// }
+/// ```
+/// into:
+/// ```
+/// %0 = memref.alloc() : memref<5x4x128xf32>
+/// scf.for %iv = %c1 to %c1024 step %c3 {
+///   %s = arith.subi %iv, %c1 : index
+///   %d = arith.divsi %s, %c3 : index
+///   %i = arith.remsi %d, %c5 : index
+///   %sv = memref.subview %0[%i, 0, 0] [1, 4, 128] [1, 1, 1] :
+///     memref<5x4x128xf32> to memref<4x128xf32, #map0>
+///   memref.copy %1, %sv : memref<4x128xf32> to memref<4x128xf32, #map0>
+///   "some_use"(%sv) : (memref<4x128xf32, $map0>) -> ()
+/// }
+/// ```
+bool multiBuffering(memref::AllocOp allocOp, unsigned multiplier);
+
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 /// Creates an instance of the ExpandOps pass that legalizes memref dialect ops
 /// to be convertible to LLVM. For example, `memref.reshape` gets converted to
 /// `memref_reinterpret_cast`.
 std::unique_ptr<Pass> createExpandOpsPass();

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

 add_mlir_dialect_library(MLIRMemRefTransforms
   ComposeSubView.cpp
   ExpandOps.cpp
   FoldSubViewOps.cpp
+  MultiBuffering.cpp
   NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen

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

+//===----------- MultiBuffering.cpp ---------------------------------------===//
+//
+// 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 file implements multi buffering transformation.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/Dominance.h"
+
+using namespace mlir;
+
+/// Return true if the op fully overwrite the given `buffer` value.
+static bool overrideBuffer(Operation *op, Value buffer) {
+  return isa<memref::CopyOp>(op) && cast<memref::CopyOp>(op).target() == buffer;
+}
+
+/// Replace the uses of `oldOp` with the given `val` and for subview uses
+/// propagate the type change. Changing the memref type may require propagating
+/// it through subview ops so we cannot just do a replaceAllUse but need to
+/// propagate the type change.
+static void replaceUsesAndPropagateType(Operation *oldOp, Value val,
+                                        OpBuilder &builder) {
+  SmallVector<Operation *> opToDelete;
+  SmallVector<OpOperand *> operandsToReplace;
+  for (OpOperand &use : oldOp->getUses()) {
+    auto subviewUse = dyn_cast<memref::SubViewOp>(use.getOwner());
+    if (subviewUse) {
+      builder.setInsertionPoint(subviewUse);
+      Type newType = memref::SubViewOp::inferRankReducedResultType(
+          subviewUse.getType().getRank(), val.getType().cast<MemRefType>(),
+          extractFromI64ArrayAttr(subviewUse.static_offsets()),
+          extractFromI64ArrayAttr(subviewUse.static_sizes()),
+          extractFromI64ArrayAttr(subviewUse.static_strides()));
+      Value newSubview = builder.create<memref::SubViewOp>(
+          subviewUse->getLoc(), newType.cast<MemRefType>(), val,
+          subviewUse.getMixedOffsets(), subviewUse.getMixedSizes(),
+          subviewUse.getMixedStrides());
+      replaceUsesAndPropagateType(subviewUse, newSubview, builder);
+      opToDelete.push_back(use.getOwner());
+      continue;
+    }
+    // Save the operand to and replace outside the loop to not invalidate the
+    // iterator.
+    operandsToReplace.push_back(&use);
+  }
+  for (OpOperand *operand : operandsToReplace)
+    operand->set(val);
+  // Clean up old subview ops.
+  for (Operation *op : opToDelete)
+    op->erase();
+}
+
+/// Transformation to do multi-buffering/array expansion to remove dependencies
+/// on the temporary allocation between consecutive loop iterations.
+// This is not a pattern as it requires propagating the new memref type to its
+// uses and requires updating subview ops.
+bool mlir::memref::multiBuffering(memref::AllocOp allocOp,
+                                  unsigned multiplier) {
+  DominanceInfo dom(allocOp->getParentOp());
+  scf::ForOp candidateLoop;
+  for (Operation *user : allocOp->getUsers()) {
+    auto parentLoop = user->getParentOfType<scf::ForOp>();
+    if (!parentLoop)
+      return false;
+    /// Make sure there is no loop carried dependency on the allocation.
+    if (!overrideBuffer(user, allocOp.getResult()))
+      continue;
+    // If the user doesn't dominate all the other users keep looking.
+    if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {
+          return otherUser->getNumRegions() != 0 ||
+                 !dom.dominates(user, otherUser);
+        }))
+      continue;
+    candidateLoop = parentLoop;
+    break;
+  }
+  if (!candidateLoop)
+    return false;
+  SmallVector<int64_t, 4> newShape(1, multiplier);
+  ArrayRef<int64_t> oldShape = allocOp.getType().getShape();
+  newShape.append(oldShape.begin(), oldShape.end());
+  auto newMemref = MemRefType::get(newShape, allocOp.getType().getElementType(),
+                                   MemRefLayoutAttrInterface(),
+                                   allocOp.getType().getMemorySpace());
+  OpBuilder builder(allocOp);
+  Location loc = allocOp->getLoc();
+  auto newAlloc = builder.create<memref::AllocOp>(loc, newMemref);
+  builder.setInsertionPoint(candidateLoop.getBody(),
+                            candidateLoop.getBody()->begin());
+  // Calculate the iteration index = ((iv - inital_val) / step) % multiplier.
+  AffineExpr induc = getAffineDimExpr(0, allocOp.getContext());
+  AffineExpr init = getAffineDimExpr(1, allocOp.getContext());
+  AffineExpr step = getAffineDimExpr(2, allocOp.getContext());
+  AffineExpr expr = ((induc - init).floorDiv(step)) % multiplier;
+  auto map = AffineMap::get(3, 0, expr);
+  std::array<Value, 3> operands = {candidateLoop.getInductionVar(),
+                                   candidateLoop.getLowerBound(),
+                                   candidateLoop.getStep()};
+  Value bufferIndex = builder.create<AffineApplyOp>(loc, map, operands);
+  SmallVector<OpFoldResult> offsets, sizes, strides;
+  offsets.push_back(bufferIndex);
+  offsets.append(oldShape.size(), builder.getIndexAttr(0));
+  strides.append(oldShape.size() + 1, builder.getIndexAttr(1));
+  sizes.push_back(builder.getIndexAttr(1));
+  for (int64_t size : oldShape)
+    sizes.push_back(builder.getIndexAttr(size));
+  auto dstMemref =
+      memref::SubViewOp::inferRankReducedResultType(
+          allocOp.getType().getRank(), newMemref, offsets, sizes, strides)
+          .cast<MemRefType>();
+  Value subview = builder.create<memref::SubViewOp>(loc, dstMemref, newAlloc,
+                                                    offsets, sizes, strides);
+  replaceUsesAndPropagateType(allocOp, subview, builder);
+  allocOp.erase();
+  return true;
+}

mlir/test/Dialect/MemRef/multibuffering.mlir

+// RUN: mlir-opt %s -allow-unregistered-dialect -test-multi-buffering=multiplier=5 -split-input-file | FileCheck %s
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 128 + s0 + d1)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (((d0 - d1) floordiv d2) mod 5)>
+
+// CHECK-LABEL: func @multi_buffer
+func @multi_buffer(%a: memref<1024x1024xf32>) {
+// CHECK-DAG: %[[A:.*]] = memref.alloc() : memref<5x4x128xf32>
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+  %0 = memref.alloc() : memref<4x128xf32>
+  %c1024 = arith.constant 1024 : index
+  %c1 = arith.constant 1 : index
+  %c3 = arith.constant 3 : index
+// CHECK: scf.for %[[IV:.*]] = %[[C1]]
+  scf.for %arg2 = %c1 to %c1024 step %c3 {
+// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]], %[[C1]], %[[C3]])
+// CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0, 0] [1, 4, 128] [1, 1, 1] : memref<5x4x128xf32> to memref<4x128xf32, #[[$MAP0]]>
+   %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
+    memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
+// CHECK: memref.copy %{{.*}}, %[[SV]] : memref<4x128xf32, #{{.*}}> to memref<4x128xf32, #[[$MAP0]]>
+   memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
+// CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, #[[$MAP0]]>) -> ()
+    "some_use"(%0) : (memref<4x128xf32>) -> ()
+// CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, #[[$MAP0]]>) -> ()
+   "some_use"(%0) : (memref<4x128xf32>) -> ()
+  }
+  return
+}
+
+// -----
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 128 + s0 + d1)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (((d0 - d1) floordiv d2) mod 5)>
+
+// CHECK-LABEL: func @multi_buffer_subview_use
+func @multi_buffer_subview_use(%a: memref<1024x1024xf32>) {
+// CHECK-DAG: %[[A:.*]] = memref.alloc() : memref<5x4x128xf32>
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+  %0 = memref.alloc() : memref<4x128xf32>
+  %c1024 = arith.constant 1024 : index
+  %c1 = arith.constant 1 : index
+  %c3 = arith.constant 3 : index
+// CHECK: scf.for %[[IV:.*]] = %[[C1]]
+  scf.for %arg2 = %c1 to %c1024 step %c3 {
+// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]], %[[C1]], %[[C3]])
+// CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0, 0] [1, 4, 128] [1, 1, 1] : memref<5x4x128xf32> to memref<4x128xf32, #[[$MAP0]]>
+   %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
+    memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
+// CHECK: memref.copy %{{.*}}, %[[SV]] : memref<4x128xf32, #{{.*}}> to memref<4x128xf32, #[[$MAP0]]>
+   memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
+// CHECK: %[[SV1:.*]] = memref.subview %[[SV]][0, 1] [4, 127] [1, 1] : memref<4x128xf32, #[[$MAP0]]> to memref<4x127xf32, #[[$MAP0]]>
+   %s = memref.subview %0[0, 1] [4, 127] [1, 1] :
+      memref<4x128xf32> to memref<4x127xf32, affine_map<(d0, d1) -> (d0 * 128 + d1 + 1)>>
+// CHECK: "some_use"(%[[SV1]]) : (memref<4x127xf32, #[[$MAP0]]>) -> ()
+   "some_use"(%s) : (memref<4x127xf32, affine_map<(d0, d1) -> (d0 * 128 + d1 + 1)>>) -> ()
+// CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, #[[$MAP0]]>) -> ()
+   "some_use"(%0) : (memref<4x128xf32>) -> ()
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @multi_buffer_negative
+func @multi_buffer_negative(%a: memref<1024x1024xf32>) {
+// CHECK-NOT: %{{.*}} = memref.alloc() : memref<5x4x128xf32>
+//     CHECK: %{{.*}} = memref.alloc() : memref<4x128xf32>
+  %0 = memref.alloc() : memref<4x128xf32>
+  %c1024 = arith.constant 1024 : index
+  %c0 = arith.constant 0 : index
+  %c3 = arith.constant 3 : index
+  scf.for %arg2 = %c0 to %c1024 step %c3 {
+   "blocking_use"(%0) : (memref<4x128xf32>) -> ()
+   %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
+    memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
+   memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
+   "some_use"(%0) : (memref<4x128xf32>) -> ()
+  }
+  return
+}

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

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRMemRefTestPasses
   TestComposeSubView.cpp
+  TestMultiBuffering.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
+  MLIRMemRef
   MLIRMemRefTransforms
   MLIRTestDialect
   )
 
 target_include_directories(MLIRMemRefTestPasses
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../Test
   ${CMAKE_CURRENT_BINARY_DIR}/../Test
   )

mlir/test/lib/Dialect/MemRef/TestMultiBuffering.cpp

+//===- TestComposeSubView.cpp - Test composed subviews --------------------===//
+//
+// 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/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+struct TestMultiBufferingPass
+    : public PassWrapper<TestMultiBufferingPass, OperationPass<FuncOp>> {
+  TestMultiBufferingPass() = default;
+  TestMultiBufferingPass(const TestMultiBufferingPass &pass)
+      : PassWrapper(pass) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect>();
+  }
+
+  StringRef getArgument() const final { return "test-multi-buffering"; }
+  StringRef getDescription() const final {
+    return "Test multi buffering transformation";
+  }
+  void runOnOperation() override;
+  Option<unsigned> multiplier{
+      *this, "multiplier",
+      llvm::cl::desc(
+          "Decide how many versions of the buffer should be created,"),
+      llvm::cl::init(2)};
+};
+
+void TestMultiBufferingPass::runOnOperation() {
+  SmallVector<memref::AllocOp> allocs;
+  getOperation().walk(
+      [&allocs](memref::AllocOp alloc) { allocs.push_back(alloc); });
+  for (memref::AllocOp alloc : allocs)
+    multiBuffering(alloc, multiplier);
+}
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestMultiBuffering() {
+  PassRegistration<TestMultiBufferingPass>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestGpuSerializeToHsacoPass();
 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestComposeSubView();
+void registerTestMultiBuffering();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgControlFuseByExpansion();
 void registerTestLinalgDistribution();

 #endif
   mlir::test::registerTestComprehensiveFunctionBufferize();
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestExpandTanhPass();
   mlir::test::registerTestComposeSubView();
+  mlir::test::registerTestMultiBuffering();
   mlir::test::registerTestGpuParallelLoopMappingPass();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLinalgCodegenStrategy();
   mlir::test::registerTestLinalgControlFuseByExpansion();
   mlir::test::registerTestLinalgDistribution();
   mlir::test::registerTestLinalgElementwiseFusion();

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

deps = [
         ":AffineUtils",
         ":ArithmeticDialect",
         ":ArithmeticTransforms",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefDialect",
         ":MemRefPassIncGen",
         ":Pass",
+        ":SCFDialect",
         ":StandardOps",
         ":Support",
         ":TensorDialect",
         ":Transforms",
         ":VectorOps",
         "//llvm:Support",
     ],
 )

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

 cc_library(
     name = "TestMemRef",
     srcs = glob(["lib/Dialect/MemRef/*.cpp"]),
     defines = ["MLIR_CUDA_CONVERSIONS_ENABLED"],
     includes = ["lib/Dialect/Test"],
     deps = [
         ":TestDialect",
         "//mlir:Affine",
+        "//mlir:MemRefDialect",
         "//mlir:MemRefTransforms",
         "//mlir:Pass",
         "//mlir:Transforms",
     ],
 )
 
 cc_library(
     name = "TestSCF",