[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

191

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

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

412

maybe this:

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

422

static_cast<uint64_t>(elements.getNumElements())

mlir/test/Dialect/Shape/canonicalize.mlir

226

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

[silvas]

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

193

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)

210

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

412

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

210

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

193

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

226

Fixed entire file.

[jpienaar]

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

193

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

193

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.
   }];
   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<int64_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<int64_t> GetExtentOp::getConstantDim() {
+  if (auto constSizeOp = dim().getDefiningOp<ConstSizeOp>()) {
+    return constSizeOp.value().getLimitedValue();
+  }
+  return llvm::None;
+}
+
 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<int64_t> dim = getConstantDim();
+  if (!dim.hasValue())
     return nullptr;
-  return elements.getValue({dimToGet});
+  if (dim.getValue() >= elements.getNumElements())
+    return nullptr;
+  return elements.getValue({(uint64_t)dim.getValue()});
+}
+
+void GetExtentOp::build(OpBuilder &builder, OperationState &result, Value shape,
+                        int64_t dim) {
+  auto loc = result.location;
+  auto dimAttr = builder.getIndexAttr(dim);
+  Value dimValue = builder.create<ConstSizeOp>(loc, dimAttr);
+  build(builder, result, shape, dimValue);
 }
 
 //===----------------------------------------------------------------------===//
 // NumElementsOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
 

mlir/test/Dialect/Shape/canonicalize.mlir

 // RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize <%s | FileCheck %s --dump-input=fail
 
-// -----
 // CHECK-LABEL: func @f
 func @f(%arg0: tensor<2x3x4xf32>) -> !shape.shape {
   // CHECK: shape.const_shape [2, 3, 4]
   %0 = "shape.shape_of"(%arg0) : (tensor<2x3x4xf32>) -> !shape.shape
   return %0 : !shape.shape
 }
 
 // -----
+
 // 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 [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 [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 [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 {
   // CHECK: shape.const_shape [7, 2]
   %0 = shape.const_shape [1, 2]
   %1 = shape.const_shape [7, 1]
   %2 = "shape.broadcast"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   return %2 : !shape.shape
 }
 
 // -----
+
 // Incompatible shapes. No folding.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: shape.broadcast
   %0 = shape.const_shape [2]
   %1 = shape.const_shape [7]
   %2 = "shape.broadcast"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   return %2 : !shape.shape
 }
 
 // -----
+
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: shape.const_shape [0, 1, 2, 3]
   %lhs = shape.const_shape [0, 1]
   %rhs = shape.const_shape [2, 3]
   %0 = "shape.concat"(%lhs, %rhs) : (!shape.shape, !shape.shape) -> !shape.shape
   return %0 : !shape.shape
 }
 
 // -----
+
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> tensor<2xindex> {
   // CHECK: constant dense<[0, 1]> : tensor<2xindex>
   %cs = shape.const_shape [0, 1]
   %0 = "shape.to_extent_tensor"(%cs) : (!shape.shape) -> tensor<2xindex>
   return %0 : tensor<2xindex>
 }
 
 // -----
+
 // Basic case.
 // CHECK-LABEL: func @f()
 func @f() -> !shape.shape {
   // CHECK: shape.const_shape [3, 5, 11]
   %e0 = constant 3 : index
   %e1 = constant 5 : index
   %e2 = constant 11 : index
   %ret = shape.from_extents %e0, %e1, %e2
   return %ret : !shape.shape
 }
 
+// -----
+
 // 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
 }
 
 // -----
+
 // Cast constant size to index and fold it away.
 // CHECK-LABEL: func @const_size_to_index
 func @const_size_to_index() -> index {
   // CHECK-NOT: shape.index_cast
   %cs = shape.const_size 123
   // CHECK: constant 123 : index
   %ci = shape.size_to_index %cs
   return %ci : index
 }
 
 // -----
+
 // Cast constant index to size and fold it away.
 // CHECK-LABEL: func @const_index_to_size
 func @const_index_to_size() -> !shape.size {
   // CHECK-NOT: index_cast
   %ci = constant 123 : index
   // CHECK: shape.const_size 123
   %cs = shape.index_to_size %ci
   return %cs : !shape.size
 }
 
 // -----
+
 // Cast constant index to size, then back, and fold it away.
 // CHECK-LABEL: func @const_index_to_size_to_index
 func @const_index_to_size_to_index() -> index {
   // CHECK-NOT: shape.index_cast
   %ci0 = constant 123 : index
   %cs0 = shape.index_to_size %ci0
   // CHECK: %[[CI:.*]] = constant 123 : index
   // CHECK-NEXT: return %[[CI]] : index
   %ci1 = shape.size_to_index %cs0
   return %ci1 : index
 }
 
 // -----
+
 // No folding.
 // CHECK-LABEL: func @nonfoldable_size_to_index
 func @nonfoldable_size_to_index(%cs : !shape.size) -> index {
   // CHECK: shape.size_to_index
   %ci = shape.size_to_index %cs
   return %ci : index
 }
 
 // -----
+
 // No folding.
 // CHECK-LABEL: func @nonfoldable_index_to_size
 func @nonfoldable_index_to_size(%ci : index) -> !shape.size {
   // CHECK: shape.index_to_size
   %cs = shape.index_to_size %ci
   return %cs : !shape.size
 }
 
 // -----
+
 // Fold number of elements computation.
 // CHECK-LABEL: func @num_elements
 func @num_elements() -> !shape.size {
   // CHECK-NOT: shape.const_shape
   %shape = shape.const_shape [4, 5, 6]
   // CHECK-NOT: shape.num_elements
   %num_elements = shape.num_elements %shape
   // CHECK: %[[NUM:.*]] = shape.const_size 120
   // CHECK-NEXT: return %[[NUM]] : !shape.size
   return %num_elements : !shape.size
 }
 
 // -----
+
 // No folding.
 // 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) {
func @f(%arg0: !shape.shape, %arg1: !shape.shape) {
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.cstr_eq %arg0, %arg1
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
 
 // -----
+
 // assuming_all with known passing witnesses can be folded
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.const_witness true
   %1 = shape.const_witness true
   %2 = shape.const_witness true
   %3 = shape.assuming_all %0, %1, %2
   "consume.witness"(%3) : (!shape.witness) -> ()
   return
 }
 
 // -----
+
 // assuming_all should not be removed if not all witnesses are statically passing.
 //
 // Additionally check that the attribute is moved to the end as this op is
 // commutative.
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: %[[UNKNOWN:.*]] = "test.source"
   // CHECK-NEXT: shape.assuming_all %[[UNKNOWN]]
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.const_witness true
   %1 = "test.source"() : () -> !shape.witness
   %2 = shape.assuming_all %0, %1
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
+
 // any can be replaced with a constant input if it has one.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: %[[CS:.*]] = shape.const_shape
   // CHECK-NEXT: return %[[CS]]
   %0 = shape.const_shape [2, 3, 4]
   %1 = shape.any %0, %arg0
   return %1 : !shape.shape
 }
 
 
 // -----
+
 // Folding of any with partially constant operands is not yet implemented.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: shape.any
   // CHECK-NEXT: return %[[CS]]
   %1 = shape.any %arg0, %arg1
   return %1 : !shape.shape
 }
 
 // -----
+
 // assuming with a known passing witness can be removed
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: source
   // CHECK-NEXT: sink
   // CHECK-NEXT: return
   %0 = shape.const_witness true
   %1 = shape.assuming %0 -> index {
     %2 = "test.source"() : () -> (index)
     shape.assuming_yield %2 : index
   }
   "test.sink"(%1) : (index) -> ()
   return
 }
 
 // -----
+
 // assuming without a known passing passing witness cannot be removed
 // CHECK-LABEL: func @f
 func @f() {
   // CHECK-NEXT: test.source
   // CHECK-NEXT: shape.assuming
   // CHECK-NEXT:   test.source
   // CHECK-NEXT:   shape.assuming_yield
   // CHECK-NEXT: }