[mlir] Add reifyReturnShape to shaped type OpInterface

[jpienaar]

Summary

This call results in inserting operations that compute the return shape
dynamically for the operation.

[pifon2a]

mlir/include/mlir/Analysis/InferTypeOpInterface.td

108

nit:

/*methodBody=*/[{}],
/*defaultImplementation=*/[{ return failure(); }]

mlir/test/lib/TestDialect/TestDialect.cpp

322

nit: maybe smth like:

auto sval = operandType.dyn_cast<ShapedType>();
if (!sval) {
    return emitOptionalError(location, "only shaped type operands allowed");
}
int64_t dim = sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamicSize;

323–324

just out of curiosity: is reserve() really needed for 1 element?

325

why 17?

333

will smth like that work?

shapes = {builder.createOrFold<mlir::DimOp>(getLoc(), getOperand(0), 0)};

[jpienaar]

Address review comments

[jpienaar]

mlir/test/lib/TestDialect/TestDialect.cpp

323–324

Not really, force of habit :)

325

To make it less likely that it accidentally happened :) No real reasion, completely arbitrary but I was using standard bitwidths as inputs and ensure it does get built.

[jpienaar]

Landed in e706533f0a52b2eb8929d4004d1daccd53feda29, removed the review line by accident.

mlir/include/mlir/Analysis/InferTypeOpInterface.h

 // `InferTypeOpInterface.td`.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_INFERTYPEOPINTERFACE_H_
 #define MLIR_ANALYSIS_INFERTYPEOPINTERFACE_H_
 
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 
 /// ShapedTypeComponents that represents the components of a ShapedType.

mlir/include/mlir/Analysis/InferTypeOpInterface.td

StaticInterfaceMethod<
       /*args=*/(ins "MLIRContext*":$context,
                     "Optional<Location>":$location,
                     "ValueRange":$operands,
                     "ArrayRef<NamedAttribute>":$attributes,
                     "RegionRange":$regions,
                     "SmallVectorImpl<ShapedTypeComponents>&":
                       $inferedReturnShapes)
     >,
+    InterfaceMethod<
+      /*desc=*/[{Reify the shape computation for the operation.
+
+      Insert operations using the given OpBulder that computes the result shape.
+      }],
+      /*retTy=*/"LogicalResult",
+      /*methodName=*/"reifyReturnTypeShapes",
+      /*args=*/(ins "OpBuilder&":$builder,
+                    "SmallVectorImpl<Value>&":$reifiedReturnShapes),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{ return failure(); }]
+    >,
   ];
 }
 
 #endif // MLIR_INFERTYPEOPINTERFACE

mlir/test/lib/TestDialect/TestDialect.cpp

 //===- TestDialect.cpp - MLIR Dialect for Testing -------------------------===//
 //
 // 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 "TestDialect.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/StringSwitch.h"
 
LogicalResult mlir::OpWithInferTypeInterfaceOp::inferReturnTypes(
   }
   inferedReturnTypes.assign({operands[0].getType()});
   return success();
 }
 
 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::inferReturnTypeComponents(
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     ArrayRef<NamedAttribute> attributes, RegionRange regions,
-    SmallVectorImpl<ShapedTypeComponents> &inferedComponents) {
-  // Create return type consisting of the first element of each shape of the
-  // input operands or unknown for unranked operand.
-  std::vector<int64_t> shape;
-  shape.reserve(operands.size());
-  for (auto operandType : operands.getTypes()) {
-    if (auto sval = operandType.dyn_cast<ShapedType>()) {
-      if (sval.hasRank())
-        shape.push_back(sval.getShape().front());
-      else
-        shape.push_back(ShapedType::kDynamicSize);
-    } else {
-      return emitOptionalError(location, "only shaped type operands allowed");
-    }
-  }
-  inferedComponents.reserve(1);
+    SmallVectorImpl<ShapedTypeComponents> &inferedReturnShapes) {
+  // Create return type consisting of the last element of the first operand.
+  auto operandType = *operands.getTypes().begin();
+  auto sval = operandType.dyn_cast<ShapedType>();
+  if (!sval) {
+  int64_t dim =
+      sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamicSize;
   auto type = IntegerType::get(17, context);
-  inferedComponents.emplace_back(shape, type);
+  inferedReturnShapes.push_back(ShapedTypeComponents({dim}, type));
+  return success();
+}
+
+LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(
+    OpBuilder &builder, llvm::SmallVectorImpl<Value> &shapes) {
+  shapes = SmallVector<Value, 1>{
+      builder.createOrFold<mlir::DimOp>(getLoc(), getOperand(0), 0)};
   return success();
 }
 
 // Static initialization for Test dialect registration.
 static mlir::DialectRegistration<mlir::TestDialect> testDialect;
 
 #include "TestOpEnums.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

mlir/test/lib/TestDialect/TestPatterns.cpp

for (int j = 0; j < e; ++j) {
         // TODO(jpienaar): Expand to regions.
         OpTy::build(&b, state, values, op->getAttrs());
         (void)b.createOperation(state);
       }
     }
   }
 }
 
+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)))
+    return;
+  for (auto it : llvm::enumerate(shapes))
+    op->emitRemark() << "value " << it.index() << ": "
+                     << it.value().getDefiningOp();
+}
+
 struct TestReturnTypeDriver : public FunctionPass<TestReturnTypeDriver> {
   void runOnFunction() override {
     if (getFunction().getName() == "testCreateFunctions") {
       std::vector<Operation *> ops;
       // Collect ops to avoid triggering on inserted ops.
       for (auto &op : getFunction().getBody().front())
         ops.push_back(&op);
       // Generate test patterns for each, but skip terminator.
       for (auto *op : llvm::makeArrayRef(ops).drop_back()) {
         // Test create method of each of the Op classes below. The resultant
         // output would be in reverse order underneath `op` from which
         // the attributes and regions are used.
         invokeCreateWithInferedReturnType<OpWithInferTypeInterfaceOp>(op);
         invokeCreateWithInferedReturnType<OpWithShapedTypeInferTypeInterfaceOp>(
             op);
       };
       return;
     }
+    if (getFunction().getName() == "testReifyFunctions") {
+      std::vector<Operation *> ops;
+      // Collect ops to avoid triggering on inserted ops.
+      for (auto &op : getFunction().getBody().front())
+        if (isa<OpWithShapedTypeInferTypeInterfaceOp>(op))
+          ops.push_back(&op);
+      // Generate test patterns for each, but skip terminator.
+      for (auto *op : ops)
+        reifyReturnShape(op);
+    }
   }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // Legalization Driver.
 //===----------------------------------------------------------------------===//
 

mlir/test/mlir-tblgen/return-types.mlir

 // RUN: mlir-opt %s -test-return-type -split-input-file -verify-diagnostics | FileCheck %s --dump-input-on-failure
 
 // CHECK-LABEL: testCreateFunctions
 // This function tests invoking the create method with different inference
 // methods. The attributes of the ops inside are used to test creation.
 func @testCreateFunctions(%arg0 : tensor<10xf32>, %arg1 : tensor<20xi32>) {
 // CHECK: "test.no_attributes"
   %good = "test.no_attributes"(%arg0, %arg0) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
 // CHECK: "test.op_with_shaped_type_infer_type_if"
-// CHECK-SAME: (tensor<10xf32>, tensor<10xf32>) -> tensor<10x10xi17>
+// CHECK-SAME: (tensor<10xf32>, tensor<10xf32>) -> tensor<10xi17>
 // CHECK: "test.op_with_shaped_type_infer_type_if"
-// CHECK-SAME: (tensor<10xf32>, tensor<20xi32>) -> tensor<10x20xi17>
+// CHECK-SAME: (tensor<10xf32>, tensor<20xi32>) -> tensor<10xi17>
 // CHECK: "test.op_with_shaped_type_infer_type_if"
-// CHECK-SAME: (tensor<20xi32>, tensor<10xf32>) -> tensor<20x10xi17>
+// CHECK-SAME: (tensor<20xi32>, tensor<10xf32>) -> tensor<20xi17>
 // CHECK: "test.op_with_shaped_type_infer_type_if"
-// CHECK-SAME: (tensor<20xi32>, tensor<20xi32>) -> tensor<20x20xi17>
+// CHECK-SAME: (tensor<20xi32>, tensor<20xi32>) -> tensor<20xi17>
 // CHECK: "test.op_with_infer_type_if"
 // CHECK-SAME: (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
 // CHECK: "test.op_with_infer_type_if"
 // CHECK-SAME: (tensor<20xi32>, tensor<20xi32>) -> tensor<20xi32>
   return
 }
 
 // -----
 
 func @testReturnTypeOpInterface(%arg0 : tensor<10xf32>) {
   // expected-error@+1 {{incompatible with return type}}
   %bad = "test.op_with_infer_type_if"(%arg0, %arg0) : (tensor<10xf32>, tensor<10xf32>) -> tensor<*xf32>
   return
 }
 
 // -----
 
 func @testReturnTypeOpInterfaceMismatch(%arg0 : tensor<10xf32>, %arg1 : tensor<20xf32>) {
   // expected-error@+1 {{operand type mismatch}}
   %bad = "test.op_with_infer_type_if"(%arg0, %arg1) : (tensor<10xf32>, tensor<20xf32>) -> tensor<*xf32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: testReifyFunctions
+func @testReifyFunctions(%arg0 : tensor<10xf32>, %arg1 : tensor<20xf32>) {
+  // expected-remark@+1 {{constant 10}}
+  %0 = "test.op_with_shaped_type_infer_type_if"(%arg0, %arg1) : (tensor<10xf32>, tensor<20xf32>) -> tensor<10xi17>
+  // expected-remark@+1 {{constant 20}}
+  %1 = "test.op_with_shaped_type_infer_type_if"(%arg1, %arg0) : (tensor<20xf32>, tensor<10xf32>) -> tensor<20xi17>
+  return
+}