[mlir][MemRef] Add patterns to extract address computations

[qcolombet]

Summary

This patch adds patterns to rewrite memory accesses such that the resulting
accesses are only using a base pointer.
E.g.,

mlir
memref.load %base[%off0, ...]

Will be rewritten in:

mlir
%new_base = memref.subview %base[%off0,...][1,...][1,...]
memref.load %new_base[%c0,...]

The idea behind these patterns is to offer a way to more gradually lower
address computations.

These patterns are the exact opposite of FoldMemRefAliasOps.
I've implemented the support of only three operations in this patch:

• memref.load
• memref.store
• nvgpu.ldmatrix

Going forward we may want to provide an interface for these rewritings (and
the ones in FoldMemRefAliasOps.)
One step at a time!

[nicolasvasilache]

Thanks, LGTM, make sense to keep the interface work for later.

One note though, we also need those on vector.transfer soonish.

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

65

Similary to https://reviews.llvm.org/D146624 we want all these populate to live in Transform/Transforms.h

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

7

It would be nice to find a name that somehow ties these logically to FoldMemRefAliasOps.cpp but I odn't have a better suggestion atm.

[qcolombet]

• Add support for vector.transfer_read/write
• Move the populate function in its own header (we'll want to move others as well in an NFC patch)
• Fix the dropping of the nontemporal flag (and add a test for that)
• Refactor getSrcMemRef to not choke on tensor type
• Refactor getViewSizeForEachDim to use the same implementation for ldmatrix, vector.transfer_read/write
• Add tests for vector.transfer_read/write

[qcolombet]

• Add missing vector dialect dependency to the test-extract-address-computations pass

[qcolombet]

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

7

Yeah, I can't find anything else either.
Maybe we could rename the other one?

No good idea there either though!

[nicolasvasilache]

thanks for adding vector.transfer too, LGTM!

[nicolasvasilache]

sorry clicked too soon, added some comments

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

35 ↗

(On Diff #508034)

also move the other populates here ?

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

119

For my own education, could you please point me to explanations about \see and other similar things I have seen you use over time ?
I'd happily follow suit and try to improve my own comments.

224

llvm::function_ref instad of C style funptrs ?

247

return rewriter.notifyMatchFailure(op, "...") everywhere plz, this greatly simplifies understanding when debugging

270

Add a good comment that strides are handled properly by composition of subviews here please.
This is not true for the reverse transformation and the precondition for vector.transfer checks and fails to match if strides are not all one: the forward and inverse problems are not symmetrical unfortunately.

mlir/test/Dialect/MemRef/extract-address-computations.mlir

288

some examples with strides please that mirror my previous comment

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

30

can we just wrap this in a transform op and avoid creating unused code?

[qcolombet]

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

35 ↗

(On Diff #508034)

I thought we would do that in an NFC commit.
Which I'll do after that one land.

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

119

Here is the complete doc for the doxygen "keywords":
https://www.doxygen.nl/manual/commands.html

The ones I use that may not be that common are:

• \pre, \post: Document a pre-post condition
• \see: Create a link to another item (class, function, etc.) also documented in doxygen

But maybe what I think is common or not is biased :P (\note, \code, \p, etc.).

224

Good point, I always forgot about these!
(I start with std::function, then remember that templates don't like them and resort to plain C. I need to update my mental model :)).

247

Good point!

270

👍

mlir/test/Dialect/MemRef/extract-address-computations.mlir

288

👍

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

30

Sure thing.

[qcolombet]

• Remove the test pass
• Add a tranform op instead
• Add a comment for strides when creating the subview
• Add a test for non-1 strides (only for vector.transfer_write, @nicolasvasilache let me know if you want more)
• Report failures with notifyMatchFailure

[qcolombet]

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

224

Alright looks like we can't use them either. They suffer the same issue as std::function:

error: non-type template parameters of class type only available with ‘-std=c++20’ or ‘-std=gnu++20’

Am I missing something?

mlir/test/Dialect/MemRef/extract-address-computations.mlir

288

Only did one for vector.transfer_write (see test_transfer_write_op_with_strides).
Let me know if you want some for the others.

[nicolasvasilache]

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

35 ↗

(On Diff #508034)

sgtm!

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

224

oh well ..

mlir/test/Dialect/MemRef/extract-address-computations.mlir

288

that's great, thanks!

[qcolombet]

• Fix bazel build

[qcolombet]

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

35 ↗

(On Diff #508034)

Done in faafd26c4d58

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

 
 /// 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
 /// ops into `patterns`.
 void populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns);
 
+/// Appends patterns for extracting address computations from the instructions
+/// with memory accesses such that these memory accesses use only a base
+/// pointer.
+///
+/// For instance,
+/// ```mlir
+/// memref.load %base[%off0, ...]
+/// ```
+///
+/// Will be rewritten in:
+/// ```mlir
+/// %new_base = memref.subview %base[%off0,...][1,...][1,...]
+/// memref.load %new_base[%c0,...]
+/// ```
+void populateExtractAddressComputationsPatterns(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

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
   MultiBuffer.cpp
   NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
   RuntimeOpVerification.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineUtils
   MLIRArithDialect
   MLIRArithTransforms
   MLIRBufferizationDialect
   MLIRFuncDialect
   MLIRGPUOps
   MLIRInferTypeOpInterface
   MLIRLoopLikeInterface
   MLIRMemRefDialect
+  MLIRNVGPUDialect
   MLIRPass
   MLIRTensorDialect
   MLIRTransforms
   MLIRVectorDialect
 )
 

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

+//===- ExtractAddressCmoputations.cpp - Extract address computations  -----===//
+//
+// 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 transformation pass rewrites loading/storing from/to a memref with
+/// offsets into loading/storing from/to a subview and without any offset on
+/// the instruction itself.
+//
+//===----------------------------------------------------------------------===//
+
+#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/NVGPU/IR/NVGPUDialect.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+
+using namespace mlir;
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Helper functions for the `load base[off0...]`
+//  => `load (subview base[off0...])[0...]` pattern.
+//===----------------------------------------------------------------------===//
+
+// Matches getSrcMemRef specs for LoadOp.
+// \see LoadStoreLikeOpRewriter.
+static Value getLoadOpSrcMemRef(memref::LoadOp loadOp) {
+  return loadOp.getMemRef();
+}
+
+// Matches rebuildOpFromAddressAndIndices specs for LoadOp.
+// \see LoadStoreLikeOpRewriter.
+static memref::LoadOp rebuildLoadOp(RewriterBase &rewriter,
+                                    memref::LoadOp loadOp, Value srcMemRef,
+                                    ArrayRef<Value> indices) {
+  Location loc = loadOp.getLoc();
+  return rewriter.create<memref::LoadOp>(loc, srcMemRef, indices);
+}
+
+// Matches getViewSizeForEachDim specs for LoadOp.
+// \see LoadStoreLikeOpRewriter.
+static SmallVector<OpFoldResult>
+getLoadOpViewSizeForEachDim(RewriterBase &rewriter, memref::LoadOp loadOp) {
+  MemRefType ldTy = loadOp.getMemRefType();
+  unsigned loadRank = ldTy.getRank();
+  return SmallVector<OpFoldResult>(loadRank, rewriter.getIndexAttr(1));
+}
+
+//===----------------------------------------------------------------------===//
+// Helper functions for the `store val, base[off0...]`
+//  => `store val, (subview base[off0...])[0...]` pattern.
+//===----------------------------------------------------------------------===//
+
+// Matches getSrcMemRef specs for StoreOp.
+// \see LoadStoreLikeOpRewriter.
+static Value getStoreOpSrcMemRef(memref::StoreOp storeOp) {
+  return storeOp.getMemRef();
+}
+
+// Matches rebuildOpFromAddressAndIndices specs for StoreOp.
+// \see LoadStoreLikeOpRewriter.
+static memref::StoreOp rebuildStoreOp(RewriterBase &rewriter,
+                                      memref::StoreOp storeOp, Value srcMemRef,
+                                      ArrayRef<Value> indices) {
+  Location loc = storeOp.getLoc();
+  return rewriter.create<memref::StoreOp>(loc, storeOp.getValueToStore(),
+                                          srcMemRef, indices);
+}
+
+// Matches getViewSizeForEachDim specs for StoreOp.
+// \see LoadStoreLikeOpRewriter.
+static SmallVector<OpFoldResult>
+getStoreOpViewSizeForEachDim(RewriterBase &rewriter, memref::StoreOp storeOp) {
+  MemRefType ldTy = storeOp.getMemRefType();
+  unsigned loadRank = ldTy.getRank();
+  return SmallVector<OpFoldResult>(loadRank, rewriter.getIndexAttr(1));
+}
+
+//===----------------------------------------------------------------------===//
+// Helper functions for the `ldmatrix base[off0...]`
+//  => `ldmatrix (subview base[off0...])[0...]` pattern.
+//===----------------------------------------------------------------------===//
+
+// Matches getSrcMemRef specs for LdMatrixOp.
+// \see LoadStoreLikeOpRewriter.
+static Value getLdMatrixOpSrcMemRef(nvgpu::LdMatrixOp ldMatrixOp) {
+  return ldMatrixOp.getSrcMemref();
+}
+
+// Matches rebuildOpFromAddressAndIndices specs for LdMatrixOp.
+// \see LoadStoreLikeOpRewriter.
+static nvgpu::LdMatrixOp rebuildLdMatrixOp(RewriterBase &rewriter,
+                                           nvgpu::LdMatrixOp ldMatrixOp,
+                                           Value srcMemRef,
+                                           ArrayRef<Value> indices) {
+  Location loc = ldMatrixOp.getLoc();
+  return rewriter.create<nvgpu::LdMatrixOp>(
+      loc, ldMatrixOp.getResult().getType(), srcMemRef, indices,
+      ldMatrixOp.getTranspose(), ldMatrixOp.getNumTiles());
+}
+
+// Matches getViewSizeForEachDim specs for LdMatrixOp.
+// \see LoadStoreLikeOpRewriter.
+static SmallVector<OpFoldResult>
+getLdMatrixOpViewSizeForEachDim(RewriterBase &rewriter,
+                                nvgpu::LdMatrixOp ldMatrixOp) {
+  Location loc = ldMatrixOp.getLoc();
+  auto extractStridedMetadataOp =
+      rewriter.create<memref::ExtractStridedMetadataOp>(
+          loc, ldMatrixOp.getSrcMemref());
+  SmallVector<OpFoldResult> srcSizes =
+      extractStridedMetadataOp.getConstifiedMixedSizes();
+  SmallVector<OpFoldResult> indices =
+      getAsOpFoldResult(ldMatrixOp.getIndices());
+  SmallVector<OpFoldResult> finalSizes;
+
+  AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
+  AffineExpr s1 = rewriter.getAffineSymbolExpr(1);
+
+  for (auto [srcSize, indice] : llvm::zip(srcSizes, indices)) {
+    finalSizes.push_back(makeComposedFoldedAffineApply(rewriter, loc, s0 - s1,
+                                                       {srcSize, indice}));
+  }
+  return finalSizes;
+}
+
+/// Rewrite a store/load-like op so that all its indices are zeros.
+/// E.g., %ld = memref.load %base[%off0]...[%offN]
+/// =>
+/// %new_base = subview %base[%off0,.., %offN][1,..,1][1,..,1]
+/// %ld = memref.load %new_base[0,..,0] :
+///    memref<1x..x1xTy, strided<[1,..,1], offset: ?>>
+///
+/// `getSrcMemRef` returns the source memref for the given load-like operation.
+///
+/// `getViewSizeForEachDim` returns the sizes of view that is going to feed
+/// new operation. This must return one size per dimension of the view.
+/// The sizes of the view needs to be at least as big as what is actually
+/// going to be accessed. Use the provided `loadStoreOp` to get the right
+/// sizes.
+///
+/// Using the given rewriter, `rebuildOpFromAddressAndIndices` creates a new
+/// LoadStoreLikeOp that reads from srcMemRef[indices].
+/// The returned operation will be used to replace loadStoreOp.
+template <typename LoadStoreLikeOp, Value (*getSrcMemRef)(LoadStoreLikeOp),
+          LoadStoreLikeOp (*rebuildOpFromAddressAndIndices)(
+              RewriterBase & /*rewriter*/, LoadStoreLikeOp /*loadStoreOp*/,
+              Value /*srcMemRef*/, ArrayRef<Value> /*indices*/),
+          SmallVector<OpFoldResult> (*getViewSizeForEachDim)(
+              RewriterBase & /*rewriter*/, LoadStoreLikeOp /*loadStoreOp*/)>
+struct LoadStoreLikeOpRewriter : public OpRewritePattern<LoadStoreLikeOp> {
+  using OpRewritePattern<LoadStoreLikeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(LoadStoreLikeOp loadStoreLikeOp,
+                                PatternRewriter &rewriter) const override {
+    Value srcMemRef = getSrcMemRef(loadStoreLikeOp);
+    auto ldStTy = srcMemRef.getType().cast<MemRefType>();
+    unsigned loadStoreRank = ldStTy.getRank();
+    // Don't waste compile time if there is nothing to rewrite.
+    if (loadStoreRank == 0)
+      return failure();
+
+    // If our load already has only zeros as indices there is nothing
+    // to do.
+    SmallVector<OpFoldResult> indices =
+        getAsOpFoldResult(loadStoreLikeOp.getIndices());
+    if (std::all_of(indices.begin(), indices.end(),
+                    [](const OpFoldResult &opFold) {
+                      return isConstantIntValue(opFold, 0);
+                    })) {
+      return failure();
+    }
+
+    // Create the array of ones of the right size.
+    SmallVector<OpFoldResult> ones(loadStoreRank, rewriter.getIndexAttr(1));
+    SmallVector<OpFoldResult> sizes =
+        getViewSizeForEachDim(rewriter, loadStoreLikeOp);
+    assert(sizes.size() == loadStoreRank &&
+           "Expected one size per load dimension");
+    Location loc = loadStoreLikeOp.getLoc();
+    auto subview =
+        rewriter.create<memref::SubViewOp>(loc, /*source=*/srcMemRef,
+                                           /*offsets=*/indices,
+                                           /*sizes=*/sizes, /*strides=*/ones);
+    // Rewrite the load/store with the subview as the base pointer.
+    SmallVector<Value> zeros(loadStoreRank,
+                             rewriter.create<arith::ConstantIndexOp>(loc, 0));
+    LoadStoreLikeOp newLoadStore = rebuildOpFromAddressAndIndices(
+        rewriter, loadStoreLikeOp, subview.getResult(), zeros);
+    rewriter.replaceOp(loadStoreLikeOp, newLoadStore->getResults());
+    return success();
+  }
+};
+} // namespace
+
+void memref::populateExtractAddressComputationsPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<LoadStoreLikeOpRewriter<
+                   memref::LoadOp,
+                   /*getSrcMemRef=*/getLoadOpSrcMemRef,
+                   /*rebuildOpFromAddressAndIndices=*/rebuildLoadOp,
+                   /*getViewSizeForEachDim=*/getLoadOpViewSizeForEachDim>,
+               LoadStoreLikeOpRewriter<
+                   memref::StoreOp,
+                   /*getSrcMemRef=*/getStoreOpSrcMemRef,
+                   /*rebuildOpFromAddressAndIndices=*/rebuildStoreOp,
+                   /*getViewSizeForEachDim=*/getStoreOpViewSizeForEachDim>,
+               LoadStoreLikeOpRewriter<
+                   nvgpu::LdMatrixOp,
+                   /*getSrcMemRef=*/getLdMatrixOpSrcMemRef,
+                   /*rebuildOpFromAddressAndIndices=*/rebuildLdMatrixOp,
+                   /*getViewSizeForEachDim=*/getLdMatrixOpViewSizeForEachDim>>(
+      patterns.getContext());
+}

mlir/test/Dialect/MemRef/extract-address-computations.mlir

+// RUN: mlir-opt -test-extract-address-computations %s --split-input-file | FileCheck %s
+// Simple test: check that we extract the address computation of a load into
+// a dedicated subview.
+// The resulting load will be loading from the subview and have only indices
+// set to zero.
+
+// CHECK-LABEL: @test(
+// CHECK-SAME: %[[BASE:[^:]*]]: memref{{[^,]*}},
+// CHECK-SAME: %[[DYN_OFFSET:.*]]: index)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET]], 0, 8] [1, 1, 1] [1, 1, 1] : memref<2x16x16xf32> to memref<1x1x1xf32, strided<[256, 16, 1], offset: ?>>
+// CHECK: %[[LOADED_VAL:.*]] = memref.load %[[SUBVIEW]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32, strided<[256, 16, 1], offset: ?>>
+// CHECK: return %[[LOADED_VAL]] : f32
+func.func @test(%base : memref<2x16x16xf32>, %offset : index) -> f32 {
+  %c0 = arith.constant 0 : index
+  %c8 = arith.constant 8 : index
+  %loaded_val = memref.load %base[%offset, %c0, %c8] : memref<2x16x16xf32>
+  return %loaded_val : f32
+}
+
+// -----
+
+// Simple test: check that we extract the address computation of a store into
+// a dedicated subview.
+// The resulting store will use the address from the subview and have only
+// indices set to zero.
+
+// CHECK-LABEL: @test_store(
+// CHECK-SAME: %[[BASE:[^:]*]]: memref{{[^,]*}},
+// CHECK-SAME: %[[DYN_OFFSET:.*]]: index)
+// CHECK-DAG: %[[CF0:.*]] = arith.constant 0.0{{0*e\+00}} : f32
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET]], 0, 8] [1, 1, 1] [1, 1, 1] : memref<2x16x16xf32> to memref<1x1x1xf32, strided<[256, 16, 1], offset: ?>>
+// CHECK: memref.store %[[CF0]], %[[SUBVIEW]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32, strided<[256, 16, 1], offset: ?>>
+// CHECK: return
+func.func @test_store(%base : memref<2x16x16xf32>, %offset : index) -> () {
+  %cf0 = arith.constant 0.0 : f32
+  %c0 = arith.constant 0 : index
+  %c8 = arith.constant 8 : index
+  memref.store %cf0, %base[%offset, %c0, %c8] : memref<2x16x16xf32>
+  return
+}
+
+// -----
+
+// For this test, we made the source memref fully dynamic.
+// The gist of the check remains the same as the simple test:
+// The address computation is extracted into its own subview.
+// CHECK-LABEL: @testWithLoop(
+// CHECK-SAME: %[[BASE:[^:]*]]: memref
+// CHECK:  %[[SUM_ALL:.*]] = arith.constant 0.0{{0*e\+00}} : f32
+// CHECK:  %[[C0:.*]] = arith.constant 0 : index
+// CHECK:  %[[C1:.*]] = arith.constant 1 : index
+// CHECK:  %[[C2:.*]] = arith.constant 2 : index
+// CHECK:  %[[UPPER_BOUND0:.*]] = memref.dim %[[BASE]], %[[C0]] : memref<?x?x?xf32,
+// CHECK:  %[[UPPER_BOUND1:.*]] = memref.dim %[[BASE]], %[[C1]] : memref<?x?x?xf32,
+// CHECK:  %[[UPPER_BOUND2:.*]] = memref.dim %[[BASE]], %[[C2]] : memref<?x?x?xf32,
+// CHECK:  %[[SUM_RES2:.*]] = scf.for %[[IV2:.*]] = %[[C0]] to %[[UPPER_BOUND2]] step %[[C1]] iter_args(%[[SUM_ITER2:.*]] = %[[SUM_ALL]]) -> (f32) {
+// CHECK:    %[[SUM_RES1:.*]] = scf.for %[[IV1:.*]] = %[[C0]] to %[[UPPER_BOUND1]] step %[[C1]] iter_args(%[[SUM_ITER1:.*]] = %[[SUM_ITER2]]) -> (f32) {
+// CHECK:      %[[SUM_RES0:.*]] = scf.for %[[IV0:.*]] = %[[C0]] to %[[UPPER_BOUND0]] step %[[C1]] iter_args(%[[SUM_ITER0:.*]] = %[[SUM_ITER1]]) -> (f32) {
+// CHECK:        %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[IV0]], %[[IV1]], %[[IV2]]] [1, 1, 1] [1, 1, 1] : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to memref<1x1x1xf32, strided<[?, ?, ?], offset: ?>>
+// CHECK:        %[[LOADED_VAL:.*]] = memref.load %[[SUBVIEW]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32, strided<[?, ?, ?], offset: ?>>
+// CHECK:        %[[RES:.*]] = arith.addf %[[LOADED_VAL]], %[[SUM_ITER2]] : f32
+// CHECK:        scf.yield %[[RES]] : f32
+// CHECK:      }
+// CHECK:      scf.yield %[[SUM_RES0]] : f32
+// CHECK:    }
+// CHECK:    scf.yield %[[SUM_RES1]] : f32
+// CHECK:  }
+// CHECK:  return %[[SUM_RES2]] : f32
+func.func @testWithLoop(%base : memref<?x?x?xf32, strided<[?,?,?], offset: ?>>) -> f32 {
+  %sum_all = arith.constant 0.0 : f32
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %upper_bound0 = memref.dim %base, %c0 : memref<?x?x?xf32, strided<[?,?,?], offset: ?>>
+  %upper_bound1 = memref.dim %base, %c1 : memref<?x?x?xf32, strided<[?,?,?], offset: ?>>
+  %upper_bound2 = memref.dim %base, %c2 : memref<?x?x?xf32, strided<[?,?,?], offset: ?>>
+  %sum_res2 = scf.for %iv2 = %c0 to %upper_bound2 step %c1 iter_args(%sum_iter2 = %sum_all) -> (f32) {
+    %sum_res1 = scf.for %iv1 = %c0 to %upper_bound1 step %c1 iter_args(%sum_iter1 = %sum_iter2) -> (f32) {
+      %sum_res0 = scf.for %iv0 = %c0 to %upper_bound0 step %c1 iter_args(%sum_iter0 = %sum_iter1) -> (f32) {
+        %loaded_val = memref.load %base[%iv0, %iv1, %iv2] : memref<?x?x?xf32, strided<[?,?,?], offset: ?>>
+        %res = arith.addf %loaded_val, %sum_iter2 : f32
+        scf.yield %res : f32
+      }
+      scf.yield %sum_res0 : f32
+    }
+    scf.yield %sum_res1 : f32
+  }
+  return %sum_res2 : f32
+}
+
+// -----
+
+// Simple test: check that we extract the address computation of a ldmatrix into
+// a dedicated subview.
+// The resulting ldmatrix will loaded from with subview and have only indices set
+// to zero.
+// Also the sizes of the view are adjusted to `original size - offset`.
+
+// CHECK-DAG: #[[$FOUR_MINUS_OFF_MAP:.*]] = affine_map<()[s0] -> (-s0 + 4)>
+// CHECK-DAG: #[[$THIRTY_TWO_MINUS_OFF_MAP:.*]] = affine_map<()[s0] -> (-s0 + 32)>
+// CHECK-LABEL: @test_ldmatrix(
+// CHECK-SAME: %[[BASE:[^:]*]]: memref<{{[^,]*}}, 3>,
+// CHECK-SAME: %[[DYN_OFFSET0:[^:]*]]: index,
+// CHECK-SAME: %[[DYN_OFFSET1:[^:]*]]: index,
+// CHECK-SAME: %[[DYN_OFFSET2:[^:]*]]: index)
+// CHECK-DAG: %[[DYN_SIZE0:.*]] = affine.apply #[[$FOUR_MINUS_OFF_MAP]]()[%[[DYN_OFFSET0]]]
+// CHECK-DAG: %[[DYN_SIZE1:.*]] = affine.apply #[[$THIRTY_TWO_MINUS_OFF_MAP]]()[%[[DYN_OFFSET1]]]
+// CHECK-DAG: %[[DYN_SIZE2:.*]] = affine.apply #[[$THIRTY_TWO_MINUS_OFF_MAP]]()[%[[DYN_OFFSET2]]]
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET0]], %[[DYN_OFFSET1]], %[[DYN_OFFSET2]]] [%[[DYN_SIZE0]], %[[DYN_SIZE1]], %[[DYN_SIZE2]]] [1, 1, 1] : memref<4x32x32xf16, 3> to memref<?x?x?xf16, strided<[1024, 32, 1], offset: ?>, 3>
+// CHECK: %[[LOADED_VAL:.*]] = nvgpu.ldmatrix %[[SUBVIEW]][%[[C0]], %[[C0]], %[[C0]]] {numTiles = 4 : i32, transpose = false} : memref<?x?x?xf16, strided<[1024, 32, 1], offset: ?>, 3> -> vector<4x2xf16>
+// CHECK: return %[[LOADED_VAL]] : vector<4x2xf16>
+func.func @test_ldmatrix(%base : memref<4x32x32xf16, 3>,
+    %offset0 : index, %offset1: index, %offset2: index)
+    -> vector<4x2xf16> {
+  %loaded_val = nvgpu.ldmatrix
+    %base[%offset0, %offset1, %offset2]
+    {numTiles = 4 : i32, transpose = false}
+      : memref<4x32x32xf16, 3> -> vector<4x2xf16>
+  return %loaded_val : vector<4x2xf16>
+}
+
+// -----
+
+// Same as test_ldmatrix but with fully dynamic memref.
+
+// CHECK-DAG: #[[$A_MINUS_B_MAP:.*]] = affine_map<()[s0, s1] -> (s0 - s1)>
+// CHECK-LABEL: @test_ldmatrix(
+// CHECK-SAME: %[[BASE:[^:]*]]: memref<{{[^,]*}}, 3>,
+// CHECK-SAME: %[[DYN_OFFSET0:[^:]*]]: index,
+// CHECK-SAME: %[[DYN_OFFSET1:[^:]*]]: index,
+// CHECK-SAME: %[[DYN_OFFSET2:[^:]*]]: index)
+// CHECK-DAG: {{.*}}, {{.*}}, %[[DYN_SIZES:.*]]:3, {{.*}} = memref.extract_strided_metadata %[[BASE]]
+// CHECK-DAG: %[[DYN_SIZE0:.*]] = affine.apply #[[$A_MINUS_B_MAP]]()[%[[DYN_SIZES]]#0, %[[DYN_OFFSET0]]]
+// CHECK-DAG: %[[DYN_SIZE1:.*]] = affine.apply #[[$A_MINUS_B_MAP]]()[%[[DYN_SIZES]]#1, %[[DYN_OFFSET1]]]
+// CHECK-DAG: %[[DYN_SIZE2:.*]] = affine.apply #[[$A_MINUS_B_MAP]]()[%[[DYN_SIZES]]#2, %[[DYN_OFFSET2]]]
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview %[[BASE]][%[[DYN_OFFSET0]], %[[DYN_OFFSET1]], %[[DYN_OFFSET2]]] [%[[DYN_SIZE0]], %[[DYN_SIZE1]], %[[DYN_SIZE2]]] [1, 1, 1] : memref<?x?x?xf16, 3> to memref<?x?x?xf16, strided<[?, ?, 1], offset: ?>, 3>
+// CHECK: %[[LOADED_VAL:.*]] = nvgpu.ldmatrix %[[SUBVIEW]][%[[C0]], %[[C0]], %[[C0]]] {numTiles = 4 : i32, transpose = false} : memref<?x?x?xf16, strided<[?, ?, 1], offset: ?>, 3> -> vector<4x2xf16>
+// CHECK: return %[[LOADED_VAL]] : vector<4x2xf16>
+func.func @test_ldmatrix(%base : memref<?x?x?xf16, 3>,
+    %offset0 : index, %offset1: index, %offset2: index)
+    -> vector<4x2xf16> {
+  %loaded_val = nvgpu.ldmatrix
+    %base[%offset0, %offset1, %offset2]
+    {numTiles = 4 : i32, transpose = false}
+      : memref<?x?x?xf16, 3> -> vector<4x2xf16>
+  return %loaded_val : vector<4x2xf16>
+}
+

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

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

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

+//===- TestExtractAddressComputations.cpp - Test extract addr computations-===//
+//
+// 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 a pass to test extract address computations patterns.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#define PASS_NAME "test-extract-address-computations"
+
+using namespace mlir;
+
+namespace {
+
+struct TestExtractAddressComputations
+    : public PassWrapper<TestExtractAddressComputations,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestExtractAddressComputations)
+
+  StringRef getArgument() const final { return PASS_NAME; }
+  StringRef getDescription() const final {
+    return "Tests extract address computations patterns.";
+  }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry
+        .insert<AffineDialect, memref::MemRefDialect, nvgpu::NVGPUDialect>();
+  }
+  TestExtractAddressComputations() = default;
+  TestExtractAddressComputations(const TestExtractAddressComputations &pass)
+      : PassWrapper(pass){};
+
+  void runOnOperation() override;
+};
+
+} // namespace
+
+void TestExtractAddressComputations::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  memref::populateExtractAddressComputationsPatterns(patterns);
+  if (failed(
+          applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) {
+    return signalPassFailure();
+  }
+}
+
+namespace mlir {
+namespace test {
+void registerTestExtractAddressComputations() {
+  PassRegistration<TestExtractAddressComputations>();
+}
+} // namespace test
+} // namespace mlir

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

 void registerTestDataLayoutQuery();
 void registerTestDeadCodeAnalysisPass();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDialectConversionPasses();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandMathPass();
+void registerTestExtractAddressComputations();
 void registerTestFooAnalysisPass();
 void registerTestComposeSubView();
 void registerTestMultiBuffering();
 void registerTestIntRangeInference();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();

 #endif
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDataLayoutPropagation();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDeadCodeAnalysisPass();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestExpandMathPass();
+  mlir::test::registerTestExtractAddressComputations();
   mlir::test::registerTestFooAnalysisPass();
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
   mlir::test::registerTestIntRangeInference();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLastModifiedPass();

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

deps = [
         ":DialectUtils",
         ":FuncDialect",
         ":GPUDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":MemRefPassIncGen",
+        ":NVGPUDialect",
         ":Pass",
         ":RuntimeVerifiableOpInterface",
         ":TensorDialect",
         ":Transforms",
         ":VectorDialect",
         "//llvm:Support",
     ],
 )