Add a pattern to combine composed subview ops

[agrue]

Fix relevant CMakeLists.txt

[mehdi_amini]

mlir/lib/Dialect/StandardOps/Transforms/ComposeSubView.cpp

76

FYI LLVM is still using C++14, this won't build with all of our supported platforms unfortunately (it is neat though!)

[agrue]

Fix clang-tidy warnings

[mravishankar]

mlir/lib/Dialect/StandardOps/Transforms/ComposeSubView.cpp

69

Not sure about the math for strides != 1. It can have weird corner cases (I am not saying it doesnt work, or I have a specific flaw I see in the code here, but just an intuition). Unless really needed, maybe restrict to strides = 1 case in source and dest subviews?

117

If one of the offsets is a constant then you can construct the affine map with the constant already. So do something like this

AffineExpr expr;
SmallVector<Value> operands;
auto addValueOrAttr = [&](OpFoldResult valueOrAttr) {
  if (auto attr = valueOrAttr.dyn_cast<Attribute>()) {
    expr = expr + attr.cast<IntegerAttr>().getInt();
  } else {
    expr = expr + getAffineSymbolExpr(operands.size());
    operands.push_back(valueOrAttr.get<Value>());
  }
};
AffineMap map = AffineMap::get(0, operands.size(), expr);

[agrue]

Address review comments

[agrue]

mlir/lib/Dialect/StandardOps/Transforms/ComposeSubView.cpp

69

Makes sense. We only need strides == 1. Done.

117

I see, thanks. Done.

[mravishankar]

Thanks. Just a few minor nits, this time around.

mlir/lib/Dialect/StandardOps/Transforms/ComposeSubView.cpp

69

Just to make it cleaner. Just add the strides == 1 check outside of the loop. You dont need to iterate over it here. Makes it simpler to read. Also you can just create the strides for the result later on.

85

All the stride stuff is better checked before. You just need this check

if (llvm::any_of(strides, [](OpFoldResult &valueOrAttr) {
  Attribute attr = valueOrAttr.dyn_cast<Attribute>();
  return attr && attr.cast<IntegerAttr>().getInt() == 1;
}

So you ca

98

Nit: I think you can drop the if here and get result below using

Value result = affineMapToValues(rewriter, loc, AffineMap::get(0, affineApplyOperands.size(), expr), affineApplyOperands)[0];

https://github.com/llvm/llvm-project/blob/fe5c4a06a49c697399323cc55e2135a5452ed616/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h#L28

More readable (to me).

[mravishankar]

LGTM-ing this based on the strides being checked ahead of time (instead of in the loop)

mlir/lib/Dialect/StandardOps/Transforms/ComposeSubView.cpp

98

Dropping this for now.

[agrue]

Address further review comments.

[agrue]

Thanks. Can someone else please land this patch when it's ready? I don't have permissions.

[vinograd47]

Why it is located in MLIRStandardOpsTransforms, while operating in MemRef dialect? I would expect it to be a part of MLIRMemRefTransforms.

mlir/include/mlir/Dialect/StandardOps/Transforms/ComposeSubView.h

+//===- ComposeSubView.h - Combining composed subview 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Patterns for combining composed subview ops.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_STANDARDOPS_TRANSFORMS_COMPOSESUBVIEW_H_
+#define MLIR_DIALECT_STANDARDOPS_TRANSFORMS_COMPOSESUBVIEW_H_
+
+namespace mlir {
+
+// Forward declarations.
+class MLIRContext;
+class RewritePatternSet;
+using OwningRewritePatternList = RewritePatternSet;
+
+void populateComposeSubViewPatterns(OwningRewritePatternList &patterns,
+                                    MLIRContext *context);
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_STANDARDOPS_TRANSFORMS_COMPOSESUBVIEW_H_

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

 add_mlir_dialect_library(MLIRStandardOpsTransforms
   Bufferize.cpp
+  ComposeSubView.cpp
   DecomposeCallGraphTypes.cpp
   ExpandOps.cpp
   FuncBufferize.cpp
   FuncConversions.cpp
   TensorConstantBufferize.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps/Transforms

mlir/lib/Dialect/StandardOps/Transforms/ComposeSubView.cpp

+//===- ComposeSubView.cpp - Combining composed subview 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains patterns for combining composed subview ops (i.e. subview
+// of a subview becomes a single subview).
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/StandardOps/Transforms/ComposeSubView.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+
+namespace {
+
+// Replaces a subview of a subview with a single subview. Only supports subview
+// ops with static sizes and static strides (both static and dynamic offsets
+// are supported).
+struct ComposeSubViewOpPattern : public OpRewritePattern<memref::SubViewOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::SubViewOp op,
+                                PatternRewriter &rewriter) const override {
+    // 'op' is the 'SubViewOp' we're rewriting. 'source_op' is the op that
+    // produces the input of the op we're rewriting (for 'SubViewOp' the input
+    // is called the "source" value). We can only combine them if both 'op' and
+    // 'source_op' are 'SubViewOp'.
+    auto source_op = op.source().getDefiningOp<memref::SubViewOp>();
+    if (!source_op)
+      return failure();
+
+    // A 'SubViewOp' can be "rank-reducing" by eliminating dimensions of the
+    // output memref that are statically known to be equal to 1. We do not
+    // allow 'source_op' to be a rank-reducing subview because then our two
+    // 'SubViewOp's would have different numbers of offset/size/stride
+    // parameters (just difficult to deal with, not impossible if we end up
+    // needing it).
+    if (source_op.getSourceType().getRank() != source_op.getType().getRank()) {
+      return failure();
+    }
+
+    // Offsets, sizes and strides ValueOrAttr for the combined 'SubViewOp'.
+    SmallVector<OpFoldResult> offsets, sizes, strides;
+    auto op_offset_iter = op.offsets().begin();
+    auto source_offset_iter = source_op.offsets().begin();
+
+    // An affine map that sums two dimension parameters.
+    auto *ctx = rewriter.getContext();
+    auto binary_add_map = AffineMap::get(
+        /*dimCount=*/2, /*symbolCount=*/0,
+        getAffineDimExpr(0, ctx) + getAffineDimExpr(1, ctx));
+
+    // The rules for calculating the new offsets, sizes and strides are:
+    // * Multiple subview offsets for a given dimension compose additively.
+    //   ("Offset by m" followed by "Offset by n" == "Offset by m + n")
+    // * Multiple strides for a given dimension compose multiplicatively.
+    //   ("Stride by m" followed by "Stride by n" == "Stride by m * n")
+    // * Multiple sizes for a given dimension compose by taking the size of the
+    //   final subview and ignoring the rest. ("Take m values" followed by "Take
+    //   n values" == "Take n values") This size must also be the smallest one
+    //   by definition (a subview needs to be the same size as or smaller than
+    //   its source along each dimension; presumably subviews that are larger
+    //   than their sources are disallowed by validation).
+    for (auto [op_offset, source_offset, op_size, op_stride, source_stride] :
+         llvm::zip(op.getMixedOffsets(), source_op.getMixedOffsets(),
+                   op.getMixedSizes(), op.getMixedStrides(),
+                   source_op.getMixedStrides())) {
+      // We only support static sizes and static strides.
+      if (op_size.is<Value>() || op_stride.is<Value>() ||
+          source_stride.is<Value>()) {
+        return failure();
+      }
+
+      sizes.push_back(op_size);
+      strides.push_back(rewriter.getI64IntegerAttr(
+          op_stride.get<Attribute>().cast<IntegerAttr>().getInt() *
+          source_stride.get<Attribute>().cast<IntegerAttr>().getInt()));
+
+      Value opOffset = op_offset.dyn_cast<Value>(),
+            sourceOffset = source_offset.dyn_cast<Value>();
+      // If both offsets are static we can simply calculate the combined
+      // offset statically.
+      if (!opOffset && !sourceOffset) {
+        offsets.push_back(rewriter.getI64IntegerAttr(
+            op_offset.get<Attribute>().cast<IntegerAttr>().getInt() +
+            source_offset.get<Attribute>().cast<IntegerAttr>().getInt()));
+        continue;
+      }
+
+      // If either offset attribute is dynamic, we must emit an additional
+      // affine transformation to add the two offsets together dynamically. Any
+      // static offsets must also be materialized as runtime constants.
+      Attribute opOffsetAttr = op_offset.dyn_cast<Attribute>(),
+                sourceOffsetAttr = source_offset.dyn_cast<Attribute>();
+      Value op_offset_for_sum =
+          opOffsetAttr
+              ? rewriter.create<ConstantIndexOp>(
+                    op.getLoc(), opOffsetAttr.cast<IntegerAttr>().getInt())
+              : *op_offset_iter++;
+      Value source_offset_for_sum =
+          sourceOffsetAttr
+              ? rewriter.create<ConstantIndexOp>(
+                    op.getLoc(), sourceOffsetAttr.cast<IntegerAttr>().getInt())
+              : *source_offset_iter++;
+      Value result = rewriter.create<AffineApplyOp>(
+          op.getLoc(), binary_add_map,
+          ValueRange{op_offset_for_sum, source_offset_for_sum});
+      offsets.push_back(result);
+    }
+
+    // This replaces 'op' but leaves 'source_op' alone; if it no longer has any
+    // uses it can be removed by a (separate) dead code elimination pass.
+    rewriter.replaceOpWithNewOp<memref::SubViewOp>(op, source_op.source(),
+                                                   offsets, sizes, strides);
+    return success();
+  }
+};
+
+} // namespace
+
+void populateComposeSubViewPatterns(OwningRewritePatternList &patterns,
+                                    MLIRContext *context) {
+  patterns.insert<ComposeSubViewOpPattern>(context);
+}
+
+} // namespace mlir

mlir/test/Transforms/compose-subview.mlir

+// RUN: mlir-opt -allow-unregistered-dialect %s -test-compose-subview -split-input-file | FileCheck %s
+
+// CHECK: [[MAP:#.*]] = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 3456)
+#map0 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 2304)>
+#map1 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 3456)>
+
+func @main(%input: memref<4x1024xf32>) -> memref<1x128xf32, #map1> {
+  //      CHECK: subview %arg0[3, 384] [1, 128] [1, 1]
+  // CHECK-SAME: memref<4x1024xf32> to memref<1x128xf32, [[MAP]]>
+  %0 = memref.subview %input[2, 256] [2, 256] [1, 1] : memref<4x1024xf32> to memref<2x256xf32, #map0>
+  %1 = memref.subview %0[1, 128] [1, 128] [1, 1] : memref<2x256xf32, #map0> to memref<1x128xf32, #map1>
+  return %1 : memref<1x128xf32, #map1>
+}
+
+// -----
+
+// CHECK: [[MAP:#.*]] = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 3745)
+#map0 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 1536)>
+#map1 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 2688)>
+#map2 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 3745)>
+
+func @main(%input: memref<4x1024xf32>) -> memref<1x10xf32, #map2> {
+  //      CHECK: subview %arg0[3, 673] [1, 10] [1, 1]
+  // CHECK-SAME: memref<4x1024xf32> to memref<1x10xf32, [[MAP]]>
+  %0 = memref.subview %input[1, 512] [3, 256] [1, 1] : memref<4x1024xf32> to memref<3x256xf32, #map0>
+  %1 = memref.subview %0[1, 128] [2, 128] [1, 1] : memref<3x256xf32, #map0> to memref<2x128xf32, #map1>
+  %2 = memref.subview %1[1, 33] [1, 10] [1, 1] : memref<2x128xf32, #map1> to memref<1x10xf32, #map2>
+  return %2 : memref<1x10xf32, #map2>
+}
+
+// -----
+
+// CHECK: [[MAP:#.*]] = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)
+#map = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>
+
+func @main(%input: memref<4x1024xf32>) -> memref<1x128xf32, #map> {
+  // CHECK: [[CST_3:%.*]] = constant 3 : index
+  %cst_1 = constant 1 : index
+  %cst_2 = constant 2 : index
+  //      CHECK: subview %arg0{{\[}}[[CST_3]], 384] [1, 128] [1, 1]
+  // CHECK-SAME: memref<4x1024xf32> to memref<1x128xf32, [[MAP]]>
+  %0 = memref.subview %input[%cst_2, 256] [2, 256] [1, 1] : memref<4x1024xf32> to memref<2x256xf32, #map>
+  %1 = memref.subview %0[%cst_1, 128] [1, 128] [1, 1] : memref<2x256xf32, #map> to memref<1x128xf32, #map>
+  return %1 : memref<1x128xf32, #map>
+}
+
+// -----
+
+// CHECK: [[MAP:#.*]] = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)
+#map = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>
+
+func @main(%input: memref<4x1024xf32>) -> memref<1x128xf32, #map> {
+  // CHECK: [[CST_3:%.*]] = constant 3 : index
+  %cst_2 = constant 2 : index
+  // CHECK: [[CST_384:%.*]] = constant 384 : index
+  %cst_128 = constant 128 : index
+  //      CHECK: subview %arg0{{\[}}[[CST_3]], [[CST_384]]] [1, 128] [1, 1]
+  // CHECK-SAME: memref<4x1024xf32> to memref<1x128xf32, [[MAP]]>
+  %0 = memref.subview %input[%cst_2, 256] [2, 256] [1, 1] : memref<4x1024xf32> to memref<2x256xf32, #map>
+  %1 = memref.subview %0[1, %cst_128] [1, 128] [1, 1] : memref<2x256xf32, #map> to memref<1x128xf32, #map>
+  return %1 : memref<1x128xf32, #map>
+}

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestAffineLoopParametricTiling.cpp
   TestExpandTanh.cpp
   TestCallGraph.cpp
   TestDecomposeCallGraphTypes.cpp
   TestConstantFold.cpp
   TestConvVectorization.cpp
   TestConvertCallOp.cpp
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp
   TestDataLayoutQuery.cpp
   TestDominance.cpp
   TestDynamicPipeline.cpp
+  TestComposeSubView.cpp
   TestLoopFusion.cpp
   TestGpuMemoryPromotion.cpp
   TestGpuParallelLoopMapping.cpp
   TestGpuRewrite.cpp
   TestInlining.cpp
   TestLinalgCodegenStrategy.cpp
   TestLinalgFusionTransforms.cpp
   TestLinalgHoisting.cpp

mlir/test/lib/Transforms/TestComposeSubView.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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to test the composed subview patterns.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/StandardOps/Transforms/ComposeSubView.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+struct TestComposeSubViewPass
+    : public PassWrapper<TestComposeSubViewPass, FunctionPass> {
+  void runOnFunction() override;
+  void getDependentDialects(DialectRegistry &registry) const override;
+};
+
+void TestComposeSubViewPass::getDependentDialects(
+    DialectRegistry &registry) const {
+  registry.insert<AffineDialect>();
+}
+
+void TestComposeSubViewPass::runOnFunction() {
+  OwningRewritePatternList patterns(&getContext());
+  populateComposeSubViewPatterns(patterns, &getContext());
+  (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+}
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestComposeSubView() {
+  PassRegistration<TestComposeSubViewPass> pass(
+      "test-compose-subview", "Test combining composed subviews");
+}
+} // namespace test
+} // namespace mlir

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

 void registerTestGpuSerializeToCubinPass();
 void registerTestGpuSerializeToHsacoPass();
 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDialect(DialectRegistry &);
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
+void registerTestComposeSubView();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestIRVisitorsPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgFusionTransforms();
 void registerTestLinalgTensorFusionTransforms();
 void registerTestLinalgGreedyFusion();
 void registerTestLinalgHoisting();
#if MLIR_ROCM_CONVERSIONS_ENABLED
   test::registerTestGpuSerializeToHsacoPass();
 #endif
   test::registerTestConvVectorization();
   test::registerTestDecomposeCallGraphTypes();
   test::registerTestDataLayoutQuery();
   test::registerTestDominancePass();
   test::registerTestDynamicPipelinePass();
   test::registerTestExpandTanhPass();
+  test::registerTestComposeSubView();
   test::registerTestGpuParallelLoopMappingPass();
   test::registerTestIRVisitorsPass();
   test::registerTestInterfaces();
   test::registerTestLinalgCodegenStrategy();
   test::registerTestLinalgFusionTransforms();
   test::registerTestLinalgTensorFusionTransforms();
   test::registerTestLinalgGreedyFusion();
   test::registerTestLinalgHoisting();