[MLIR][Shape] Add constant folding to `shape.get_shape`

[frgossen]

Summary

Add constant folding for the shape.get_shape operation of the shape dialect.

Depends On D82028

[frgossen]

Rebase

[frgossen]

Rebase

[frgossen]

Rebase

[frgossen]

Rebase

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

def Shape_GetSizeOp : Shape_Op<"get_size", [NoSideEffect]> {
   let description = [{
     Returns the size of the shape, i.e. the number of extents.
   }];
 
   let arguments = (ins Shape_ShapeType:$shape);
   let results = (outs Shape_SizeType:$size);
 
   let assemblyFormat = "attr-dict $shape";
+  let hasFolder = 1;
 }
 
 def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
   let summary = "Creates a dimension tensor from a shape";
   let description = [{
     Converts a shape to a 1D integral tensor of extents. The number of elements
     in the tensor equals the rank of the shape, and the elements equal the
     extents of the shape.

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

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);
 }
 
 //===----------------------------------------------------------------------===//
+// GetSizeOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult GetSizeOp::fold(ArrayRef<Attribute> operands) {
+  auto shape = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
+  if (!shape)
+    return {};
+  int64_t size = shape.end() - shape.begin();
+  Builder builder(getContext());
+  return builder.getIndexAttr(size);
+}
+
+//===----------------------------------------------------------------------===//
 // NumElementsOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
 
   // Fold only when argument constant.
   Attribute shape = operands[0];
   if (!shape)

mlir/test/Dialect/Shape/canonicalize.mlir

 func @f(%arg0 : !shape.shape) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.cstr_broadcastable %arg0, %arg0
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
+
+// -----
+
+// Fold `get_size` based on statically shaped tensor.
+// CHECK-LABEL: @fold_size
+func @fold_size(%arg : tensor<2x3x4xf32>) -> !shape.size {
+  // CHECK-DAG: %[[RESULT:.*]] = shape.const_size 3
+  // CHECK-DAG: return %[[RESULT]] : !shape.size
+  %shape = shape.shape_of %arg : tensor<2x3x4xf32>
+  %size = shape.get_size %shape
+  return %size : !shape.size
+}
+
+// -----
+
+// Fold `get_size` based on constant shape.
+// CHECK-LABEL: @fold_size
+func @fold_size() -> !shape.size {
+  // CHECK-DAG: %[[RESULT:.*]] = shape.const_size 5
+  // CHECK-DAG: return %[[RESULT]] : !shape.size
+  %shape = shape.const_shape [3, 4, 5, 6, 7]
+  %size = shape.get_size %shape
+  return %size : !shape.size
+}
+
+// -----
+
+// Do not fold `get_size` if shape is dynamic.
+// CHECK-LABEL: @dont_fold_size
+// CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) -> !shape.size
+func @dont_fold_size(%shape : !shape.shape) -> !shape.size {
+  // CHECK-DAG: %[[RESULT:.*]] = shape.get_size %[[SHAPE]]
+  // CHECK-DAG: return %[[RESULT]] : !shape.size
+  %size = shape.get_size %shape
+  return %size : !shape.size
+}