[mlir] Rename ShapeTypeConversion to ShapeBufferize

[silvas]

Summary

Once we have tensor_to_memref ops suitable for type materializations,
this pass can be split into a generic type conversion pattern.

Part of the refactor discussed in:
https://llvm.discourse.group/t/what-is-the-strategy-for-tensor-memref-conversion-bufferization/1938/17

[silvas]

update

[silvas]

arc is messing with me.... somehow crossed the streams with another commit... please hold

[silvas]

update

[silvas]

ok, fixed.

[herhut]

Thanks!

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

9

Why has this moved up?

23

Also change this name?

[silvas]

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

9

I think it is a clang-format heuristic based on the file name

23

I will do that in a follow-up. There is nothing specific to bufferization here. But without the materialization infra in place (I'll do that next), it's not possible to split it out nicely.

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

 // 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();
+// Bufferizes shape dialect ops.
+//
+// Note that most shape dialect ops must be converted to std before
+// bufferization happens, as they are intended to be bufferized at the std
+// level.
+std::unique_ptr<FunctionPass> createShapeBufferizePass();
 
 //===----------------------------------------------------------------------===//
 // 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 ShapeToShapeLowering : FunctionPass<"shape-to-shape-lowering"> {
   let summary = "Legalize Shape dialect to be convertible to Standard";
   let constructor = "mlir::createShapeToShapeLowering()";
 }
 
 // TODO: 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()";
+def ShapeBufferize : FunctionPass<"shape-bufferize"> {
+  let summary = "Bufferize the shape dialect.";
+  let constructor = "mlir::createShapeBufferizePass()";
 }
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES

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

-//====----- ShapeTypeConversion.cpp - Shape Type Conversions ----*- C++-*--===//
+//====----- Bufferize.cpp - Bufferization of shape ops  ---------*- 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 "mlir/Transforms/Bufferize.h"
 #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/Bufferize.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 {
matchAndRewrite(shape::AssumingOp assumingOp, ArrayRef<Value> operands,
     rewriter.replaceOp(assumingOp, newAssumingOp.getResults());
     rewriter.inlineRegionBefore(assumingOp.doRegion(), newAssumingOp.doRegion(),
                                 newAssumingOp.doRegion().end());
 
     return success();
   }
 };
 
-struct ShapeTensorToMemrefPass
-    : public ShapeTensorToMemrefBase<ShapeTensorToMemrefPass> {
+struct ShapeBufferizePass : public ShapeBufferizeBase<ShapeBufferizePass> {
   void runOnFunction() override {
     MLIRContext &ctx = getContext();
 
     OwningRewritePatternList patterns;
     BufferAssignmentTypeConverter converter;
     populateShapeTypeConversionPatterns(&ctx, converter, patterns);
 
     ConversionTarget target(getContext());

 // TODO: 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
+// ShapeBufferizePass construction
 //===----------------------------------------------------------------------===//
 
-std::unique_ptr<FunctionPass> mlir::createShapeTensorToMemrefPass() {
-  return std::make_unique<ShapeTensorToMemrefPass>();
+std::unique_ptr<FunctionPass> mlir::createShapeBufferizePass() {
+  return std::make_unique<ShapeBufferizePass>();
 }

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

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

mlir/test/Dialect/Shape/bufferize.mlir

-// RUN: mlir-opt -split-input-file -shape-tensor-to-memref <%s | FileCheck %s
+// RUN: mlir-opt -split-input-file -shape-bufferize <%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
 }
-
-