[MLIR] Allow ShapedTypeComponents with attribute in inferReturnTensorTypes

[Smit]

Summary

Originally, inferReturnTensorTypes didn't support shaped type components
containing an attribute just because there wasn't any motivating use-case.
Removing that limitation and using it to set the encoding attribute for
RankedTensorType.

Updated the existing test to set result attribute based on the first operand,
if available.

Signed-off-by: Smit Hinsu <[e-mail]>

[Smit]

Allow ShapedTypeComponents with attribute in inferReturnTensorTypes

[jpienaar]

Looks good in general (I'll land in a bit) with regular caveat that encodings are open and downstream users may need to be more restrictive.

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

1144

Elide single line braces.

[Smit]

Removing single line braces

[Smit]

Rebased

[Smit]

Merged with original commit

mlir/lib/Interfaces/InferTypeOpInterface.cpp

LogicalResult mlir::detail::inferReturnTensorTypes(
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   SmallVector<ShapedTypeComponents, 2> retComponents;
   if (failed(componentTypeFn(context, location, operands, attributes, regions,
                              retComponents)))
     return failure();
   for (const auto &shapeAndType : retComponents) {
-    assert(shapeAndType.getAttribute() == nullptr && "attribute not supported");
-    assert(shapeAndType.getElementType() &&
-           "element type required to construct tensor");
-    if (shapeAndType.hasRank())
-      inferredReturnTypes.push_back(RankedTensorType::get(
-          shapeAndType.getDims(), shapeAndType.getElementType()));
-    else
+    Type element_ty = shapeAndType.getElementType();
+    assert(element_ty && "element type required to construct tensor");
+
+    Attribute attr = shapeAndType.getAttribute();
+    if (shapeAndType.hasRank()) {
       inferredReturnTypes.push_back(
-          UnrankedTensorType::get(shapeAndType.getElementType()));
+          RankedTensorType::get(shapeAndType.getDims(), element_ty, attr));
+    } else {
+      assert(attr == nullptr && "attribute not supported");
+      inferredReturnTypes.push_back(UnrankedTensorType::get(element_ty));
+    }
   }
   return success();
 }
 
 LogicalResult mlir::detail::verifyInferredResultTypes(Operation *op) {
   SmallVector<Type, 4> inferredReturnTypes(op->getResultTypes());
   auto retTypeFn = cast<InferTypeOpInterface>(op);
   return retTypeFn.refineReturnTypes(op->getContext(), op->getLoc(),
                                      op->getOperands(), op->getAttrDictionary(),
                                      op->getRegions(), inferredReturnTypes);
 }

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

LogicalResult OpWithShapedTypeInferTypeInterfaceOp::inferReturnTypeComponents(
   auto operandType = operands.front().getType();
   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;
   auto type = IntegerType::get(context, 17);
-  inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type));
+
+  Attribute encoding;
+  if (auto ranked_ty = sval.dyn_cast<RankedTensorType>()) {
+    encoding = ranked_ty.getEncoding();
+  }
+  inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type, encoding));
   return success();
 }
 
 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(
     OpBuilder &builder, ValueRange operands,
     llvm::SmallVectorImpl<Value> &shapes) {
   shapes = SmallVector<Value, 1>{
       builder.createOrFold<tensor::DimOp>(getLoc(), operands.front(), 0)};

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

 // RUN: mlir-opt %s -test-return-type -split-input-file -verify-diagnostics | FileCheck %s
 
 // 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.func @testCreateFunctions(%arg0 : tensor<10xf32>, %arg1 : tensor<20xi32>) {
+func.func @testCreateFunctions(%arg0 : tensor<10xf32, !test.smpla>, %arg1 : tensor<20xi32>) {
 // CHECK: "test.no_attributes"
-  %good = "test.no_attributes"(%arg0, %arg0) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
+  %good = "test.no_attributes"(%arg0, %arg0) : (tensor<10xf32, !test.smpla>, tensor<10xf32, !test.smpla>) -> tensor<10xf32, !test.smpla>
 // CHECK: "test.op_with_shaped_type_infer_type_if"
-// CHECK-SAME: (tensor<10xf32>, tensor<10xf32>) -> tensor<10xi17>
+// CHECK-SAME: (tensor<10xf32, !test.smpla>, tensor<10xf32, !test.smpla>) -> tensor<10xi17, !test.smpla>
 // CHECK: "test.op_with_shaped_type_infer_type_if"
-// CHECK-SAME: (tensor<10xf32>, tensor<20xi32>) -> tensor<10xi17>
+// CHECK-SAME: (tensor<10xf32, !test.smpla>, tensor<20xi32>) -> tensor<10xi17, !test.smpla>
 // CHECK: "test.op_with_shaped_type_infer_type_if"
-// CHECK-SAME: (tensor<20xi32>, tensor<10xf32>) -> tensor<20xi17>
+// CHECK-SAME: (tensor<20xi32>, tensor<10xf32, !test.smpla>) -> tensor<20xi17>
 // CHECK: "test.op_with_shaped_type_infer_type_if"
 // CHECK-SAME: (tensor<20xi32>, tensor<20xi32>) -> tensor<20xi17>
 // CHECK: "test.op_with_infer_type_if"
-// CHECK-SAME: (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
+// CHECK-SAME: (tensor<10xf32, !test.smpla>, tensor<10xf32, !test.smpla>) -> tensor<10xf32, !test.smpla>
 // CHECK: "test.op_with_infer_type_if"
 // CHECK-SAME: (tensor<20xi32>, tensor<20xi32>) -> tensor<20xi32>
   return
 }
 
 // -----
 
 func.func @testReturnTypeOpInterface(%arg0 : tensor<10xf32>) {