[mlir][shape] Split out structural type conversions for shape dialect.

[silvas]

Summary

A "structural" type conversion is one where the underlying ops are
completely agnostic to the actual types involved and simply need to update
their types. An example of this is shape.assuming -- the shape.assuming op
and the corresponding shape.assuming_yield op need to update their types
accordingly to the TypeConverter, but otherwise don't care what type
conversions are happening.

Also, the previous conversion code would not correctly materialize
conversions for the shape.assuming_yield op. This should have caused a
verification failure, but shape.assuming's verifier wasn't calling
RegionBranchOpInterface::verifyTypes (which for reasons can't be called
automatically as part of the trait verification, and requires being
called manually). This patch also adds that verification.

[silvas]

update

[jpienaar]

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

52

according ?

52

This sounds weird: they do it following X, but they don't care what they do.

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

2

nit: Shape

[silvas]

address feedback

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

52

Good catch! Updated the wording.

[tpopp]

Thanks for the explanation regarding the AssumingYieldOp also. I found it strange that it was working with that before and just assumed that ops didn't need to be updated if only their inputs changed.

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

19

If you have a strong desire to update this op in place like you did with IfOp, I wouldn't object :)

Not necessary at all, but I might follow up and do that if you don't because it looks nicer and avoids sharp edges in some DialectConversion infrastructure.

[herhut]

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

1

Which patterns do you expect to end up in this file? Like shape.shape_of tensor to memref conversion? Or would that also be structural?

[silvas]

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

1

shape.shape_of would not be structural.

Structural refers to ops that have operands/results that behave analogous to successor operands / block arguments -- one Value transfers to another Value as a pure identity, so no special knowledge of the type is needed (much as any type can be passed through to a successor block argument). That is, a structural type conversion is defined as being correct for any type conversion.

So shape.shape_of would not fit that category, since it does care about the type of its operands (such a conversion would not work, for example, for a tensor<2xi32> -> i64 conversion, as that would result in an invalid shape.shape_of op).

One observation on shape.shape_of is that its result type is independent of whether its argument is a tensor or a memref. And conversion patterns are mostly about converting results -- so as part of any bufferization flow, tensors will be converted to memrefs, and the shape.shape_of operand will naturally shift from tensor to memref (when its def gets converted, or in the "composable passes" case, the materialization back to tensor is elided as part of a final "full conversion" that eliminates the materializations). The shape.shape_of will remain valid when that happens, and, no conversion pattern is even needed.

One exception to the "conversion patterns are mostly about converting results" is e.g. "yield" patterns that are just there to trigger materializations to the target type (e.g. memref in tensor->memref conversion) so that they are consistent with the enclosing op.

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

19

I'm still parsing through the best way to do this. I think that updateRootInPlace is the best option, but the conversion infra has a bug that makes it actually not work as intended with the current state of the dialect conversion infra.

The approach as written here is actually the "best" given the current state of the conversion infrastructure. The test case would fail with updateRootInPlace, due to the tensor_cast on the shape.assuming result not getting materialized.

A workaround for that bug (which unfortunately introduces another issue), which I used in D89757 was to do "op.clone()" and then replace the op with a munged clone (rather than munging in place). However, if I did that here, the test case would fail, because the cloning process clones the "test.source" in the body, which then the conversion infra thinks it needs to legalize (as if you had rewriter.create'd it), even though we are doing a partial conversion. The approach here dodges that second issue because inlineRegionBefore doesn't create new ops, so the test.source just naturally moves into the new shape.assuming and the conversion infra doesn't feel compelled to legalize it. There might be a conversion infrastructure bug here or it might be working as intended.

I'll be following up with River to get the conversion infra fixed and update these tests to a single idiomatic way of doing it (which I think will be the updateRootInPlace approach).

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

let description = [{
     ready to execute.
   }];
   let arguments = (ins Shape_WitnessType:$witness);
   let regions = (region SizedRegion<1>:$doRegion);
   let results = (outs Variadic<AnyType>:$results);
 
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
+  let verifier = [{ return RegionBranchOpInterface::verifyTypes(*this); }];
 
   let extraClassDeclaration = [{
     // Inline the region into the region containing the AssumingOp and delete
     // the AssumingOp.
     //
     // This does no checks on the inputs to the AssumingOp.
     static void inlineRegionIntoParent(AssumingOp &op, PatternRewriter &rewriter);
   }];

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 BufferizeTypeConverter;
+class ConversionTarget;
+class TypeConverter;
 } // 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,
-                                         BufferizeTypeConverter &converter,
-                                         OwningRewritePatternList &patterns);
+/// Populates patterns for shape dialect structural type conversions and sets up
+/// the provided ConversionTarget with the appropriate legality configuration
+/// for the ops to get converted properly.
+///
+/// A "structural" type conversion is one where the underlying ops are
+/// completely agnostic to the actual types involved and simply need to update
+/// their types consistently. An example of this is shape.assuming -- the
+/// shape.assuming op and the corresponding shape.assuming_yield op need to have
+/// consistent types, but the exact types don't matter. So all that we need to
+/// do for a structural type conversion is to update both of their types
+/// consistently to the new types prescribed by the TypeConverter.
+void populateShapeStructuralTypeConversionsAndLegality(
+    MLIRContext *context, TypeConverter &typeConverter,
+    OwningRewritePatternList &patterns, ConversionTarget &target);
+
 // 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();
 
 //===----------------------------------------------------------------------===//

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

 //====----- 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
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 
 using namespace mlir;
-using namespace mlir::shape;
 
 namespace {
-// Propagate tensor to memref conversions through shape.assuming ops.
-class TypeConversionAssumingOpConverter
-    : public BufferizeOpConversionPattern<shape::AssumingOp> {
-public:
-  using BufferizeOpConversionPattern<
-      shape::AssumingOp>::BufferizeOpConversionPattern;
-
-  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());
-
-    rewriter.replaceOp(assumingOp, newAssumingOp.getResults());
-    rewriter.inlineRegionBefore(assumingOp.doRegion(), newAssumingOp.doRegion(),
-                                newAssumingOp.doRegion().end());
-
-    return success();
-  }
-};
-
 struct ShapeBufferizePass : public ShapeBufferizeBase<ShapeBufferizePass> {
   void runOnFunction() override {
     MLIRContext &ctx = getContext();
 
     OwningRewritePatternList patterns;
-    BufferizeTypeConverter converter;
-    populateShapeTypeConversionPatterns(&ctx, converter, patterns);
-
+    BufferizeTypeConverter typeConverter;
     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);
-    });
+    populateBufferizeMaterializationLegality(target);
+    populateShapeStructuralTypeConversionsAndLegality(&ctx, typeConverter,
+                                                      patterns, target);
 
-    if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
+    if (failed(applyPartialConversion(getFunction(), target, patterns)))
       signalPassFailure();
   }
 };
-
 } // namespace
 
-/// Populates `patterns` with the conversion patterns of tensor->memref.
-//
-// TODO: Change this to work generally with any type conversions.
-void mlir::populateShapeTypeConversionPatterns(
-    MLIRContext *context, BufferizeTypeConverter &converter,
-    OwningRewritePatternList &patterns) {
-  patterns.insert<TypeConversionAssumingOpConverter>(context, converter);
-}
-
-//===----------------------------------------------------------------------===//
-// ShapeBufferizePass construction
-//===----------------------------------------------------------------------===//
-
 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
   ShapeToShapeLowering.cpp
+  StructuralTypeConversions.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ShapeOps/Transforms
 
   DEPENDS
   MLIRShapeTransformsIncGen
   )
 
 target_link_libraries(MLIRShapeOpsTransforms
   PUBLIC
   MLIRIR
   MLIRPass
   MLIRShape
   MLIRSupport
   MLIRTransforms
   )

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

+//===- StructuralTypeConversions.cpp - Shape structural type conversions --===//
+//
+// 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 "PassDetail.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Shape/Transforms/Passes.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+using namespace mlir::shape;
+
+namespace {
+class ConvertAssumingOpTypes : public OpConversionPattern<AssumingOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(AssumingOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    SmallVector<Type, 2> newResultTypes;
+    newResultTypes.reserve(op.getNumResults());
+    for (auto result : op.getResults()) {
+      auto originalType = result.getType();
+      Type convertedType = getTypeConverter()->convertType(originalType);
+      newResultTypes.push_back(convertedType);
+    }
+
+    auto newAssumingOp =
+        rewriter.create<AssumingOp>(op.getLoc(), newResultTypes, op.witness());
+
+    rewriter.replaceOp(op, newAssumingOp.getResults());
+    rewriter.inlineRegionBefore(op.doRegion(), newAssumingOp.doRegion(),
+                                newAssumingOp.doRegion().end());
+
+    return success();
+  }
+};
+} // namespace
+
+namespace {
+class ConvertAssumingYieldOpTypes
+    : public OpConversionPattern<AssumingYieldOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(AssumingYieldOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    rewriter.replaceOpWithNewOp<AssumingYieldOp>(op, operands);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateShapeStructuralTypeConversionsAndLegality(
+    MLIRContext *context, TypeConverter &typeConverter,
+    OwningRewritePatternList &patterns, ConversionTarget &target) {
+  patterns.insert<ConvertAssumingOpTypes, ConvertAssumingYieldOpTypes>(
+      typeConverter, context);
+  target.addDynamicallyLegalOp<AssumingOp>([&](AssumingOp op) {
+    return typeConverter.isLegal(op.getResultTypes());
+  });
+  target.addDynamicallyLegalOp<AssumingYieldOp>([&](AssumingYieldOp op) {
+    return typeConverter.isLegal(op.getOperandTypes());
+  });
+}

mlir/test/Dialect/Shape/bufferize.mlir

 // 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() {
+
+// CHECK-LABEL:   func @shape_assuming() {
+// CHECK:           %[[WTRUE:.*]] = shape.const_witness true
+// CHECK:           %[[MEMREF:.*]] = shape.assuming %[[WTRUE]] -> (memref<2xf16>) {
+// CHECK:             %[[TENSOR_VAL:.*]] = "test.source"() : () -> tensor<2xf16>
+// CHECK:             %[[YIELDED_MEMREF:.*]] = tensor_to_memref %[[TENSOR_VAL]] : memref<2xf16>
+// CHECK:             shape.assuming_yield %[[YIELDED_MEMREF]] : memref<2xf16>
+// CHECK:           }
+// CHECK:           %[[TENSOR:.*]] = tensor_load %[[MEMREF:.*]] : memref<2xf16>
+// CHECK:           "test.sink"(%[[TENSOR]]) : (tensor<2xf16>) -> ()
+// CHECK:           return
+// CHECK:         }
+func @shape_assuming() {
   %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
 }