[mlir][MemRef] Expose utility to invert subview index mapping

[christopherbate]

Summary

The FoldSubView transformation contains a utility that takes a set of
indices in the domain of a memref.subview result and transform them
into the domain of the memref.subview's source. This function is
useful more widely, especially when "subview folding" needs to be
implemented in other dialects. It is therefore moved under
MemRef/Utils.h and renamed to a memref::invertSubViewIndexMapping.

[nicolasvasilache]

@christopherbate note https://reviews.llvm.org/D128986 which also supports expand / collapse.
I suspect we can refactor at finer granularity and make the individual parts more composable (i.e. make OpFoldResult and AffineApply happen under the hood rather than have to maintain all this if/else everywhere).

mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h

39

invert is a misnomer IMO, this is better thought of as folding through a subview.

[christopherbate]

Closing in favor of introducing interface. Here is the draft diff: https://reviews.llvm.org/D134393

mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h

 // either by passes, optimization sequences, or in turn by other transformation
 // utilities.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H
 #define MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H
 
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace mlir {
+namespace memref {
+
+/// Given the 'indices' indexing into a memref that is the result
+/// of a subview op, returns the indices w.r.t to the source memref of the
+/// subview op. For example
+///
+/// %0 = ... : memref<12x42xf32>
+/// %1 = memref.subview %0[%arg0, %arg1][4, 4][%stride0, %stride1]
+///         memref<12x42xf32> to memref<4x4xf32, #map0>
+/// %2 = memref.load %1[%i0, %i1] : memref<4x4xf32, #map0>
+///
+/// The reference to %1[%i0, %i1] can be replaced with a reference to
+/// %0 if we invert the indexing transform represented by the subview:
+///
+/// memref.load %0[%arg0 + %i0 * %stride0][%arg1 + %i1 * %stride1] :
+///          memref<12x42xf32>
+LogicalResult invertSubViewIndexMapping(Location loc, PatternRewriter &rewriter,
+                                        memref::SubViewOp subViewOp,
+                                        ValueRange indices,
+                                        SmallVectorImpl<Value> &sourceIndices);
+
+} // namespace memref
+} // namespace mlir
+
 #endif // MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H

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

add_mlir_dialect_library(MLIRMemRefTransforms
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineUtils
   MLIRArithmeticDialect
   MLIRFuncDialect
   MLIRInferTypeOpInterface
   MLIRLoopLikeInterface
   MLIRMemRefDialect
+  MLIRMemRefUtils
   MLIRPass
   MLIRTensorDialect
   MLIRTransforms
   MLIRVectorDialect
 )
 

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Utility functions
 //===----------------------------------------------------------------------===//
 
-/// Given the 'indices' of an load/store operation where the memref is a result
-/// of a subview op, returns the indices w.r.t to the source memref of the
-/// subview op. For example
-///
-/// %0 = ... : memref<12x42xf32>
-/// %1 = subview %0[%arg0, %arg1][][%stride1, %stride2] : memref<12x42xf32> to
-///          memref<4x4xf32, offset=?, strides=[?, ?]>
-/// %2 = load %1[%i1, %i2] : memref<4x4xf32, offset=?, strides=[?, ?]>
-///
-/// could be folded into
-///
-/// %2 = load %0[%arg0 + %i1 * %stride1][%arg1 + %i2 * %stride2] :
-///          memref<12x42xf32>
-static LogicalResult
-resolveSourceIndices(Location loc, PatternRewriter &rewriter,
-                     memref::SubViewOp subViewOp, ValueRange indices,
-                     SmallVectorImpl<Value> &sourceIndices) {
-  SmallVector<OpFoldResult> mixedOffsets = subViewOp.getMixedOffsets();
-  SmallVector<OpFoldResult> mixedSizes = subViewOp.getMixedSizes();
-  SmallVector<OpFoldResult> mixedStrides = subViewOp.getMixedStrides();
-
-  SmallVector<Value> useIndices;
-  // Check if this is rank-reducing case. Then for every unit-dim size add a
-  // zero to the indices.
-  unsigned resultDim = 0;
-  llvm::SmallBitVector unusedDims = subViewOp.getDroppedDims();
-  for (auto dim : llvm::seq<unsigned>(0, subViewOp.getSourceType().getRank())) {
-    if (unusedDims.test(dim))
-      useIndices.push_back(rewriter.create<arith::ConstantIndexOp>(loc, 0));
-    else
-      useIndices.push_back(indices[resultDim++]);
-  }
-  if (useIndices.size() != mixedOffsets.size())
-    return failure();
-  sourceIndices.resize(useIndices.size());
-  for (auto index : llvm::seq<size_t>(0, mixedOffsets.size())) {
-    SmallVector<Value> dynamicOperands;
-    AffineExpr expr = rewriter.getAffineDimExpr(0);
-    unsigned numSymbols = 0;
-    dynamicOperands.push_back(useIndices[index]);
-
-    // Multiply the stride;
-    if (auto attr = mixedStrides[index].dyn_cast<Attribute>()) {
-      expr = expr * attr.cast<IntegerAttr>().getInt();
-    } else {
-      dynamicOperands.push_back(mixedStrides[index].get<Value>());
-      expr = expr * rewriter.getAffineSymbolExpr(numSymbols++);
-    }
-
-    // Add the offset.
-    if (auto attr = mixedOffsets[index].dyn_cast<Attribute>()) {
-      expr = expr + attr.cast<IntegerAttr>().getInt();
-    } else {
-      dynamicOperands.push_back(mixedOffsets[index].get<Value>());
-      expr = expr + rewriter.getAffineSymbolExpr(numSymbols++);
-    }
-    Location loc = subViewOp.getLoc();
-    sourceIndices[index] = rewriter.create<AffineApplyOp>(
-        loc, AffineMap::get(1, numSymbols, expr), dynamicOperands);
-  }
-  return success();
-}
-
 /// Helpers to access the memref operand for each op.
 template <typename LoadOrStoreOpTy>
 static Value getMemRefOperand(LoadOrStoreOpTy op) {
   return op.getMemref();
 }
 
 static Value getMemRefOperand(vector::TransferReadOp op) {
   return op.getSource();

 LoadOpOfSubViewFolder<OpTy>::matchAndRewrite(OpTy loadOp,
                                              PatternRewriter &rewriter) const {
   auto subViewOp =
       getMemRefOperand(loadOp).template getDefiningOp<memref::SubViewOp>();
   if (!subViewOp)
     return failure();
 
   SmallVector<Value, 4> sourceIndices;
-  if (failed(resolveSourceIndices(loadOp.getLoc(), rewriter, subViewOp,
+  if (failed(invertSubViewIndexMapping(loadOp.getLoc(), rewriter, subViewOp,
                                        loadOp.getIndices(), sourceIndices)))
     return failure();
 
   replaceOp(loadOp, subViewOp, sourceIndices, rewriter);
   return success();
 }
 
 template <typename OpTy>
 LogicalResult
 StoreOpOfSubViewFolder<OpTy>::matchAndRewrite(OpTy storeOp,
                                               PatternRewriter &rewriter) const {
   auto subViewOp =
       getMemRefOperand(storeOp).template getDefiningOp<memref::SubViewOp>();
   if (!subViewOp)
     return failure();
 
   SmallVector<Value, 4> sourceIndices;
-  if (failed(resolveSourceIndices(storeOp.getLoc(), rewriter, subViewOp,
+  if (failed(invertSubViewIndexMapping(storeOp.getLoc(), rewriter, subViewOp,
                                        storeOp.getIndices(), sourceIndices)))
     return failure();
 
   replaceOp(storeOp, subViewOp, sourceIndices, rewriter);
   return success();
 }
 
 void memref::populateFoldSubViewOpPatterns(RewritePatternSet &patterns) {
   patterns.add<LoadOpOfSubViewFolder<AffineLoadOp>,

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

 add_mlir_dialect_library(MLIRMemRefUtils
   MemRefUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/inlude/mlir/Dialect/MemRefDialect
+
+  DEPENDS
+  MLIRAffineDialect
+  MLIRArithmeticDialect
+  MLIRMemRefDialect
 )
 

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

 //===- MemRefUtils.cpp - Utilities to support the MemRef dialect ----------===//
 //
 // 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 utilities for the MemRef dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+
+using namespace mlir;
+using namespace mlir::memref;
+
+LogicalResult
+memref::invertSubViewIndexMapping(Location loc, PatternRewriter &rewriter,
+                                  SubViewOp subViewOp, ValueRange indices,
+                                  SmallVectorImpl<Value> &sourceIndices) {
+  SmallVector<OpFoldResult> mixedOffsets = subViewOp.getMixedOffsets();
+  SmallVector<OpFoldResult> mixedSizes = subViewOp.getMixedSizes();
+  SmallVector<OpFoldResult> mixedStrides = subViewOp.getMixedStrides();
+
+  SmallVector<Value> useIndices;
+  // Check if this is rank-reducing case. Then for every unit-dim size add a
+  // zero to the indices.
+  unsigned resultDim = 0;
+  llvm::SmallBitVector unusedDims = subViewOp.getDroppedDims();
+  for (auto dim : llvm::seq<unsigned>(0, subViewOp.getSourceType().getRank())) {
+    if (unusedDims.test(dim))
+      useIndices.push_back(rewriter.create<arith::ConstantIndexOp>(loc, 0));
+    else
+      useIndices.push_back(indices[resultDim++]);
+  }
+  if (useIndices.size() != mixedOffsets.size())
+    return failure();
+  sourceIndices.resize(useIndices.size());
+  for (auto index : llvm::seq<size_t>(0, mixedOffsets.size())) {
+    SmallVector<Value> dynamicOperands;
+    AffineExpr expr = rewriter.getAffineDimExpr(0);
+    unsigned numSymbols = 0;
+    dynamicOperands.push_back(useIndices[index]);
+
+    // Multiply the stride;
+    if (auto attr = mixedStrides[index].dyn_cast<Attribute>()) {
+      expr = expr * attr.cast<IntegerAttr>().getInt();
+    } else {
+      dynamicOperands.push_back(mixedStrides[index].get<Value>());
+      expr = expr * rewriter.getAffineSymbolExpr(numSymbols++);
+    }
+
+    // Add the offset.
+    if (auto attr = mixedOffsets[index].dyn_cast<Attribute>()) {
+      expr = expr + attr.cast<IntegerAttr>().getInt();
+    } else {
+      dynamicOperands.push_back(mixedOffsets[index].get<Value>());
+      expr = expr + rewriter.getAffineSymbolExpr(numSymbols++);
+    }
+    Location loc = subViewOp.getLoc();
+    sourceIndices[index] = rewriter.create<AffineApplyOp>(
+        loc, AffineMap::get(1, numSymbols, expr), dynamicOperands);
+  }
+  return success();
+}