[mlir][shape] Use IndexElementsAttr in Shape dialect.

[silvas]

Summary

Index is the proper type for storing shapes when constant folding, so
this fixes the previous code (which was using i64).

[herhut]

Much nicer indeed!

[jpienaar]

Thanks, good cleanup

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

let description = [{
     Creates a !shape.shape with rank given by the length of `shape` and with
     dimension sizes given by the values of `shape`.
 
     ```mlir
     %0 = shape.const_shape []
     %1 = shape.const_shape [1, 2, 3]
     ```
   }];
-  let arguments = (ins I64ElementsAttr:$shape);
+  let arguments = (ins IndexElementsAttr:$shape);
   let results = (outs Shape_ShapeType:$result);
 
   // TODO: Move this to main so that all shape ops implement these.
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let hasFolder = 1;
 }
 
let arguments = (ins
     Confined<I64Attr, [IntNonNegative]>:$dim
   );
   let results = (outs Shape_SizeType:$extent);
   let assemblyFormat = "$shape `,` $dim attr-dict";
 
   let builders = [
     // Builder that allows passing a simple integer instead of an IntegerAttr.
     OpBuilder<
-      [{
-        OpBuilder &builder, OperationState &result,
-        Value shape, int64_t dim
-      }],
-      [{
-        build(builder, result, shape, builder.getI64IntegerAttr(dim));
-      }]
+      [{OpBuilder &builder, OperationState &result, Value shape, int64_t dim}],
+      [{build(builder, result, shape, builder.getI64IntegerAttr(dim));}]
     >
   ];
 
   let hasFolder = 1;
 }
 
 def Shape_JoinOp : Shape_Op<"join", []> {
   let summary = "Returns the least general shape.size of its operands";

mlir/lib/Dialect/Shape/IR/Shape.cpp

OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
   auto rhsShape = llvm::to_vector<6>(
       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
   SmallVector<int64_t, 6> resultShape;
   // If the shapes are not compatible, we can't fold it.
   // TODO: Fold to an "error".
   if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
     return nullptr;
   Builder builder(getContext());
-  return builder.getI64TensorAttr(resultShape);
+  return builder.getIndexTensorAttr(resultShape);
 }
 
 //===----------------------------------------------------------------------===//
 // ConstShapeOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, ConstShapeOp &op) {
   p << "shape.const_shape ";
static ParseResult parseConstShapeOp(OpAsmParser &parser,
   SmallVector<int64_t, 6> ints;
   for (Attribute extent : extentsArray) {
     IntegerAttr attr = extent.dyn_cast<IntegerAttr>();
     if (!attr)
       return failure();
     ints.push_back(attr.getInt());
   }
   Builder &builder = parser.getBuilder();
-  result.addAttribute("shape", builder.getI64TensorAttr(ints));
+  result.addAttribute("shape", builder.getIndexTensorAttr(ints));
 
   result.types.push_back(ShapeType::get(builder.getContext()));
   return success();
 }
 
 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shape(); }
 
 //===----------------------------------------------------------------------===//

 
 OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
   if (llvm::any_of(operands, [](Attribute a) { return !a; }))
     return nullptr;
   SmallVector<int64_t, 6> extents;
   for (auto attr : operands)
     extents.push_back(attr.cast<IntegerAttr>().getInt());
   Builder builder(getContext());
-  return builder.getI64TensorAttr(extents);
+  return builder.getIndexTensorAttr(extents);
 }
 
 //===----------------------------------------------------------------------===//
 // GetExtentOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 GetExtentOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                               ValueRange operands, DictionaryAttr attributes,
                               RegionRange regions,
                               SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(SizeType::get(context));
   return success();
 }
 
 OpFoldResult GetExtentOp::fold(ArrayRef<Attribute> operands) {
   auto elements = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
   if (!elements)
     return nullptr;
   uint64_t dimToGet = dim().getLimitedValue();
   // TODO: Constant fold this to some kind of constant error.
   if (dimToGet >= (uint64_t)elements.getNumElements())
     return nullptr;
-  // This is a little inconvenient because getValue returns an IntegerAttr
-  // that is not of IndexType, but the result here needs to be of
-  // IndexType.
-  // TODO: Make ConstShapeOp hold an tensor of index instead of i64.
-  Builder builder(getContext());
-  return builder.getIntegerAttr(
-      builder.getIndexType(),
-      elements.getValue<IntegerAttr>({dimToGet}).getInt());
+  return elements.getValue({dimToGet});
 }
 
 //===----------------------------------------------------------------------===//
 // ShapeOfOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ShapeOfOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                             ValueRange operands, DictionaryAttr attributes,
                             RegionRange regions,
                             SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(ShapeType::get(context));
   return success();
 }
 
 OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
   auto type = getOperand().getType().dyn_cast<ShapedType>();
   if (!type || !type.hasStaticShape())
     return nullptr;
   Builder builder(getContext());
-  return builder.getI64TensorAttr(type.getShape());
+  return builder.getIndexTensorAttr(type.getShape());
 }
 
 //===----------------------------------------------------------------------===//
 // SplitAtOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 SplitAtOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
LogicalResult SplitAtOp::fold(ArrayRef<Attribute> operands,
   // Verify that the split point is in the correct range.
   // TODO: Constant fold to an "error".
   int64_t rank = shape.size();
   if (!(-rank <= splitPoint && splitPoint <= rank))
     return failure();
   if (splitPoint < 0)
     splitPoint += shape.size();
   Builder builder(operands[0].getContext());
-  results.push_back(builder.getI64TensorAttr(shape.take_front(splitPoint)));
-  results.push_back(builder.getI64TensorAttr(shape.drop_front(splitPoint)));
+  results.push_back(builder.getIndexTensorAttr(shape.take_front(splitPoint)));
+  results.push_back(builder.getIndexTensorAttr(shape.drop_front(splitPoint)));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ConcatOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
OpFoldResult ConcatOp::fold(ArrayRef<Attribute> operands) {
   auto lhsShape = llvm::to_vector<6>(
       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
   auto rhsShape = llvm::to_vector<6>(
       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
   SmallVector<int64_t, 6> resultShape;
   resultShape.append(lhsShape.begin(), lhsShape.end());
   resultShape.append(rhsShape.begin(), rhsShape.end());
   Builder builder(getContext());
-  return builder.getI64TensorAttr(resultShape);
+  return builder.getIndexTensorAttr(resultShape);
 }
 
 //===----------------------------------------------------------------------===//
 // ToExtentTensorOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ToExtentTensorOp::fold(ArrayRef<Attribute> operands) {
   if (!operands[0])

mlir/test/Dialect/Shape/canonicalize.mlir

 
 // -----
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> (!shape.shape, !shape.shape) {
   // CHECK: shape.const_shape [2, 3]
   // CHECK: shape.const_shape [4, 5]
   %c2 = constant 2 : i32
-  %0 = "shape.const_shape"() {shape = dense<[2, 3, 4, 5]> : tensor<4xi64>} : () -> !shape.shape
+  %0 = shape.const_shape [2, 3, 4, 5]
   %head, %tail = "shape.split_at"(%0, %c2) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
   return %head, %tail : !shape.shape, !shape.shape
 
 }
 
 // -----
 // Negative split point.
 // CHECK-LABEL: func @f
 func @f() -> (!shape.shape, !shape.shape) {
   // CHECK: shape.const_shape [2, 3, 4]
   // CHECK: shape.const_shape [5]
   %c-1 = constant -1 : i32
-  %0 = "shape.const_shape"() {shape = dense<[2, 3, 4, 5]> : tensor<4xi64>} : () -> !shape.shape
+  %0 = shape.const_shape [2, 3, 4, 5]
   %head, %tail = "shape.split_at"(%0, %c-1) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
   return %head, %tail : !shape.shape, !shape.shape
 }
 
 // -----
 // Out of range split point. No folding.
 // CHECK-LABEL: func @f
 func @f() -> (!shape.shape, !shape.shape) {
   // CHECK: shape.split_at
   %c5 = constant 5 : i32
-  %0 = "shape.const_shape"() {shape = dense<[2, 3, 4, 5]> : tensor<4xi64>} : () -> !shape.shape
+  %0 = shape.const_shape [2, 3, 4, 5]
   %head, %tail = "shape.split_at"(%0, %c5) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
   return %head, %tail : !shape.shape, !shape.shape
 }
 
 // -----
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {