[MLIR][Linalg] Handle Attribute in InitTensorOp

[chelini]

Summary

In some cases, the result of an initTensorOp may have an attribute.
However, the Attribute was not passed to inferResultType, failing the
verifier. Therefore, propagate the Attribute to inferResultType.

[mehdi_amini]

Can you add a test please?

[chelini]

Add test case

[chelini]

In D117192#3239756, @mehdi_amini wrote:

Can you add a test please?

Thanks, forgot to git add it.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

static StringRef getStaticSizesAttrName() {
       return "static_sizes";
     }
 
     RankedTensorType getType() {
       return getResult().getType().cast<RankedTensorType>(); }
 
     // Infer the shape of the result tensor given the static shapes
     // and element type of the result tensor.
-    static Type inferResultType(ArrayRef<int64_t> staticSizes, Type elementType);
+    static Type inferResultType(ArrayRef<int64_t> staticSizes, Type elementType,
+                                Attribute encoding = {});
 
     // Return true if the size of the tensor is dynamic at `idx`
     bool isDynamicSize(unsigned idx) {
       APInt v = *(static_sizes().getAsValueRange<IntegerAttr>().begin() + idx);
       return ShapedType::isDynamic(v.getSExtValue());
     }
 
     // Assert that the size of the result tensor is static at `idx`

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

if (failed(verifyListOfOperandsOrIntegers(op, "sizes", resultType.getRank(),
                                             op.static_sizes(), op.sizes(),
                                             ShapedType::isDynamic)))
     return failure();
 
   if (op.static_sizes().size() != static_cast<unsigned>(resultType.getRank()))
     return op->emitError("expected ")
            << resultType.getRank() << " sizes values";
 
-  Type expectedType =
-      InitTensorOp::inferResultType(staticSizes, resultType.getElementType());
+  Type expectedType = InitTensorOp::inferResultType(
+      staticSizes, resultType.getElementType(), resultType.getEncoding());
   if (resultType != expectedType) {
     return op.emitError("specified type ")
            << resultType << " does not match the inferred type "
            << expectedType;
   }
   return success();
 }
 
 Type InitTensorOp::inferResultType(ArrayRef<int64_t> staticSizes,
-                                   Type elementType) {
-  return RankedTensorType::get(staticSizes, elementType);
+                                   Type elementType, Attribute encoding) {
+  return RankedTensorType::get(staticSizes, elementType, encoding);
 }
 
 namespace {
 /// Change the type of the result of a `linalg.init_tensor` by making the result
 /// type statically sized along dimension that in the original operation where
 /// defined as dynamic, but the size was defined using a `constant` op. For
 /// example
 ///

mlir/test/Dialect/Linalg/roundtrip.mlir

 // CHECK-LABEL: func @named_ops
 //       CHECK:   linalg.batch_matmul
 //       CHECK:   linalg.batch_matmul
 //       CHECK:   linalg.batch_matmul
 //       CHECK:   linalg.batch_matmul
 
 // -----
 
+#attr = {"foo"}
 func @init_tensor(%arg0 : index, %arg1 : index)
 {
   %0 = linalg.init_tensor [3, 42] : tensor<3x42xf32>
   %1 = linalg.init_tensor [4, %arg0, %arg1, 5] : tensor<4x?x?x5xf32>
+  %2 = linalg.init_tensor [2, 2] : tensor<2x2xf32, #attr>
   return
 }
 // CHECK-LABEL: func @init_tensor
 //       CHECK:   linalg.init_tensor [3, 42] : tensor<3x42xf32>
 //       CHECK:   linalg.init_tensor [4, %{{.*}}, %{{.*}}, 5] : tensor<4x?x?x5xf32>
+//       CHECK:   linalg.init_tensor [2, 2] : tensor<2x2xf32, {foo}>
 
 // -----
 
 func @fill_tensor(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> {
   %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
   %1 = linalg.fill(%arg2, %0) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }