rename transform to reify-result-shapes

Author: fabianmcg

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

@@ -182,13 +182,40 @@ def ResolveShapedTypeResultDimsPass : Pass<"resolve-shaped-type-result-dims"> {
   ];
 }
 
-def InferStaticShapesPass : Pass<"infer-static-shapes"> {
-  let summary = "Resolve memref.dim of result values";
+def ReifyResultShapesPass : Pass<"reify-result-shapes"> {
+  let summary = "Reifies the results of all `ReifyRankedShapedTypeOpInterface` operations";
   let description = [{
-    The pass resolves memref.dim of result of operations that
-    implement the `InferShapedTypeOpInterface` or
-    `ReifyRankedShapedTypeOpInterface` in terms of shapes of its
-    operands.
+    This pass reifies the shapes of every `ReifyRankedShapedTypeOpInterface`
+    operation with ranked `memref` and `tensor` results. Replacing the
+    operations with their reified versions, and inserting casts when results
+    shapes are updated.
+
+    Example:
+    ```mlir
+    #map = affine_map<(d0) -> (-d0 + 256)>
+    func.func @func(%arg0: f32, %arg1: index, %arg2: tensor<64x?x64xf32>) -> tensor<1x?x64xf32> {
+      %0 = affine.apply #map(%arg1)
+      %extracted_slice = tensor.extract_slice %arg2[0, 0, 0] [1, %arg1, 64] [1, 1, 1] : tensor<64x?x64xf32> to tensor<1x?x64xf32>
+      %padded = tensor.pad %extracted_slice low[0, 0, 0] high[0, %0, 0] {
+      ^bb0(%arg3: index, %arg4: index, %arg5: index):
+        tensor.yield %arg0 : f32
+      } : tensor<1x?x64xf32> to tensor<1x?x64xf32>
+      return %padded : tensor<1x?x64xf32>
+    }
+
+    // mlir-opt --reify-result-shapes
+    #map = affine_map<()[s0] -> (-s0 + 256)>
+    func.func @func(%arg0: f32, %arg1: index, %arg2: tensor<64x?x64xf32>) -> tensor<1x?x64xf32> {
+      %0 = affine.apply #map()[%arg1]
+      %extracted_slice = tensor.extract_slice %arg2[0, 0, 0] [1, %arg1, 64] [1, 1, 1] : tensor<64x?x64xf32> to tensor<1x?x64xf32>
+      %padded = tensor.pad %extracted_slice low[0, 0, 0] high[0, %0, 0] {
+      ^bb0(%arg3: index, %arg4: index, %arg5: index):
+        tensor.yield %arg0 : f32
+      } : tensor<1x?x64xf32> to tensor<1x256x64xf32>
+      %cast = tensor.cast %padded : tensor<1x256x64xf32> to tensor<1x?x64xf32>
+      return %cast : tensor<1x?x64xf32>
+    }
+    ```
   }];
   let dependentDialects = [
     "affine::AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect"

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

@@ -23,6 +23,7 @@ class RewritePatternSet;
 class RewriterBase;
 class Value;
 class ValueRange;
+class ReifyRankedShapedTypeOpInterface;
 
 namespace arith {
 class WideIntEmulationConverter;
@@ -213,6 +214,17 @@ memref::AllocaOp allocToAlloca(
     RewriterBase &rewriter, memref::AllocOp alloc,
     function_ref<bool(memref::AllocOp, memref::DeallocOp)> filter = nullptr);
 
+/// Reifies the results of `op`, potentially replacing `op` with a reified
+/// version. Returns `failure` if `mlir::reifyResultShapes` returned failure,
+/// otherwise it always succeeds. Users of this transform should always expect
+/// it to modify the IR, even when it fails. If any of the result types changes,
+/// the transform will insert cast operations to the old type to keep the IR
+/// consistent.
+///
+/// Note: This transform only works on ranked `memref` or `tensor` results,
+/// other types are ignored.
+LogicalResult reifyOpResultShapes(RewriterBase &rewriter,
+                                  ReifyRankedShapedTypeOpInterface op);
 } // namespace memref
 } // namespace mlir
 

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

@@ -13,6 +13,7 @@ add_mlir_dialect_library(MLIRMemRefTransforms
   IndependenceTransforms.cpp
   MultiBuffer.cpp
   NormalizeMemRefs.cpp
+  ReifyResultShapes.cpp
   ResolveShapedTypeResultDims.cpp
   RuntimeOpVerification.cpp
 

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

@@ -0,0 +1,144 @@
+//===- ReifyResultShapes.cpp - Reify result shapes ------------------------===//
+//
+// 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 transform reifies result shapes of `ReifyRankedShapedTypeOpInterface`
+// operations with ranked `memref` and `tensor` results.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "llvm/Support/InterleavedRange.h"
+
+#define DEBUG_TYPE "reify-result-shapes"
+#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
+
+namespace mlir {
+namespace memref {
+#define GEN_PASS_DEF_REIFYRESULTSHAPESPASS
+#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
+} // namespace memref
+} // namespace mlir
+
+using namespace mlir;
+
+LogicalResult
+mlir::memref::reifyOpResultShapes(RewriterBase &rewriter,
+                                  ReifyRankedShapedTypeOpInterface op) {
+  LLVM_DEBUG({ DBGS() << " reifying op: " << op << "\n"; });
+  // Get the reified out shapes.
+  ReifiedRankedShapedTypeDims reifiedResultShapes;
+  if (failed(mlir::reifyResultShapes(rewriter, op, reifiedResultShapes)) ||
+      reifiedResultShapes.empty()) {
+    return op.emitError() << "failed to get the reified shapes";
+  }
+
+  bool modified = false;
+  // Compute the new output types.
+  SmallVector<Type> outTypes;
+  for (const auto &[oldTy, reifiedShape] :
+       llvm::zip(op->getResultTypes(), reifiedResultShapes)) {
+    // Skip if it's not a memref or tensor type.
+    if (!isa<RankedTensorType, MemRefType>(oldTy)) {
+      outTypes.push_back(oldTy);
+      continue;
+    }
+
+    ShapedType shapedTy = dyn_cast<ShapedType>(oldTy);
+
+    SmallVector<int64_t> shape = llvm::to_vector(shapedTy.getShape());
+    for (auto &&[dim, ofr] : llvm::zip_equal(shape, reifiedShape)) {
+      std::optional<int64_t> maybeCst = getConstantIntValue(ofr);
+      // If the reified dim is dynamic set it appropriately.
+      if (!maybeCst.has_value()) {
+        dim = ShapedType::kDynamic;
+        continue;
+      }
+      // Set the static dim.
+      dim = *maybeCst;
+    }
+
+    // If the shape didn't change continue.
+    if (shape == shapedTy.getShape()) {
+      outTypes.push_back(oldTy);
+      continue;
+    }
+    modified = true;
+    outTypes.push_back(shapedTy.cloneWith(shape, shapedTy.getElementType()));
+  }
+
+  // Return if we don't need to update.
+  if (!modified) {
+    LLVM_DEBUG({ DBGS() << "- op doesn't require update\n"; });
+    return success();
+  }
+
+  LLVM_DEBUG({
+    DBGS() << "- oldTypes: " << llvm::interleaved_array(op->getResultTypes())
+           << " \n";
+    DBGS() << "- outTypes: " << llvm::interleaved_array(outTypes) << " \n";
+  });
+
+  // We now have outTypes that need to be turned to cast ops.
+  Location loc = op->getLoc();
+  SmallVector<Value> newResults;
+  Operation *newOp = rewriter.clone(*op);
+  for (auto [reifiedTy, oldRes] : llvm::zip(outTypes, op->getResults())) {
+    OpResult newRes = newOp->getResult(oldRes.getResultNumber());
+    Type oldTy = oldRes.getType();
+    // Continue if the type remained invariant or is not shaped.
+    if (oldTy == reifiedTy || !isa<MemRefType, RankedTensorType>(oldTy)) {
+      newResults.push_back(newRes);
+      continue;
+    }
+
+    // Update the type.
+    newRes.setType(reifiedTy);
+    if (isa<RankedTensorType>(reifiedTy)) {
+      newResults.push_back(rewriter.create<tensor::CastOp>(loc, oldTy, newRes));
+    } else {
+      assert(isa<MemRefType>(reifiedTy) && "expected a memref type");
+      newResults.push_back(rewriter.create<memref::CastOp>(loc, oldTy, newRes));
+    }
+  }
+
+  LLVM_DEBUG({
+    DBGS() << "- reified results " << llvm::interleaved_array(newResults)
+           << "\n";
+  });
+  rewriter.replaceOp(op, newResults);
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// Pass registration
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct ReifyResultShapesPass final
+    : public memref::impl::ReifyResultShapesPassBase<ReifyResultShapesPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void ReifyResultShapesPass::runOnOperation() {
+  SmallVector<ReifyRankedShapedTypeOpInterface> ops;
+  getOperation()->walk(
+      [&](ReifyRankedShapedTypeOpInterface op) { ops.push_back(op); });
+  IRRewriter rewriter(&getContext());
+  for (ReifyRankedShapedTypeOpInterface op : ops) {
+    rewriter.setInsertionPoint(op);
+    if (failed(memref::reifyOpResultShapes(rewriter, op)))
+      return signalPassFailure();
+  }
+}

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

@@ -20,22 +20,13 @@
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/BuiltinTypeInterfaces.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/Value.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "llvm/Support/ErrorHandling.h"
-#include "llvm/Support/InterleavedRange.h"
-
-#define DEBUG_TYPE "resolve-shaped-type"
-#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
 
 namespace mlir {
 namespace memref {
 #define GEN_PASS_DEF_RESOLVERANKEDSHAPETYPERESULTDIMSPASS
 #define GEN_PASS_DEF_RESOLVESHAPEDTYPERESULTDIMSPASS
-#define GEN_PASS_DEF_INFERSTATICSHAPESPASS
 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
 } // namespace memref
 } // namespace mlir
@@ -114,99 +105,6 @@ struct DimOfReifyRankedShapedTypeOpInterface : public OpRewritePattern<OpTy> {
   }
 };
 
-struct ReifyToInferStaticShapePattern
-    : public OpInterfaceRewritePattern<ReifyRankedShapedTypeOpInterface> {
-  using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
-
-  LogicalResult matchAndRewrite(ReifyRankedShapedTypeOpInterface op,
-                                PatternRewriter &rewriter) const override {
-    LLVM_DEBUG(
-        { DBGS() << "ReifyToInferStaticShapePattern on " << op << "\n"; });
-
-    bool rewriteToMoreStatic = false;
-    ReifiedRankedShapedTypeDims reifiedResultShapes;
-    if (failed(reifyResultShapes(rewriter, op, reifiedResultShapes)) ||
-        reifiedResultShapes.empty()) {
-      LLVM_DEBUG({ DBGS() << "reifyResultShapes failed\n"; });
-      return failure();
-    }
-
-    SmallVector<Type> newTypes;
-    for (auto [t, reifiedShape] :
-         llvm::zip(op->getResultTypes(), reifiedResultShapes)) {
-      ShapedType st = dyn_cast<ShapedType>(t);
-      if (!st)
-        continue;
-
-      SmallVector<int64_t> newShape;
-      for (const auto &[s, ofr] :
-           llvm::zip_equal(st.getShape(), reifiedShape)) {
-        std::optional<int64_t> maybeCst = getConstantIntValue(ofr);
-        // Reification does not add static information, just use existing shape.
-        if (!maybeCst.has_value()) {
-          newShape.push_back(s);
-          continue;
-        }
-        int64_t cst = *maybeCst;
-        assert((ShapedType::isDynamic(s) || s == cst) &&
-               "constants must agree!");
-        newShape.push_back(cst);
-      }
-
-      if (newShape == st.getShape()) {
-        newTypes.push_back(t);
-        continue;
-      }
-
-      rewriteToMoreStatic = true;
-      Type newType = st.cloneWith(newShape, st.getElementType());
-      newTypes.push_back(newType);
-    }
-
-    LLVM_DEBUG({
-      DBGS() << "--oldTypes: " << llvm::interleaved_array(op->getResultTypes())
-             << " \n";
-      DBGS() << "--newTypes: " << llvm::interleaved_array(newTypes) << " \n";
-    });
-    if (!rewriteToMoreStatic) {
-      LLVM_DEBUG({ DBGS() << "not more static\n"; });
-      return failure();
-    }
-
-    // We now have newTypes that need to be turned to tensor::CastOp.
-    Location loc = op->getLoc();
-    SmallVector<Value> newResults;
-    Operation *newOp = rewriter.clone(*op);
-    for (auto [nt, oldVal] : llvm::zip(newTypes, op->getResults())) {
-      Type ot = oldVal.getType();
-      OpResult newResult = newOp->getResult(oldVal.getResultNumber());
-      if (ot == nt) {
-        newResults.push_back(newResult);
-        continue;
-      }
-      newResult.setType(nt);
-      if (isa<RankedTensorType>(nt)) {
-        newResults.push_back(
-            rewriter.create<tensor::CastOp>(loc, ot, newResult));
-      } else if (isa<MemRefType>(nt)) {
-        newResults.push_back(
-            rewriter.create<memref::CastOp>(loc, ot, newResult));
-      } else {
-        llvm_unreachable("expected RankedTensorType or MemRefType");
-      }
-    }
-
-    LLVM_DEBUG({
-      op->getParentOp()->dump();
-      DBGS() << "replace op " << *op << "\n";
-      DBGS() << "with newResults " << llvm::interleaved_array(newResults)
-             << "\n\n\n\n";
-    });
-    rewriter.replaceAllOpUsesWith(op, newResults);
-    return success();
-  }
-};
-
 /// Fold dim ops of iter_args to dim ops of their respective init args. E.g.:
 ///
 /// ```
@@ -277,11 +175,6 @@ struct ResolveShapedTypeResultDimsPass final
   void runOnOperation() override;
 };
 
-struct InferStaticShapesPass final
-    : public memref::impl::InferStaticShapesPassBase<InferStaticShapesPass> {
-  void runOnOperation() override;
-};
-
 } // namespace
 
 void memref::populateResolveRankedShapedTypeResultDimsPatterns(
@@ -299,11 +192,6 @@ void memref::populateResolveShapedTypeResultDimsPatterns(
       patterns.getContext());
 }
 
-void memref::populateReifyToInferStaticShapePatterns(
-    RewritePatternSet &patterns) {
-  patterns.add<ReifyToInferStaticShapePattern>(patterns.getContext());
-}
-
 void ResolveRankedShapeTypeResultDimsPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
@@ -318,17 +206,3 @@ void ResolveShapedTypeResultDimsPass::runOnOperation() {
   if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
     return signalPassFailure();
 }
-
-void InferStaticShapesPass::runOnOperation() {
-  RewritePatternSet patterns(&getContext());
-  patterns.add<ReifyToInferStaticShapePattern>(&getContext());
-  FrozenRewritePatternSet frozenPatterns(std::move(patterns));
-
-  SmallVector<Operation *> opsToSimplify;
-  getOperation()->walk([&](ReifyRankedShapedTypeOpInterface op) {
-    opsToSimplify.push_back(op);
-  });
-  (void)applyOpPatternsGreedily(opsToSimplify, frozenPatterns,
-                                GreedyRewriteConfig().setStrictness(
-                                    GreedyRewriteStrictness::ExistingOps));
-}