Fix parsing of hex-format index dense tensor attributes.

[benvanik]

Summary

TensorLiteralParser::getHexAttr does a isIntOrIndexOrFloat check and properly handles index elements, but TensorLiteralParser::getAttr that calls into it has a mismatched check. This just makes the checks match so that index element attrs can parse when of type tensor.

[stellaraccident]

Maybe a test somewhere around here: https://github.com/llvm/llvm-project/blob/5f4c91583ee772a6ce2c4f192e25b07e6075eb00/mlir/test/IR/attribute.mlir#L444 ?

[rriddle]

LG after adding test.

[benvanik]

Added a test - there was nothing in this file testing the hex format at all, so I just used what I knew would trigger that path and round-trip as hex.

mlir/lib/Parser/AttributeParser.cpp

 /// Build a dense attribute instance with the parsed elements and the given
 /// shaped type.
 DenseElementsAttr TensorLiteralParser::getAttr(llvm::SMLoc loc,
                                                ShapedType type) {
   Type eltType = type.getElementType();
 
   // Check to see if we parse the literal from a hex string.
   if (hexStorage.hasValue() &&
-      (eltType.isIntOrFloat() || eltType.isa<ComplexType>()))
+      (eltType.isIntOrIndexOrFloat() || eltType.isa<ComplexType>()))
     return getHexAttr(loc, type);
 
   // Check that the parsed storage size has the same number of elements to the
   // type, or is a known splat.
   if (!shape.empty() && getShape() != type.getShape()) {
     p.emitError(loc) << "inferred shape of elements literal ([" << getShape()
                      << "]) does not match type ([" << type.getShape() << "])";
     return nullptr;

mlir/test/IR/attribute.mlir

func @correct_int_elements_attr_pass() {
   } : () -> ()
 
   "test.int_elements_attr"() {
     // CHECK: any_i32_attr = dense<5> : tensor<1x2x3x4xi32>,
     any_i32_attr = dense<5> : tensor<1x2x3x4xi32>,
     i32_attr = dense<5> : tensor<6xi32>
   } : () -> ()
 
+  "test.index_elements_attr"() {
+    // CHECK: any_index_attr = dense<5> : tensor<1x2x3x4xindex>,
+    any_index_attr = dense<5> : tensor<1x2x3x4xindex>,
+    index_attr = dense<5> : tensor<6xindex>
+  } : () -> ()
+
+  "test.hex_index_elements_attr"() {
+    // CHECK: hex_index_attr = dense<"0x00000C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000080C0000000000001A150000000000005078000000000000B659010000000000ECBC010000000000FEC5010000000000342902000000000046320200000000007C950200000000008E9E020000000000C401030000000000D60A0300000000000C6E0300000000001E7703000000000054DA03000000000066E30300000000009C46040000000000AE4F040000000000E4B2040000000000F6BB0400000000002C1F050000000000628100000000000098E40000000000000E0C00000000000020150000000000005678000000000000BC59010000000000F2BC01000000000004C60100000000003A290200000000004C320200000000008295020000000000949E020000000000CA01030000000000DC0A030000000000126E03000000000024770300000000005ADA0300000000006CE3030000000000A246040000000000B44F040000000000EAB2040000000000FCBB040000000000321F05000000000068810000000000009EE40000000000"> : tensor<23x5xindex>
+    hex_index_attr = dense<"0x00000C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000080C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000100C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000180C000000000000080C0000000000001A150000000000005078000000000000B659010000000000ECBC010000000000FEC5010000000000342902000000000046320200000000007C950200000000008E9E020000000000C401030000000000D60A0300000000000C6E0300000000001E7703000000000054DA03000000000066E30300000000009C46040000000000AE4F040000000000E4B2040000000000F6BB0400000000002C1F050000000000628100000000000098E40000000000000E0C00000000000020150000000000005678000000000000BC59010000000000F2BC01000000000004C60100000000003A290200000000004C320200000000008295020000000000949E020000000000CA01030000000000DC0A030000000000126E03000000000024770300000000005ADA0300000000006CE3030000000000A246040000000000B44F040000000000EAB2040000000000FCBB040000000000321F05000000000068810000000000009EE40000000000"> : tensor<23x5xindex>
+  } : () -> ()
+
   return
 }
 
 // -----
 
 func @wrong_int_elements_attr_type_fail() {
   // expected-error @+1 {{'any_i32_attr' failed to satisfy constraint: 32-bit integer elements attribute}}
   "test.int_elements_attr"() {