[mlir] verify that operand/result_segment_sizes attributes have i32 element

[ftynse]

Summary

This is an assumption that is made in numerous places in the code. In
particular, in the code generated by mlir-tblgen for operand/result accessors
in ops with attr-sized operand or result lists. Make sure to verify this
assumption.

Note that the operation traits are verified before running the custom op
verifier, which can expect the trait verifier to have passed, but some traits
may be verified before the AttrSizedOperand/ResultTrait and should not make
such assumptions.

mlir/lib/IR/Operation.cpp

 
 static LogicalResult verifyValueSizeAttr(Operation *op, StringRef attrName,
                                          bool isOperand) {
   auto sizeAttr = op->getAttrOfType<DenseIntElementsAttr>(attrName);
   if (!sizeAttr)
     return op->emitOpError("requires 1D vector attribute '") << attrName << "'";
 
   auto sizeAttrType = sizeAttr.getType().dyn_cast<VectorType>();
-  if (!sizeAttrType || sizeAttrType.getRank() != 1)
-    return op->emitOpError("requires 1D vector attribute '") << attrName << "'";
+  if (!sizeAttrType || sizeAttrType.getRank() != 1 ||
+      !sizeAttrType.getElementType().isInteger(32))
+    return op->emitOpError("requires 1D vector of i32 attribute '")
+           << attrName << "'";
 
   if (llvm::any_of(sizeAttr.getIntValues(), [](const APInt &element) {
         return !element.isNonNegative();
       }))
     return op->emitOpError("'")
            << attrName << "' attribute cannot have negative elements";
 
   size_t totalCount = std::accumulate(

mlir/test/IR/traits.mlir

 func @failedMissingOperandSizeAttr(%arg: i32) {
   // expected-error @+1 {{requires 1D vector attribute 'operand_segment_sizes'}}
   "test.attr_sized_operands"(%arg, %arg, %arg, %arg) : (i32, i32, i32, i32) -> ()
 }
 
 // -----
 
 func @failedOperandSizeAttrWrongType(%arg: i32) {
-  // expected-error @+1 {{requires 1D vector attribute 'operand_segment_sizes'}}
+  // expected-error @+1 {{requires 1D vector of i32 attribute 'operand_segment_sizes'}}
   "test.attr_sized_operands"(%arg, %arg, %arg, %arg) {operand_segment_sizes = dense<[1, 1, 1, 1]>: tensor<4xi32>} : (i32, i32, i32, i32) -> ()
 }
 
 // -----
 
 func @failedOperandSizeAttrWrongRank(%arg: i32) {
-  // expected-error @+1 {{requires 1D vector attribute 'operand_segment_sizes'}}
+  // expected-error @+1 {{requires 1D vector of i32 attribute 'operand_segment_sizes'}}
   "test.attr_sized_operands"(%arg, %arg, %arg, %arg) {operand_segment_sizes = dense<[[1, 1], [1, 1]]>: vector<2x2xi32>} : (i32, i32, i32, i32) -> ()
 }
 
 // -----
 
+func @failedOperandSizeAttrWrongElementType(%arg: i32) {
+  // expected-error @+1 {{requires 1D vector of i32 attribute 'operand_segment_sizes'}}
+  "test.attr_sized_operands"(%arg, %arg, %arg, %arg) {operand_segment_sizes = dense<[1, 1, 1, 1]>: vector<4xi64>} : (i32, i32, i32, i32) -> ()
+}
+
+// -----
+
 func @failedOperandSizeAttrNegativeValue(%arg: i32) {
   // expected-error @+1 {{'operand_segment_sizes' attribute cannot have negative elements}}
   "test.attr_sized_operands"(%arg, %arg, %arg, %arg) {operand_segment_sizes = dense<[1, 1, -1, 1]>: vector<4xi32>} : (i32, i32, i32, i32) -> ()
 }
 
 // -----
 
 func @failedOperandSizeAttrWrongTotalSize(%arg: i32) {

 func @failedMissingResultSizeAttr() {
   // expected-error @+1 {{requires 1D vector attribute 'result_segment_sizes'}}
   %0:4 = "test.attr_sized_results"() : () -> (i32, i32, i32, i32)
 }
 
 // -----
 
 func @failedResultSizeAttrWrongType() {
-  // expected-error @+1 {{requires 1D vector attribute 'result_segment_sizes'}}
+  // expected-error @+1 {{requires 1D vector of i32 attribute 'result_segment_sizes'}}
   %0:4 = "test.attr_sized_results"() {result_segment_sizes = dense<[1, 1, 1, 1]>: tensor<4xi32>} : () -> (i32, i32, i32, i32)
 }
 
 // -----
 
 func @failedResultSizeAttrWrongRank() {
-  // expected-error @+1 {{requires 1D vector attribute 'result_segment_sizes'}}
+  // expected-error @+1 {{requires 1D vector of i32 attribute 'result_segment_sizes'}}
   %0:4 = "test.attr_sized_results"() {result_segment_sizes = dense<[[1, 1], [1, 1]]>: vector<2x2xi32>} : () -> (i32, i32, i32, i32)
 }
 
 // -----
 
+func @failedResultSizeAttrWrongElementType() {
+  // expected-error @+1 {{requires 1D vector of i32 attribute 'result_segment_sizes'}}
+  %0:4 = "test.attr_sized_results"() {result_segment_sizes = dense<[1, 1, 1, 1]>: vector<4xi64>} : () -> (i32, i32, i32, i32)
+}
+
+// -----
+
 func @failedResultSizeAttrNegativeValue() {
   // expected-error @+1 {{'result_segment_sizes' attribute cannot have negative elements}}
   %0:4 = "test.attr_sized_results"() {result_segment_sizes = dense<[1, 1, -1, 1]>: vector<4xi32>} : () -> (i32, i32, i32, i32)
 }
 
 // -----
 
 func @failedResultSizeAttrWrongTotalSize() {