[mlir] Add file to implement bufferization for shape ops.

[tpopp]

Summary

This adds a shape-bufferize pass and implements the pattern for
shape.assuming.

[silvas]

mlir/lib/Dialect/Shape/Transforms/ShapeBufferize.cpp

27 ↗

(On Diff #293400)

I think this pattern is actually not specific to bufferization.

I think this pattern can be like populateFuncOpTypeConversionPattern (works with any type converter, not just a BufferAssignmentTypeConverter)

A pattern has a member typeConverter already (you can also call getTypeConverter()), which is all you need here.

I have a couple patterns like this in npcomp that I have yet to upstream and there doesn't seem to be a need for it to be tied specifically to buffer conversion: https://github.com/llvm/mlir-npcomp/blob/681c4e1d4ac879921a852e6243288fd159de222a/lib/E2E/TensorToMemref/LowerStructuralToMemref.cpp#L31

Proposal:

expose populateShapeOpsTypeConversionPatterns which takes a generic TypeConverter.
make this a regular OpConversionPattern

93 ↗

(On Diff #293400)

Is "bufferize" the consistent terminology we are going to use from now on upstream for tensor->memref conversion? I don't see that term upstream yet and just wanted to make sure we are agreed on using that term from now on.

mlir/test/Dialect/Shape/shape-bufferize.mlir

4 ↗

(On Diff #293400)

typo assumign

[pifon2a]

mlir/lib/Dialect/Shape/Transforms/ShapeBufferize.cpp

27 ↗

(On Diff #293400)

Actually, that's a very good point: do we actually need to use BufferAssignmentOpConversionPattern as a base class for ops conversion when we want to go from Tensors to MemRefs if, of course, the op is not a FuncOp, CallOp or ReturnOp. Currently, i don't think there is anything special code in BufferAssignmentOpConversionPattern. In future we might add some custom logging or smth else for the buffer assignment and then it would be useful.

So, I propose to continue using BufferAssignmentOpConversionPatter but create an issue about this. Because, if it's really not needed, then we have to refactor all of the ops that did not really need to derive from BufferAssignmentOpConversionPattern.

36 ↗

(On Diff #293400)

just curious: here you keep implementation of the method within the BufferizeAssumingOpConverter definition, but for ShapeBufferizePass you implemented runOnFunction separately.

mlir/test/Dialect/Shape/shape-bufferize.mlir

1 ↗

(On Diff #293400)

--dump-input=fail is the default value
https://llvm.org/docs/CommandGuide/FileCheck.html#cmdoption-filecheck-dump-input

7 ↗

(On Diff #293400)

nit: s/@f()/@shape_assign_returns_memref()and remove the comment

18 ↗

(On Diff #293400)

super-nit: remove these 2 empty lines.

[tpopp]

Clarify the transformation of tensors to memrefs.
Leave TODOs to generalize this the rest of the way in regards to supported type conversions.

[tpopp]

mlir/lib/Dialect/Shape/Transforms/ShapeBufferize.cpp

36 ↗

(On Diff #293400)

I don't think I had any specific reason to do them differently.

[mehdi_amini]

Title: "Add file to implement bufferization for shape ops."

You meant to write "Add *pass* to ..." ?

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

24

FYI we don't usually put usernames in TODO.

mlir/lib/Dialect/Shape/Transforms/ShapeTypeConversion.cpp

2

You have a formatting issue to fix here.

85

(same)

mlir/test/Dialect/Shape/shape-type-conversion.mlir

1

Can you please rework the test to remove the allow-unregistered-dialect? This losen the verifier in a way that isn't really desirable for tests in general.

[silvas]

mlir/test/Dialect/Shape/shape-type-conversion.mlir

1

I thought that was the preferred pattern for source/sink ops? May need to revisit my source/sink op proposal? https://llvm.discourse.group/t/rfc-add-std-source-and-std-sink-for-test-case-writing-reduction/1058/16

[mehdi_amini]

mlir/test/Dialect/Shape/shape-type-conversion.mlir

1

I make a distinction between "unregistered op" and "unregistered dialect" though.
In-tree you can write a test with test.custom_shape_sink(...) for example without registering the operation and you don't need to allow unregistered dialect for this for example.

[silvas]

mlir/test/Dialect/Shape/shape-type-conversion.mlir

1

Ah, got it!

[mehdi_amini]

mlir/test/Dialect/Shape/shape-type-conversion.mlir

1

(and it wasn't needed here actually, removed in https://github.com/llvm/llvm-project/commit/5a305f81bfc3 )

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+class BufferAssignmentTypeConverter;
+} // namespace mlir
+
+namespace mlir {
 /// Creates an instance of the ShapeToShapeLowering pass that legalizes Shape
 /// dialect to be convertible to Standard. For example, `shape.num_elements` get
 /// transformed to `shape.reduce`, which can be lowered to SCF and Standard.
 std::unique_ptr<Pass> createShapeToShapeLowering();
 
 /// Collects a set of patterns to rewrite ops within the Shape dialect.
 void populateShapeRewritePatterns(MLIRContext *context,
                                   OwningRewritePatternList &patterns);
 
 // Collects a set of patterns to replace all constraints with passing witnesses.
 // This is intended to then allow all ShapeConstraint related ops and data to
 // have no effects and allow them to be freely removed such as through
 // canonicalization and dead code elimination.
 //
 // After this pass, no cstr_ operations exist.
 void populateRemoveShapeConstraintsPatterns(OwningRewritePatternList &patterns,
                                             MLIRContext *ctx);
 std::unique_ptr<FunctionPass> createRemoveShapeConstraintsPass();
 
+void populateShapeTypeConversionPatterns(
+    MLIRContext *ctx, BufferAssignmentTypeConverter *converter,
+    OwningRewritePatternList *patterns);
+// Collects a set of patterns to replace tensors as inputs and outputs to shape
+// operations with buffers. This only modifies the shape operations.
+std::unique_ptr<FunctionPass> createShapeTensorToMemrefPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Shape/Transforms/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_

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

def RemoveShapeConstraints : FunctionPass<"remove-shape-constraints"> {
   let constructor = "mlir::createRemoveShapeConstraintsPass()";
 }
 
 def ShapeToShapeLowering : FunctionPass<"shape-to-shape-lowering"> {
   let summary = "Legalize Shape dialect to be convertible to Standard";
   let constructor = "mlir::createShapeToShapeLowering()";
 }
 
+// TODO(tpopp): Generalize this to allow any type conversions desired.
+def ShapeTensorToMemref : FunctionPass<"shape-tensor-to-memref"> {
+  let summary = "Replace tensors involving shape operations with memrefs";
+  let constructor = "mlir::createShapeTensorToMemrefPass()";
+}
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES

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

 add_mlir_dialect_library(MLIRShapeOpsTransforms
   RemoveShapeConstraints.cpp
+  ShapeTypeConversion.cpp
   ShapeToShapeLowering.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ShapeOps/Transforms
 
   DEPENDS
   MLIRShapeTransformsIncGen
   )

mlir/lib/Dialect/Shape/Transforms/ShapeTypeConversion.cpp

+//=====------- ShapeTypeConversion.cpp - Shape Type Conversions ----------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines patterns to convert types of inputs and outputs to shape
+// operations to be memrefs instead of tensors.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Shape/Transforms/Passes.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/StandardTypes.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/BufferPlacement.h"
+
+using namespace mlir;
+using namespace mlir::shape;
+
+namespace {
+// Propagate tensor to memref conversions through shape.assuming ops.
+class TypeConversionAssumingOpConverter
+    : public BufferAssignmentOpConversionPattern<shape::AssumingOp> {
+public:
+  using BufferAssignmentOpConversionPattern<
+      shape::AssumingOp>::BufferAssignmentOpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(shape::AssumingOp assumingOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    SmallVector<Type, 2> newResultTypes;
+    newResultTypes.reserve(assumingOp.getNumResults());
+    for (auto result : assumingOp.getResults()) {
+      auto originalType = result.getType();
+      Type convertedType = converter->convertType(originalType);
+      newResultTypes.push_back(convertedType);
+    }
+
+    auto newAssumingOp = rewriter.create<shape::AssumingOp>(
+        assumingOp.getLoc(), newResultTypes, assumingOp.witness());
+
+    // Handle the region transfer carefully here to avoid assertions that both
+    // operations are valid at replacement time.
+    newAssumingOp.doRegion().push_back(new Block());
+    rewriter.replaceOp(assumingOp, newAssumingOp.getResults());
+    newAssumingOp.doRegion().takeBody(assumingOp.doRegion());
+
+    return success();
+  }
+};
+
+struct ShapeTensorToMemrefPass
+    : public ShapeTensorToMemrefBase<ShapeTensorToMemrefPass> {
+  void runOnFunction() override {
+    MLIRContext &ctx = getContext();
+
+    OwningRewritePatternList patterns;
+    BufferAssignmentTypeConverter converter;
+    populateShapeTypeConversionPatterns(&ctx, &converter, &patterns);
+
+    ConversionTarget target(getContext());
+    auto isMemRefType = [](Type type) { return type.isa<BaseMemRefType>(); };
+
+    target.addDynamicallyLegalOp<AssumingOp>([&](shape::AssumingOp op) {
+      return std::all_of(op.result_type_begin(), op.result_type_end(),
+                         isMemRefType);
+    });
+
+    if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+/// Populates `patterns` with the conversion patterns of tensor->memref.
+//
+// TODO(tpopp): Change this to work generally with any type conversions.
+void mlir::populateShapeTypeConversionPatterns(
+    MLIRContext *context, BufferAssignmentTypeConverter *converter,
+    OwningRewritePatternList *patterns) {
+  patterns->insert<TypeConversionAssumingOpConverter>(context, converter);
+}
+
+//===----------------------------------------------------------------------===//
+// ShapeTensorToMemrefPass construction
+//===----------------------------------------------------------------------===//
+
+std::unique_ptr<FunctionPass> mlir::createShapeTensorToMemrefPass() {
+  return std::make_unique<ShapeTensorToMemrefPass>();
+}

mlir/test/Dialect/Shape/shape-type-conversion.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -shape-tensor-to-memref <%s | FileCheck %s
+
+// -----
+// Check that shape.assuming returns a memref.
+//
+// CHECK-LABEL: @shape_assuming_returns_memref
+func @shape_assuming_returns_memref() {
+  %0 = shape.const_witness true
+  // CHECK: shape.assuming %{{.*}} -> (memref<2xf16>) {
+  %1 = shape.assuming %0 -> (tensor<2xf16>) {
+    %2 = "test.source"() : () -> (tensor<2xf16>)
+    shape.assuming_yield %2 : tensor<2xf16>
+  }
+  "test.sink"(%1) : (tensor<2xf16>) -> ()
+  return
+}
+
+