[mlir][MemRef] Add pattern that forwards constant strided metadata.

[nicolasvasilache]

Summary

memref.extract_strided_metadata can forward constants independently of the
exsistence of other operations such as subview or reshape.

[chelini]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

16 ↗

(On Diff #462730)

Why is this needed?

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

457

Introduces the following warning: warning: extra ';' outside of a function is incompatible with C++98.

[ftynse]

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

434

Please document top-level entities.

470

(void)res;

This will break Werror bots otherwise.

mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

7–10

Any reason why this needs to be a dag?

[nicolasvasilache]

mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

7–10

I wouldn't want to encode an expected ordering of constant materialization (or expect a brittle MLIR enforced order of constants), any order is fine.

[nicolasvasilache]

Address'nRebase

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

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 namespace mlir {
 namespace memref {
 #define GEN_PASS_DEF_SIMPLIFYEXTRACTSTRIDEDMETADATA
 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
 } // namespace memref
 } // namespace mlir
assert(idx == sourceRank &&
            "We should have visited all the input dimensions");
     assert(baseIdxInResult == expandShapeRank + 2 &&
            "We should have populated all the values");
     rewriter.replaceOp(
         op, getValueOrCreateConstantIndexOp(rewriter, origLoc, results));
     return success();
   }
 };
+
+/// Helper function to perform the replacement of all constant uses of `values`
+/// by a materialized constant extracted from `maybeConstants`.
+/// `values` and `maybeConstants` are expected to have the same size.
+template <typename Container>
+bool replaceConstantUsesOf(PatternRewriter &rewriter, Location loc,
+                           Container values, ArrayRef<int64_t> maybeConstants,
+                           llvm::function_ref<bool(int64_t)> isDynamic) {
+  assert(values.size() == maybeConstants.size() &&
+         " expected values and maybeConstants of the same size");
+  bool atLeastOneReplacement = false;
+  for (auto [maybeConstant, result] : llvm::zip(maybeConstants, values)) {
+    // Don't materialize a constant if there are no uses: this would indice
+    // infinite loops in the driver.
+    if (isDynamic(maybeConstant) || result.use_empty())
+      continue;
+    Value constantVal =
+        rewriter.create<arith::ConstantIndexOp>(loc, maybeConstant);
+    for (Operation *op : llvm::make_early_inc_range(result.getUsers())) {
+      rewriter.startRootUpdate(op);
+      // updateRootInplace: lambda cannot capture structured bindings in C++17
+      // yet.
+      op->replaceUsesOfWith(result, constantVal);
+      rewriter.finalizeRootUpdate(op);
+      atLeastOneReplacement = true;
+    }
+  }
+  return atLeastOneReplacement;
+}
+
+// Forward propagate all constants information from an ExtractStridedMetadataOp.
+struct ForwardStaticMetadata
+    : public OpRewritePattern<memref::ExtractStridedMetadataOp> {
+  using OpRewritePattern<memref::ExtractStridedMetadataOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp metadataOp,
+                                PatternRewriter &rewriter) const override {
+    auto memrefType = metadataOp.getSource().getType().cast<MemRefType>();
+    SmallVector<int64_t> strides;
+    int64_t offset;
+    LogicalResult res = getStridesAndOffset(memrefType, strides, offset);
+    (void)res;
+    assert(succeeded(res) && "must be a strided memref type");
+
+    bool atLeastOneReplacement = replaceConstantUsesOf(
+        rewriter, metadataOp.getLoc(),
+        ArrayRef<TypedValue<IndexType>>(metadataOp.getOffset()),
+        ArrayRef<int64_t>(offset), ShapedType::isDynamicStrideOrOffset);
+    atLeastOneReplacement |= replaceConstantUsesOf(
+        rewriter, metadataOp.getLoc(), metadataOp.getSizes(),
+        memrefType.getShape(), ShapedType::isDynamic);
+    atLeastOneReplacement |= replaceConstantUsesOf(
+        rewriter, metadataOp.getLoc(), metadataOp.getStrides(), strides,
+        ShapedType::isDynamicStrideOrOffset);
+
+    return success(atLeastOneReplacement);
+  }
+};
 } // namespace
 
 void memref::populateSimplifyExtractStridedMetadataOpPatterns(
     RewritePatternSet &patterns) {
-  patterns.add<ExtractStridedMetadataOpSubviewFolder,
-               ExtractStridedMetadataOpExpandShapeFolder>(
+  patterns
+      .add<ExtractStridedMetadataOpSubviewFolder,
+           ExtractStridedMetadataOpExpandShapeFolder, ForwardStaticMetadata>(
           patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 

mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

 // RUN: mlir-opt --simplify-extract-strided-metadata -split-input-file %s -o - | FileCheck %s
 
+// CHECK-LABEL: func @extract_strided_metadata_constants
+//  CHECK-SAME: (%[[ARG:.*]]: memref<5x4xf32, strided<[4, 1], offset: 2>>)
+func.func @extract_strided_metadata_constants(%base: memref<5x4xf32, strided<[4, 1], offset: 2>>)
+    -> (memref<f32>, index, index, index, index, index) {
+  //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+  //   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+  //   CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+  //   CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index
+  
+  //       CHECK: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]]
+  %base_buffer, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %base :
+    memref<5x4xf32, strided<[4,1], offset:2>>
+    -> memref<f32>, index, index, index, index, index
+
+  // CHECK: %[[BASE]], %[[C2]], %[[C5]], %[[C4]], %[[C4]], %[[C1]]
+  return %base_buffer, %offset, %sizes#0, %sizes#1, %strides#0, %strides#1 :
+    memref<f32>, index, index, index, index, index
+}
+
 // -----
 
 // Check that we simplify extract_strided_metadata of subview to
 // base_buf, base_offset, base_sizes, base_strides = extract_strided_metadata
 // strides = base_stride_i * subview_stride_i
 // offset = base_offset + sum(subview_offsets_i * strides_i).
 //
 // This test also checks that we don't create useless arith operations