Add `shape.get_extent`.

[silvas]

Summary

This op extracts an extent from a shape.

This also is the first op which constant folds to shape.const_size,
which revealed that sahpe.const_size needs a folder (ConstantLike ops
seem to always need folders for the constant folding infra to work).

[silvas]

Update commit message.

[jpienaar]

Nice, thanks

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

214

This fits on one line

[I'm actually not sure why this one isn't generated by default, we have the generation for unwrapped types that I would have expected to trigger here ...]

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

290

I'm not convinced of this unless we add this check to this op's verify method.

294

What is the type? E.g., are we storing it wrong?

Is only shape.const_shape expected here? If not, then this conversion (and whether it is signed/unsigned/signless) seem to need to support more cases (based on the requirement of getInt()).

[tpopp]

Just a note of a typo in your commit message "sahpe". I see you updated the diff with an updated commit message, so I'm guessing it's fixed, but arc doesn't propagate updated commit messages unfortunately.

[tpopp]

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

298

There is a getIndexAtrr(int64_t) method that you could use here.

[jpienaar]

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

290

And this is not UB: this is an error. It can be folded to constant size error.

[silvas]

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

294

Good catch. We were using I64ElementsAttr for representing shapes when constant folding, but we should have been using IndexElementsAttr (which doesn't exist yet). I'll add IndexElementsAttr to core and fix this as part of the upgrade the shape dialect to using that.

298

it is just elements.getValue({dimToGet}) after the IndexElementsAttr fix (patches will be uploaded soon).

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

let description = [{
     %x = shape.const_size 10
     ```
   }];
 
   let arguments = (ins IndexAttr:$value);
   let results = (outs Shape_SizeType:$result);
 
   let assemblyFormat = "attr-dict $value";
+  let hasFolder = 1;
 }
 
 def Shape_FromExtentsOp : Shape_Op<"from_extents", [
     NoSideEffect,
     DeclareOpInterfaceMethods<InferTypeOpInterface>
     ]> {
   let summary = "Creates a shape from extents";
   let description = [{
def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", []> {
   }];
 
   let arguments = (ins Shape_ShapeType:$input);
   let results = (outs IndexTensor:$result);
 
   let hasFolder = 1;
 }
 
+def Shape_GetExtentOp : Shape_Op<"get_extent",
+    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let summary = "Gets the specified extent from a shape";
+  let description = [{
+    Gets the extent indexed by `dim` from `shape`.
+
+    If the shape is an error, it returns an error size.
+  }];
+  let arguments = (ins
+    Shape_ShapeType:$shape,
+    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));
+      }]
+    >
+  ];
+
+  let hasFolder = 1;
+}
+
 def Shape_JoinOp : Shape_Op<"join", []> {
   let summary = "Returns the least general shape.size of its operands";
   let description = [{
     An operation that computes the least general shape of input operands. This
     effectively asserts that corresponding static dimensions are equal. The
     behavior is to match each element of the `shape.type` and propagate the most
     restrictive information, returning an invalid shape if there are
     contradictory requirements. E.g., using pseudo code

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

 ConstSizeOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                               ValueRange operands, DictionaryAttr attributes,
                               RegionRange regions,
                               SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(SizeType::get(context));
   return success();
 }
 
+OpFoldResult ConstSizeOp::fold(ArrayRef<Attribute>) { return valueAttr(); }
+
 //===----------------------------------------------------------------------===//
 // FromExtentsOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult FromExtentsOp::inferReturnTypes(
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(ShapeType::get(context));
   return success();
 }
 
 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);
 }
 
 //===----------------------------------------------------------------------===//
+// 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());
+}
+
+//===----------------------------------------------------------------------===//
 // ShapeOfOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ShapeOfOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                             ValueRange operands, DictionaryAttr attributes,
                             RegionRange regions,
                             SmallVectorImpl<Type> &inferredReturnTypes) {

mlir/test/Dialect/Shape/canonicalize.mlir

 
 // CHECK-LABEL: func @no_fold
 func @no_fold(%arg0: index) -> !shape.shape {
   // CHECK-NOT: shape.const_shape
   %e0 = constant 3 : index
   %ret = shape.from_extents %e0, %arg0
   return %ret : !shape.shape
 }
+
+// -----
+// Canonicalization of shape.get_extent
+
+// Basic folding.
+// CHECK-LABEL: func @basic
+func @basic() -> !shape.size {
+  // CHECK: shape.const_size 2
+  %0 = shape.const_shape [0, 1, 2]
+  %1 = shape.get_extent %0, 2
+  return %1 : !shape.size
+}
+
+// Should not fold.
+// CHECK-LABEL: func @out_of_bounds
+func @out_of_bounds() -> !shape.size {
+  // CHECK: shape.const_shape
+  // CHECK: shape.get_extent
+  %0 = shape.const_shape [0, 1, 2]
+  %1 = shape.get_extent %0, 3
+  return %1 : !shape.size
+}
+
+// Should not fold.
+// CHECK-LABEL: func @not_const
+func @not_const(%arg0: !shape.shape) -> !shape.size {
+  // CHECK: shape.get_extent
+  %0 = shape.get_extent %arg0, 3
+  return %0 : !shape.size
+}