Enhance InferShapedTypeOpInterface to make it accessible during dialect conversion

[reyizero]

Summary

Original interfaces are not safe to be called during dialect conversion.
This is because some ops (e.g. dynamic_reshape(input, target_shape))
depend on the values of their operands to calculate the output shape.
However the operands may be out of reach during dialect conversion (e.g.
converting from tensor world to buffer world). This patch provides a new
kind of interface which accpets user-provided operands to solve this
problem.

[herhut]

I am +1 on the direction here.

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

159

Do we really need two versions of this function here?

I'd prefer to only have the variant with explicit operands at the interface level. What is now a default implementation of the existing interface could instead be just a helper/wrapper that can be used when the operands are the actual operands.

This has the disadvantage that we would need to refactor all implementations and users of this interface. Longer term, it has the big advantage that the interfaces cannot go out of sync.

[reyizero]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

159

Thanks! I agree it's better to have one version in the long term.

One quick question is: Where to put these helper/wrapper function? It seems more convenient if they are in the interface class (thus all subclasses no need to do wrapper again).

I don't know if it's possible to implement something like non override method inside a OpInterface?

Or we simply don't provide any wrapper function and leave it to the end users?

[jpienaar]

I think this may need a larger diff description showing exactly why an operation's operands are not appropriate but extra passed in operands are. This sounds like wanting to call the function while the IR is invalid, but if the IR is invalid then many of these functions could be wrong (and one would also need to probably avoid using accessors generated on the operation).

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

166

Bufferization is an exceptional case of usage of dialect conversion, it is the odd one out compared to most uses of dialect conversion.

[reyizero]

Add more descriptions.

[reyizero]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

166

Thanks. When the matchAndRewrite method of a conversion pattern is called, I think the op itself is still valid (e.g. attributes, output types) except that the operands may have been converted to other types, leading calling getOperand unsafely. To solve this problem, this interface follows the design of the conversion pattern, that is, accepting passed in operands to avoid calling getOperand directly inside the interface implementation.

[reyizero]

merge two interfaces into a signle one

[reyizero]

@herhut Hi, I remove the original interface and only keep the one accepting passed in operands. Could you take a look again?

[reyizero]

merge two interfaces into a signle one

[reyizero]

@jpienaar Hi, I add some additional descriptions. Could you help to take a look again?

[herhut]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

166

It is true that bufferization is special in many ways but I think the general contract in DialectConversion is that one shall not look at he ops' operands directly but only use the passed in operands arrayref.

mlir/test/lib/Dialect/Test/TestDialect.cpp

754

If you update the signature, you also need to update the inner workings of this function. it should now use operants.front() or something like that.

[reyizero]

fix

[reyizero]

mlir/test/lib/Dialect/Test/TestDialect.cpp

754

Thanks! I have fixed it.

[herhut]

There is a typo in the description. Otherwise lgtm.

[reyizero]

This comment was removed by reyizero.

[reyizero]

@herhut Hi, I don't have commit access to the llvm repo at the moment. Could you please help me to merge this patch? Thanks!

[reyizero]

@herhut @jpienaar @mehdi_amini Hi, could you help to take a look?

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

StaticInterfaceMethod<
       /*defaultImplementation=*/[{ return ::mlir::failure(); }]
     >,
     InterfaceMethod<
       /*desc=*/[{Reify the shape computation for the operation.
 
       Insert operations using the given OpBuilder that computes the
       result shape. Only one of this method or
       `reifyReturnTypeShapesPerResultDim` needs to be overriden by the
-      operation.
+      operation. This interface is supposed to be workable during dialect
+      conversion (e.g. convert from tensor world to buffer world),
+      where `getOperand` may be invalid. For example, some ops (e.g.
+      dynamic_reshape(input, target_shape)) may depend on their operands
+      to calculate the result shape. When the `matchAndRewrite ` method
+      of a conversion pattern is called, the operands of the op to convert
+      may have been converted into other types, which makes it invalid to
+      call the `getOperand` method of such op directly inside the
+      conversion pattern.  To solve this problem, this interface follows
+      the design of the conversion pattern, that is, accepting passed in
+      operands to avoid calling `getOperand` directly inside the interface
+      implementation.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"reifyReturnTypeShapes",
       /*args=*/(ins "::mlir::OpBuilder&":$builder,
+          "::mlir::ValueRange":$operands,
           "::mlir::SmallVectorImpl<Value> &":$reifiedReturnShapes),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{ return ::mlir::failure(); }]
     >,
     InterfaceMethod<
       /*desc=*/[{Reify the shape computation for the operation.
 
       Insert operations using the given OpBuilder that computes the
InterfaceMethod<
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"reifyReturnTypeShapesPerResultDim",
       /*args=*/(ins "::mlir::OpBuilder&":$builder,
           "::mlir::SmallVectorImpl<SmallVector<::mlir::Value>>&"
           :$reifiedReturnShapes),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{ return ::mlir::failure(); }]
     >
   ];
 }
 
 // Convenience class grouping together type and shaped type op interfaces for
 // ops that have tensor return types.
 class InferTensorType<list<string> overridenMethods = []> {
   list<OpTrait> traits = [
     // Op implements infer type op interface.
     InferTypeOpInterface,
     // The op will have methods implementing the ShapedType type inference
     // interface.
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface, overridenMethods>,
     // The op produces tensors and will use the ShapedType type infer interface
     // along with knowledge that it is producing Tensors to infer the type.
     NativeOpTrait<"InferTensorType">
   ];
 }
 defvar InferTensorTypeWithReify = InferTensorType<[
     "inferReturnTypeComponents", "reifyReturnTypeShapes"]>;
 
 #endif // MLIR_INFERTYPEOPINTERFACE

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

if (!shapedTypeOp)
     return nullptr;
 
   // The interface exposes two methods, one that returns the shape of all the
   // results as `Value` and other that returns the shape as a list of
   // `SmallVector<Value>`. The former takes precedence over the latter. So first
   // check if the op implements the first interface method or the second, and
   // get the value to use appropriately.
   SmallVector<Value> reifiedResultShapes;
-  if (succeeded(
-          shapedTypeOp.reifyReturnTypeShapes(builder, reifiedResultShapes))) {
+  if (succeeded(shapedTypeOp.reifyReturnTypeShapes(
+          builder, result.getOwner()->getOperands(), reifiedResultShapes))) {
     if (reifiedResultShapes.size() <= resultNumber)
       return nullptr;
     Value resultShape = reifiedResultShapes[resultNumber];
     auto resultShapeType = resultShape.getType().dyn_cast<RankedTensorType>();
     if (!resultShapeType || !resultShapeType.getElementType().isa<IndexType>())
       return nullptr;
     return builder.create<tensor::ExtractOp>(
         loc, resultShape, builder.createOrFold<ConstantIndexOp>(loc, dimIndex));

mlir/test/lib/Dialect/Test/TestDialect.cpp

LogicalResult OpWithShapedTypeInferTypeInterfaceOp::inferReturnTypeComponents(
   int64_t dim =
       sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamicSize;
   auto type = IntegerType::get(context, 17);
   inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type));
   return success();
 }
 
 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(
-    OpBuilder &builder, llvm::SmallVectorImpl<Value> &shapes) {
+    OpBuilder &builder, ValueRange operands,
+    llvm::SmallVectorImpl<Value> &shapes) {
   shapes = SmallVector<Value, 1>{
-      builder.createOrFold<memref::DimOp>(getLoc(), getOperand(0), 0)};
+      builder.createOrFold<memref::DimOp>(getLoc(), operands.front(), 0)};
   return success();
 }
 
 LogicalResult
 OpWithResultShapePerDimInterfaceOp ::reifyReturnTypeShapesPerResultDim(
     OpBuilder &builder,
     llvm::SmallVectorImpl<llvm::SmallVector<Value>> &shapes) {
   SmallVector<Value> operand1Shape, operand2Shape;

mlir/test/lib/Dialect/Test/TestPatterns.cpp

 }
 
 static void reifyReturnShape(Operation *op) {
   OpBuilder b(op);
 
   // Use permutations of 2 args as operands.
   auto shapedOp = cast<OpWithShapedTypeInferTypeInterfaceOp>(op);
   SmallVector<Value, 2> shapes;
-  if (failed(shapedOp.reifyReturnTypeShapes(b, shapes)) ||
+  if (failed(shapedOp.reifyReturnTypeShapes(b, op->getOperands(), shapes)) ||
       !llvm::hasSingleElement(shapes))
     return;
   for (auto it : llvm::enumerate(shapes)) {
     op->emitRemark() << "value " << it.index() << ": "
                      << it.value().getDefiningOp();
   }
 }