[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

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/IR/MemRefOps.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/STLExtras.h"

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();
   }
 };
+
+template <typename ResultsContainer>
+bool replaceConstantUsesOf(PatternRewriter &rewriter,
+                           memref::ExtractStridedMetadataOp metadataOp,
+                           ArrayRef<int64_t> maybeConstants,
+                           ResultsContainer results,
+                           llvm::function_ref<bool(int64_t)> isDynamic) {
+  bool atLeastOneReplacement = false;
+  for (auto [maybeConstant, result] : llvm::zip(maybeConstants, results)) {
+    // 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>(
+        metadataOp.getLoc(), 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);
+    assert(succeeded(res) && "must be a strided memref type");
+
+    bool atLeastOneReplacement = replaceConstantUsesOf(
+        rewriter, metadataOp, ArrayRef<int64_t>(offset),
+        ArrayRef<TypedValue<IndexType>>(metadataOp.getOffset()),
+        ShapedType::isDynamicStrideOrOffset);
+    atLeastOneReplacement |=
+        replaceConstantUsesOf(rewriter, metadataOp, memrefType.getShape(),
+                              metadataOp.getSizes(), ShapedType::isDynamic);
+    atLeastOneReplacement |= replaceConstantUsesOf(
+        rewriter, metadataOp, strides, metadataOp.getStrides(),
+        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