[mlir] Introduce ViewLikeConsumerOpInterface

[christopherbate]

Summary

ViewLikeConsumerOpInterface represents an operation that can "fold" a
ViewLikeInterfaceOp producer into itself for one or more of its operands.

The new interface is meant to address issues with the current MemRef
dialect fold-memref-alias-ops transformation. The transformation
implements the ability to "fold" (or "skip over"/"bypass"/"consume")
producing ops such as memref.subview, memref.collapse_shape,
and memref.expand_shape into operations such as memref.load as
well as other dialects' ops such as vector.transfer_read. So
effectively it is implementing exactly what interfaces are for, but we
now need a real interface for other dialects (in or out of tree) to
hook into this process.

The new interface is not specific to memref, it could be used for
equivalent operations on tensors.

This is also the opportunity to position some utilities that handle the
common operations (e.g. composing "offsetStrideAndSliceInterface" ops),
allowing for de-duplication of code across the tensor and memref cases.
However, due to layering issues, we have to put these utilities under
Dialect/Affine. They can't go in Dialect/Utils or
Interfaces/ViewLikeInterface due to creation of circular dependencies.

[christopherbate]

mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp

14

Pulled from recent commit: https://reviews.llvm.org/D134294

mlir/include/mlir/Dialect/Affine/ViewLikeInterfaceUtils.h

+//===- ViewLikeInterfaceUtils.h - View-like operations interface utilities-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_AFFINE_INDEXUTILS_H
+#define MLIR_DIALECT_AFFINE_INDEXUTILS_H
+
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
+
+namespace mlir {
+
+/// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use
+/// when combining a producer slice **into** a consumer slice.
+///
+/// This function performs the following computation:
+/// - Combined offsets = producer_offsets * consumer_strides + consumer_offsets
+/// - Combined sizes = consumer_sizes
+/// - Combined strides = producer_strides * consumer_strides
+LogicalResult
+mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc,
+                            ArrayRef<OpFoldResult> producerOffsets,
+                            ArrayRef<OpFoldResult> producerSizes,
+                            ArrayRef<OpFoldResult> producerStrides,
+                            const llvm::SmallBitVector &droppedProducerDims,
+                            ArrayRef<OpFoldResult> consumerOffsets,
+                            ArrayRef<OpFoldResult> consumerSizes,
+                            ArrayRef<OpFoldResult> consumerStrides,
+                            SmallVector<OpFoldResult> &combinedOffsets,
+                            SmallVector<OpFoldResult> &combinedSizes,
+                            SmallVector<OpFoldResult> &combinedStrides);
+
+/// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use
+/// when combining a `producer` slice op **into** a `consumer` slice op.
+LogicalResult
+mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc,
+                            OffsetSizeAndStrideOpInterface producer,
+                            OffsetSizeAndStrideOpInterface consumer,
+                            const llvm::SmallBitVector &droppedProducerDims,
+                            SmallVector<OpFoldResult> &combinedOffsets,
+                            SmallVector<OpFoldResult> &combinedSizes,
+                            SmallVector<OpFoldResult> &combinedStrides);
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_AFFINE_INDEXUTILS_H

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

+//===- ViewLikeConsumerImpl.h - Impl. of ViewLikeConsumerOpInterface ------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_MEMREF_TRANSFORMS_VIEWLIKECONSUMEROPINTERFACEIMPL_H
+#define MLIR_DIALECT_MEMREF_TRANSFORMS_VIEWLIKECONSUMEROPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace memref {
+void registerViewLikeConsumerOpInterfaceExternalModels(
+    DialectRegistry &registry);
+} // namespace memref
+} // namespace mlir
+
+#endif // MLIR_DIALECT_MEMREF_TRANSFORMS_VIEWLIKECONSUMEROPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/ViewLikeConsumerImpl.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
inline void registerAllDialects(DialectRegistry &registry) {
   linalg::registerTilingInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
   sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
   tensor::registerTilingInterfaceExternalModels(registry);
   vector::registerBufferizableOpInterfaceExternalModels(registry);
+  memref::registerViewLikeConsumerOpInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.
 inline void registerAllDialects(MLIRContext &context) {
   DialectRegistry registry;
   registerAllDialects(registry);
   context.appendDialectRegistry(registry);
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 /// Decompose a vector of mixed static or dynamic strides/offsets into the
 /// corresponding pair of arrays. This is the inverse function of
 /// `getMixedSizes`.
 std::pair<ArrayAttr, SmallVector<Value>>
 decomposeMixedSizes(OpBuilder &b,
                     const SmallVectorImpl<OpFoldResult> &mixedValues);
 
 class OffsetSizeAndStrideOpInterface;
-
+class RewriterBase;
 namespace detail {
 
 LogicalResult verifyOffsetSizeAndStrideOp(OffsetSizeAndStrideOpInterface op);
 
 bool sameOffsetsSizesAndStrides(
     OffsetSizeAndStrideOpInterface a, OffsetSizeAndStrideOpInterface b,
     llvm::function_ref<bool(OpFoldResult, OpFoldResult)> cmp);
 

mlir/include/mlir/Interfaces/ViewLikeInterface.td

def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface"> {
   }];
 
   let verify = [{
     return ::mlir::detail::verifyOffsetSizeAndStrideOp(
         ::mlir::cast<::mlir::OffsetSizeAndStrideOpInterface>($_op));
   }];
 }
 
+def ViewLikeConsumerOpInterface : OpInterface<"ViewLikeConsumerOpInterface"> {
+  let description = [{
+    An operation that has one or more operands that can be "fold" the producer of
+    the operand if the producer is a ViewLikeOpInterface. An example is
+    `memref.load %A [offsets]`. If `%A` is the result of `memref.subview`, then
+    we can update the op to act on the source of the subview with appropriately
+    updated offsets.
+  }];
+
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the OpOperands that can "consume" a producing ViewLikeOpInterface.
+      }],
+      /*retTy=*/"::mlir::SmallVector<::mlir::OpOperand*>",
+      /*methodName=*/"getViewLikeConsumerOperands",
+      /*args=*/(ins)
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the new indices for the given operand and producer op.
+      }],
+      /*retTy=*/"::mlir::FailureOr<::mlir::SmallVector<::mlir::OpFoldResult>>",
+      /*methodName=*/"resolveSourceIndices",
+      /*args=*/(ins "::mlir::OpBuilder&":$builder, "::mlir::OpOperand*":$operand,
+        "::mlir::Operation*":$producer)
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+
+      }],
+      /*retTy=*/"::mlir::LogicalResult",
+      /*methodName=*/"replaceViewLikeConsumerOp",
+      /*args=*/(ins 
+        "::mlir::RewriterBase&":$rewriter,
+        "::mlir::OpOperand*":$operand,
+        "::mlir::Value":$newOperand,
+        "::mlir::ArrayRef<::mlir::OpFoldResult>":$newIndices)
+    >
+  ];
+}
+
 #endif // MLIR_INTERFACES_VIEWLIKEINTERFACE

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

 add_mlir_dialect_library(MLIRAffineUtils
   LoopFusionUtils.cpp
   LoopUtils.cpp
   Utils.cpp
+  ViewLikeInterfaceUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineAnalysis
   MLIRAnalysis
   MLIRArithmeticUtils
   MLIRMemRefDialect
   MLIRTransformUtils
+  MLIRViewLikeInterface
   )

mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp

+//===- ViewLikeInterfaceUtils.cpp - View-like operations interface utils --===//
+//
+// 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/ViewLikeInterfaceUtils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/IR/AffineExpr.h"
+
+using namespace mlir;
+
+LogicalResult mlir::mergeOffsetsSizesAndStrides(
+    OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> producerOffsets,
+    ArrayRef<OpFoldResult> producerSizes,
+    ArrayRef<OpFoldResult> producerStrides,
+    const llvm::SmallBitVector &droppedProducerDims,
+    ArrayRef<OpFoldResult> consumerOffsets,
+    ArrayRef<OpFoldResult> consumerSizes,
+    ArrayRef<OpFoldResult> consumerStrides,
+    SmallVector<OpFoldResult> &combinedOffsets,
+    SmallVector<OpFoldResult> &combinedSizes,
+    SmallVector<OpFoldResult> &combinedStrides) {
+  combinedOffsets.resize(producerOffsets.size());
+  combinedSizes.resize(producerOffsets.size());
+  combinedStrides.resize(producerOffsets.size());
+
+  AffineExpr s0, s1, d0;
+  bindDims(builder.getContext(), d0);
+  bindSymbols(builder.getContext(), s0, s1);
+
+  unsigned consumerPos = 0;
+  for (auto i : llvm::seq<unsigned>(0, producerOffsets.size())) {
+    if (droppedProducerDims.test(i)) {
+      // For dropped dims, get the values from the producer.
+      combinedOffsets[i] = producerOffsets[i];
+      combinedSizes[i] = producerSizes[i];
+      combinedStrides[i] = producerStrides[i];
+      continue;
+    }
+    SmallVector<OpFoldResult> offsetSymbols, strideSymbols;
+    // The combined offset is computed as
+    //    producer_offset + consumer_offset * producer_strides.
+    combinedOffsets[i] = makeComposedFoldedAffineApply(
+        builder, loc, d0 * s0 + s1,
+        {consumerOffsets[consumerPos], producerStrides[i], producerOffsets[i]});
+    combinedSizes[i] = consumerSizes[consumerPos];
+    // The combined stride is computed as
+    //    consumer_stride * producer_stride.
+    combinedStrides[i] = makeComposedFoldedAffineApply(
+        builder, loc, d0 * s0,
+        {consumerStrides[consumerPos], producerStrides[i]});
+
+    consumerPos++;
+  }
+  return success();
+}
+
+LogicalResult mlir::mergeOffsetsSizesAndStrides(
+    OpBuilder &builder, Location loc, OffsetSizeAndStrideOpInterface producer,
+    OffsetSizeAndStrideOpInterface consumer,
+    const llvm::SmallBitVector &droppedProducerDims,
+    SmallVector<OpFoldResult> &combinedOffsets,
+    SmallVector<OpFoldResult> &combinedSizes,
+    SmallVector<OpFoldResult> &combinedStrides) {
+  SmallVector<OpFoldResult> consumerOffsets = consumer.getMixedOffsets();
+  SmallVector<OpFoldResult> consumerSizes = consumer.getMixedSizes();
+  SmallVector<OpFoldResult> consumerStrides = consumer.getMixedStrides();
+  SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets();
+  SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes();
+  SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides();
+  return mlir::mergeOffsetsSizesAndStrides(
+      builder, loc, producerOffsets, producerSizes, producerStrides,
+      droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
+      combinedOffsets, combinedSizes, combinedStrides);
+}

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

 add_mlir_dialect_library(MLIRMemRefTransforms
   ComposeSubView.cpp
   ExpandOps.cpp
   FoldMemRefAliasOps.cpp
   MultiBuffer.cpp
   NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
   SimplifyExtractStridedMetadata.cpp
+  ViewLikeConsumerImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen
 
   LINK_LIBS PUBLIC

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

if (auto affineLoadOp = dyn_cast<AffineLoadOp>(loadOp.getOperation())) {
     indices.assign(expandedIndices.begin(), expandedIndices.end());
   }
   SmallVector<Value, 4> sourceIndices;
   if (failed(resolveSourceIndicesSubView(loadOp.getLoc(), rewriter, subViewOp,
                                          indices, sourceIndices)))
     return failure();
 
   llvm::TypeSwitch<Operation *, void>(loadOp)
-      .Case<AffineLoadOp, memref::LoadOp>([&](auto op) {
+      .Case<AffineLoadOp>([&](auto op) {
         rewriter.replaceOpWithNewOp<decltype(op)>(loadOp, subViewOp.source(),
                                                   sourceIndices);
       })
       .Case([&](vector::TransferReadOp transferReadOp) {
         rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
             transferReadOp, transferReadOp.getVectorType(), subViewOp.source(),
             sourceIndices,
             getPermutationMapAttr(rewriter.getContext(), subViewOp,
llvm::TypeSwitch<Operation *, void>(storeOp)
         rewriter.replaceOpWithNewOp<decltype(op)>(
             storeOp, storeOp.getValue(), collapseShapeOp.getViewSource(),
             sourceIndices);
       })
       .Default([](Operation *) { llvm_unreachable("unexpected operation."); });
   return success();
 }
 
+struct RewriteViewLikeMemRefConsumers
+    : ::mlir::OpInterfaceRewritePattern<ViewLikeConsumerOpInterface> {
+  using OpInterfaceRewritePattern<
+      ViewLikeConsumerOpInterface>::OpInterfaceRewritePattern;
+
+  LogicalResult matchAndRewrite(ViewLikeConsumerOpInterface op,
+                                PatternRewriter &rewriter) const override {
+
+    SmallVector<OpOperand *> consumers = op.getViewLikeConsumerOperands();
+    bool madeChange = false;
+    for (OpOperand *operand : consumers) {
+      if (!operand->get().getType().isa<MemRefType>())
+        continue;
+      auto producer = operand->get().getDefiningOp<ViewLikeOpInterface>();
+      if (!producer)
+        continue;
+      FailureOr<SmallVector<OpFoldResult>> indices =
+          op.resolveSourceIndices(rewriter, operand, producer.getOperation());
+      if (failed(indices))
+        continue;
+      if (failed(op.replaceViewLikeConsumerOp(
+              rewriter, operand, producer.getViewSource(), *indices)))
+        continue;
+      madeChange = true;
+    }
+    return success(madeChange);
+  }
+};
+
 void memref::populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns) {
   patterns.add<LoadOpOfSubViewOpFolder<AffineLoadOp>,
-               LoadOpOfSubViewOpFolder<memref::LoadOp>,
                LoadOpOfSubViewOpFolder<vector::TransferReadOp>,
                StoreOpOfSubViewOpFolder<AffineStoreOp>,
                StoreOpOfSubViewOpFolder<memref::StoreOp>,
                StoreOpOfSubViewOpFolder<vector::TransferWriteOp>,
                LoadOpOfExpandShapeOpFolder<AffineLoadOp>,
                LoadOpOfExpandShapeOpFolder<memref::LoadOp>,
                StoreOpOfExpandShapeOpFolder<AffineStoreOp>,
                StoreOpOfExpandShapeOpFolder<memref::StoreOp>,
                LoadOpOfCollapseShapeOpFolder<AffineLoadOp>,
                LoadOpOfCollapseShapeOpFolder<memref::LoadOp>,
                StoreOpOfCollapseShapeOpFolder<AffineStoreOp>,
-               StoreOpOfCollapseShapeOpFolder<memref::StoreOp>>(
-      patterns.getContext());
+               StoreOpOfCollapseShapeOpFolder<memref::StoreOp>,
+               RewriteViewLikeMemRefConsumers>(patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 

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

+//===- ViewLikeConsumerImpl.cpp - Impl. of ViewLikeConsumerOpInterface ----===//
+//
+// 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/ViewLikeConsumerImpl.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
+#include "mlir/Dialect/Arithmetic/Utils/Utils.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+
+using namespace mlir;
+using namespace memref;
+
+namespace mlir {
+namespace memref {
+namespace {
+
+/// Enables `memref.load` to skip over `memref.subview` producers.
+struct LoadOpInterface
+    : public ViewLikeConsumerOpInterface::ExternalModel<LoadOpInterface,
+                                                        memref::LoadOp> {
+  SmallVector<OpOperand *> getViewLikeConsumerOperands(Operation *op) const {
+    return {&op->getOpOperand(0)};
+  }
+
+  FailureOr<SmallVector<OpFoldResult>>
+  resolveSourceIndices(Operation *op, OpBuilder &builder,
+                       ::mlir::OpOperand *operand,
+                       ::mlir::Operation *producer) const {
+    auto loadOp = cast<LoadOp>(op);
+
+    // Handle the case when the producer is a `memref.subview`.
+    if (auto ifaceOp = dyn_cast<memref::SubViewOp>(producer)) {
+      auto one = builder.getIndexAttr(1);
+      SmallVector<OpFoldResult> indices =
+          getAsOpFoldResult(loadOp.getIndices());
+      SmallVector<OpFoldResult> ones(indices.size(), one);
+      SmallVector<OpFoldResult> combinedOffsets, combinedSizes, combinedStrides;
+      if (failed(mergeOffsetsSizesAndStrides(
+              builder, op->getLoc(), ifaceOp.getMixedOffsets(),
+              ifaceOp.getMixedSizes(), ifaceOp.getMixedStrides(),
+              ifaceOp.getDroppedDims(), indices, ones, ones, combinedOffsets,
+              combinedSizes, combinedStrides)))
+        return failure();
+      return combinedOffsets;
+    }
+
+    // TODO: handle the `memref.collapse_shape` and `memref.expand_shape` cases.
+
+    return failure();
+  }
+
+  LogicalResult replaceViewLikeConsumerOp(
+      Operation *op, RewriterBase &rewriter, ::mlir::OpOperand *operand,
+      ::mlir::Value newOperand,
+      ::mlir::ArrayRef<::mlir::OpFoldResult> newIndices) const {
+    assert(operand->getOperandNumber() == 0 &&
+           "only memref operand can be replaced");
+    SmallVector<Value> newIndexVals =
+        getValueOrCreateConstantIndexOp(rewriter, op->getLoc(), newIndices);
+    rewriter.replaceOpWithNewOp<LoadOp>(op, newOperand, newIndexVals);
+    return success();
+  }
+};
+} // namespace
+} // namespace memref
+} // namespace mlir
+
+void mlir::memref::registerViewLikeConsumerOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, MemRefDialect *dialect) {
+    memref::LoadOp::attachInterface<LoadOpInterface>(*ctx);
+    ctx->loadDialect<AffineDialect>();
+  });
+}

mlir/lib/Interfaces/ViewLikeInterface.cpp

 //===- ViewLikeInterface.cpp - View-like operations in MLIR ---------------===//
 //
 // 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/Interfaces/ViewLikeInterface.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ViewLike Interfaces
 //===----------------------------------------------------------------------===//
 
 /// Include the definitions of the loop-like interfaces.

mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir

 // RUN: mlir-opt -fold-memref-alias-ops -split-input-file %s -o - | FileCheck %s
 
 func.func @fold_static_stride_subview_with_load(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) -> f32 {
   %0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] : memref<12x32xf32> to memref<4x4xf32, strided<[64, 3], offset: ?>>
   %1 = memref.load %0[%arg3, %arg4] : memref<4x4xf32, strided<[64, 3], offset: ?>>
     return %1 : f32
 }
-//  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
-//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0] -> (d0 * 3 + s0)>
+//  CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 3)>
 //      CHECK: func @fold_static_stride_subview_with_load
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: memref<12x32xf32>
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG3:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG4:[a-zA-Z0-9_]+]]: index
-//  CHECK-DAG:   %[[I1:.+]] = affine.apply #[[MAP0]](%[[ARG3]])[%[[ARG1]]]
-//  CHECK-DAG:   %[[I2:.+]] = affine.apply #[[MAP1]](%[[ARG4]])[%[[ARG2]]]
+//  CHECK-DAG:   %[[I1:.+]] = affine.apply #[[MAP0]]()[%[[ARG1]], %[[ARG3]]]
+//  CHECK-DAG:   %[[I2:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[ARG4]]]
 //      CHECK:   memref.load %[[ARG0]][%[[I1]], %[[I2]]]
 
 // -----
 
 func.func @fold_dynamic_stride_subview_with_load(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index) -> f32 {
   %0 = memref.subview %arg0[%arg1, %arg2][4, 4][%arg5, %arg6] :
     memref<12x32xf32> to memref<4x4xf32, strided<[?, ?], offset: ?>>
   %1 = memref.load %0[%arg3, %arg4] : memref<4x4xf32, strided<[?, ?], offset: ?>>
   return %1 : f32
 }
-//  CHECK-DAG: #[[MAP:.+]] = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)>
+//  CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s1 + s2 * s0)>
 //      CHECK: func @fold_dynamic_stride_subview_with_load
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: memref<12x32xf32>
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG3:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG4:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG5:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG6:[a-zA-Z0-9_]+]]: index
-//  CHECK-DAG:   %[[I1:.+]] = affine.apply #[[MAP]](%[[ARG3]])[%[[ARG5]], %[[ARG1]]]
-//  CHECK-DAG:   %[[I2:.+]] = affine.apply #[[MAP]](%[[ARG4]])[%[[ARG6]], %[[ARG2]]]
+//  CHECK-DAG:   %[[I1:.+]] = affine.apply #[[MAP]]()[%[[ARG5]], %[[ARG1]], %[[ARG3]]]
+//  CHECK-DAG:   %[[I2:.+]] = affine.apply #[[MAP]]()[%[[ARG6]], %[[ARG2]], %[[ARG4]]]
 //      CHECK:   memref.load %[[ARG0]][%[[I1]], %[[I2]]]
 
 // -----
 
 func.func @fold_static_stride_subview_with_store(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : f32) {
   %0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] :
     memref<12x32xf32> to memref<4x4xf32, strided<[64, 3], offset: ?>>
   memref.store %arg5, %0[%arg3, %arg4] : memref<4x4xf32, strided<[64, 3], offset: ?>>
func.func @fold_rank_reducing_subview_with_load
      %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index,
      %arg7 : index, %arg8 : index, %arg9 : index, %arg10: index,
      %arg11 : index, %arg12 : index, %arg13 : index, %arg14: index,
      %arg15 : index, %arg16 : index) -> f32 {
   %0 = memref.subview %arg0[%arg1, %arg2, %arg3, %arg4, %arg5, %arg6][4, 1, 1, 4, 1, 1][%arg7, %arg8, %arg9, %arg10, %arg11, %arg12] : memref<?x?x?x?x?x?xf32> to memref<4x1x4x1xf32, strided<[?, ?, ?, ?], offset: ?>>
   %1 = memref.load %0[%arg13, %arg14, %arg15, %arg16] : memref<4x1x4x1xf32, strided<[?, ?, ?, ?], offset: ?>>
   return %1 : f32
 }
-//  CHECK-DAG: #[[MAP:.+]] = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)>
+//  CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s1 + s2 * s0)>
 //      CHECK: func @fold_rank_reducing_subview_with_load
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?x?x?x?x?xf32>
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG3:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG4:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG5:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG6:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG7:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG8:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG9:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG10:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG11:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG12:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG13:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG14:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG15:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG16:[a-zA-Z0-9_]+]]: index
-//  CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
-//  CHECK-DAG:   %[[I1:.+]] = affine.apply #[[MAP]](%[[ARG13]])[%[[ARG7]], %[[ARG1]]]
-//  CHECK-DAG:   %[[I2:.+]] = affine.apply #[[MAP]](%[[C0]])[%[[ARG8]], %[[ARG2]]]
-//  CHECK-DAG:   %[[I3:.+]] = affine.apply #[[MAP]](%[[ARG14]])[%[[ARG9]], %[[ARG3]]]
-//  CHECK-DAG:   %[[I4:.+]] = affine.apply #[[MAP]](%[[ARG15]])[%[[ARG10]], %[[ARG4]]]
-//  CHECK-DAG:   %[[I5:.+]] = affine.apply #[[MAP]](%[[ARG16]])[%[[ARG11]], %[[ARG5]]]
-//  CHECK-DAG:   %[[I6:.+]] = affine.apply #[[MAP]](%[[C0]])[%[[ARG12]], %[[ARG6]]]
-//      CHECK:   memref.load %[[ARG0]][%[[I1]], %[[I2]], %[[I3]], %[[I4]], %[[I5]], %[[I6]]]
+//  CHECK-DAG:   %[[I1:.+]] = affine.apply #[[MAP]]()[%[[ARG7]], %[[ARG1]], %[[ARG13]]]
+//  CHECK-DAG:   %[[I3:.+]] = affine.apply #[[MAP]]()[%[[ARG9]], %[[ARG3]], %[[ARG14]]]
+//  CHECK-DAG:   %[[I4:.+]] = affine.apply #[[MAP]]()[%[[ARG10]], %[[ARG4]], %[[ARG15]]]
+//  CHECK-DAG:   %[[I5:.+]] = affine.apply #[[MAP]]()[%[[ARG11]], %[[ARG5]], %[[ARG16]]]
+//      CHECK:   memref.load %[[ARG0]][%[[I1]], %[[ARG2]], %[[I3]], %[[I4]], %[[I5]], %[[ARG6]]]
 
 // -----
 
 func.func @fold_vector_transfer_read_with_rank_reduced_subview(
     %arg0 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>,
     %arg1: index, %arg2 : index, %arg3 : index, %arg4: index, %arg5 : index,
     %arg6 : index) -> vector<4xf32> {
   %cst = arith.constant 0.0 : f32