[MLIR][Shape] Make dimension an operand of `get_extent`

[frgossen]

Summary

The operation get_extent now accepts the dimension as an operand and is no
longer limited to constant dimensions.
A helper function facilitates the common constant use case.

Depends On D81162

[pifon2a]

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

186

these "DeclareOp..." trait will disappear after rebase

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

395

maybe this:

if (auto constSizeOp = dyn_cast<ConstSizeOp>(dim().getDefiningOp()) {
  return constSizeOp.value().getLimitedValue();
}
return llvm::None;

405

static_cast<uint64_t>(elements.getNumElements())

mlir/test/Dialect/Shape/canonicalize.mlir

208

super-nit: empty line after //----

[silvas]

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

188

Now that we are returning !shape.size, if the shape is an error, or the dim is out of bounds or an error, then the result is an error size. (!shape.size can represent an error)

205

It's slightly confusing that the builder takes int64_t but this uses uint64_t. Can you make them consistent? (for whatever reason, it seems the preference is int64_t).

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

395

getDefiningOp<ConstSizeOp>() makes it even more succinct. Also you might want to drop the braces for consistency for these small if's.

[jpienaar]

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

205

Preference is for int64_t here yes. I'd prefer we keep that uniformly here.

[frgossen]

Rebase

[frgossen]

Apply suggestions

[frgossen]

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

188

When lowering to std, we will not represent error cases, will we?
At least in that case, the behaviour is undefined.

The way I under stand it is that the error case will be handled by "guarding" assertions (at least at the point at which we lower to std).
In that sense, get_extent has two semantic variants (?).

@herhut

mlir/test/Dialect/Shape/canonicalize.mlir

208

Fixed entire file.

[jpienaar]

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

188

That is during the lowering, the op now has explicit behavior in error input case: it propagates the error via the size type. How that semantics is maintained during lowering may change (e.g., using constraint vs assert) but that is on the lowering and not a defining characteristic of the op, the op propagates the error.

[herhut]

Thanks!

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

188

Yes, the operation should return an error if the shape is an error and also if the dimension is out of bounds. We will have to figure out how we lower this correctly. For the cases we currently care about, the lowering is correct by construction (We iterate over the rank and the shape cannot contain error values.). The latter is due to using constraints but for the former we will need some way to remove bounds checks. To begin with, I would be ok with a lowering that just neglects the bounds checking.

[frgossen]

Rebase, add missing builder

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

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]> {
   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.
+    If the shape is an error, its behavior is undefined.
   }];
   let arguments = (ins
     Shape_ShapeType:$shape,
-    Confined<I64Attr, [IntNonNegative]>:$dim
+    Shape_SizeType:$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));}]
-    >
+    // Builder that allows passing a constant dimension as a simple integer.
+    OpBuilder<"OpBuilder &builder, OperationState &result, Value shape, "
+              "int64_t dim">
   ];
 
+  let extraClassDeclaration = [{
+    /// Get the `dim` value as integer if it is constant.
+    Optional<uint64_t> getConstantDim();
+  }];
+
   let hasFolder = 1;
 }
 
 def Shape_IndexToSizeOp : Shape_Op<"index_to_size", [NoSideEffect]> {
   let summary = "Converts a standard index to a shape size";
   let description = [{
     Converts a standard index to a `shape.size`. This operation and its
     inverse, `size_to_index`, facilitate index conversion between the standard

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

OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
   Builder builder(getContext());
   return builder.getIndexTensorAttr(extents);
 }
 
 //===----------------------------------------------------------------------===//
 // GetExtentOp
 //===----------------------------------------------------------------------===//
 
+Optional<uint64_t> GetExtentOp::getConstantDim() {
+  auto constSizeOp = dyn_cast<ConstSizeOp>(dim().getDefiningOp());
+  if (!constSizeOp)
+    return {};
+  return constSizeOp.value().getLimitedValue();
+}
+
 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())
+  Optional<uint64_t> dim = getConstantDim();
+  if (!dim.hasValue())
+    return nullptr;
+  if (dim.getValue() >= (uint64_t)elements.getNumElements())
     return nullptr;
-  return elements.getValue({dimToGet});
+  return elements.getValue({dim.getValue()});
 }
 
 //===----------------------------------------------------------------------===//
 // NumElementsOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
 

mlir/test/Dialect/Shape/canonicalize.mlir

 // CHECK-LABEL: func @nonfoldable_num_elements
 func @nonfoldable_num_elements(%shape : !shape.shape) -> !shape.size {
   // CHECK-NOT: shape.const_{{.*}}
   %num_elements = shape.num_elements %shape
   return %num_elements : !shape.size
 }
 
 // -----
-
-// 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
+  %c2 = shape.const_size 2
+  %1 = shape.get_extent %0, %c2
   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
+  %c3 = shape.const_size 3
+  %1 = shape.get_extent %0, %c3
   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
+  %c3 = shape.const_size 3
+  %0 = shape.get_extent %arg0, %c3
   return %0 : !shape.size
 }
 
 
 // -----
 // cstr_eq with non-constant but known equal shapes can be removed.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) {