[shape] Basic constant folding.

[silvas]

Summary

• Implement a first constant fold for shape.shape_of (more ops coming in subsequent patches)
• Implement the right builder interfaces for ShapeType and other types
• Splits shape.constant into shape.const_size and shape.const_shape which plays better with dyn_cast and building vs one polymorphic op.

Also, fix the RUN line in ops.mlir to properly verify round-tripping.

Depends On D78752 (future patches in the constant-folding patchset will)

[silvas]

Fix include path

[silvas]

use Builder::getType

[silvas]

Fix typo.

[mehdi_amini]

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

163

I don't think we should repeat the summary in the description?

164

This needs an update.

169

This example needs an update.

185

It'd be nice to provide the long description like the other above (these will show up on the website)

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

39

Should handle SizeType here as well?

111

Having a builder and using getI64IntegerAttr(int64_t value); seems like it'd help here

[jpienaar]

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

157

Let's keep these sorted

166

This hasn't been resolved and so should still be left todo

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

88

Let's keep the ops sorted alphabetical

92

This op doesn't operate only on ShapedType so use dyn_cast instead.

111

Seems wasteful to construct an arrayattr and unique it in context just to print, an interleaveWithComma call should suffice.

[silvas]

Update.

[silvas]

Address comments.

[silvas]

update

[jpienaar]

Thanks!

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

76

I'd like to see tests showing each of these being printed (well OK with invalid one not being shown), but it wouldn't be possible without custom parser, so OK to start with the current version.

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

95

This should use ? for unknown and regular integer for everything else. Adding support for this would mean the IR can't roundtrip, but please add a TODO for follow up

125

Unknown's should be parsed as ? in the descriptions so we would need some additional logic here. E.g., -1 shouldn't be used in the textual form. So a more custom parsing would be required, could you add a TODO here too

132

This seems equally long to

if (auto attr = extent.dyn_cast<IntegerAttr>())
  ints.push_back(attr.getInt());
else
  return failure();

Don't know which one reads easier [this one feels like variable usage scope is more local and makes the two cases clearer but thats just a feelign and not a trong one]

[jpienaar]

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

175

What happens here if we have a error shape fed into ShapeOfOp?

[silvas]

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

95

See other comment.

125

I don't understand the use case for shape constants of partially unknown shape. Maybe we can discuss on discourse?

175

We check that the tensor operand has static shape. So it can't be an error.

[silvas]

landed in d1ad267a56d33c4a6556c2947c8138123cefc08d

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

let description = [{
     shapes as input, return the invalid shape if one of its operands is an
     invalid shape. This avoids flagging multiple errors for one verification
     failure. The dialect itself does not specify how errors should be combined
     (there are multiple different options, from always choosing first operand,
     concatting etc. on how to combine them).
   }];
 
   let cppNamespace = "shape";
+
+  let hasConstantMaterializer = 1;
 }
 
 def Shape_ComponentType : DialectType<ShapeDialect,
-    CPred<"$_self.isa<::mlir::shape::ComponentType>()">, "component type"> {
+    CPred<"$_self.isa<::mlir::shape::ComponentType>()">, "component type">,
+    BuildableType<"$_builder.getType<::mlir::shape::ComponentType>()"> {
   let typeDescription = [{
     `shape.element_type` represents the element type of the ShapedType. It may
     be unknown, error or regular element type supported by ShapedType.
   }];
 }
 
 def Shape_ElementType : DialectType<ShapeDialect,
-    CPred<"$_self.isa<::mlir::shape::ElementType>()">, "element type"> {
+    CPred<"$_self.isa<::mlir::shape::ElementType>()">, "element type">,
+    BuildableType<"$_builder.getType<::mlir::shape::ElementType>()"> {
   let typeDescription = [{
     `shape.element_type` represents the element type of the ShapedType. It may
     be unknown, error or regular element type supported by ShapedType.
   }];
 }
 
 def Shape_ShapeType : DialectType<ShapeDialect,
-    CPred<"$_self.isa<::mlir::shape::ShapeType>()">, "shape"> {
+    CPred<"$_self.isa<::mlir::shape::ShapeType>()">, "shape">,
+    BuildableType<"$_builder.getType<::mlir::shape::ShapeType>()"> {
   let typeDescription = [{
     `shape.type` represents either an unranked shape, a ranked shape with
     possibly unknown dimensions or an invalid shape. The rank is of type
     `shape.size` and, if rank is known, the extent is a 1D tensor of type
     `shape.size`.
 
     Shape is printed:
     * `[*]` if it is an unranked shape
     * `[?, 2]` if a rank 2 tensor with one unknown dimension
     * `[3, 4]` is a rank 2 static tensor
     * `[]` is a scalar
     * `[1]` is a rank 1 tensor with 1 element
     * `[invalid]` for an invalid shape
   }];
 }
 
 def Shape_SizeType : DialectType<ShapeDialect,
-    CPred<"$_self.isa<::mlir::shape::SizeType>()">, "size"> {
+    CPred<"$_self.isa<::mlir::shape::SizeType>()">, "size">,
+    BuildableType<"$_builder.getType<::mlir::shape::SizeType>()"> {
   let typeDescription = [{
     `shape.size` represents a non-negative integer with support for being
     unknown and invalid.
 
     Operations on `shape.size` types are specialized to handle unknown/dynamic
     value. So, for example, `<unknown> + x == <unknown>` for all non-error `x :
     !shape.size` (e.g., an unknown value does not become known due to addition).
   }];
 }
 
 def Shape_ValueShapeType : DialectType<ShapeDialect,
-    CPred<"$_self.isa<::mlir::shape::ValueShapeType>()">, "value shape"> {
+    CPred<"$_self.isa<::mlir::shape::ValueShapeType>()">, "value shape">,
+    BuildableType<"::mlir::shape::ValueShapeType::get($_builder.getContext())">
+{
   let typeDescription = [{
     `shape.value_shape` represents the value produced by an operation (this
     corresponds to `Value` in the compiler) and a shape. Conceptually this is a
     tuple of a value (potentially unknown) and `shape.type`. The value and shape
     can either or both be unknown. If both the `value` and `shape` are known,
     then the shape of `value` is conformant with `shape`.
   }];
 }
let description = [{
     broadcastable output shape possible for the given inputs.
   }];
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs,
                    OptionalAttr<StrAttr>:$error);
   let results = (outs Shape_ShapeType:$result);
 }
 
-def Shape_ConstantOp : Shape_Op<"constant", []> {
-  let summary = "Creates a shape constant";
+def Shape_ConstSizeOp : Shape_Op<"const_size",
+    [ConstantLike,
+     NoSideEffect,
+     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let summary = "Creates a constant of !shape.size type.";
   let description = [{
+    Creates a constant of !shape.size type.
     An operation that builds a size or shape from integer or array attribute.
     It allows for creating dynamically valued shapes by using `?` for unknown
     values. A constant shape specified with `*` will return an unranked shape.
 
     ```mlir
     %x = shape.constant 10 : !shape.size
     ```
   }];
 
-  // TODO(jpienaar): Change to a more specialized attribute that would
-  // encapsulate the unknown parsing while using denser packing.
-  let arguments = (ins AnyAttr:$value);
-  let results = (outs Shape_ShapeOrSizeType:$result);
+  let arguments = (ins IndexAttr:$value);
+  let results = (outs Shape_SizeType:$result);
+
+  let assemblyFormat = "attr-dict $value";
+}
+
+def Shape_ConstShapeOp : Shape_Op<"const_shape",
+    [ConstantLike,
+     NoSideEffect,
+     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let summary = "Creates a constant of !shape.shape type.";
+  let description = [{
+    Creates a constant of !shape.shape type.
+  }];
+  let arguments = (ins I64ElementsAttr:$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 verifier = [{ return ::verify(*this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
+  let hasFolder = 1;
 }
 
 def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", []> {
   let summary = "Creates a shape from a tensor of extents";
   let description = [{
     Creates a shape from a 1D integral tensor of extents. The rank of the
     resulting shape equals the number of elements in the tensor, and the
     extents match the values of the elements.
def Shape_ReduceOp : Shape_Op<"reduce", []> {
   let regions = (region SizedRegion<1>:$body);
 }
 
 def Shape_ShapeOfOp : Shape_Op<"shape_of", []> {
   let summary = "Returns shape of a value or shaped type operand";
 
   let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$arg);
   let results = (outs Shape_ShapeType:$result);
+
+  let hasFolder = 1;
 }
 
 def Shape_YieldOp : Shape_Op<"yield", [NoSideEffect, Terminator]> {
   let summary = "Returns the value to parent op";
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
   let builders = [OpBuilder<

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

 //===- Shape.cpp - MLIR Shape Operations ----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Shape/IR/Shape.h"
 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/StandardTypes.h"
+#include "mlir/IR/PatternMatch.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::shape;
 
 ShapeDialect::ShapeDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
       >();
   addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType>();
   // Allow unknown operations during prototyping and testing. As the dialect is
   // still evolving it makes it simple to start with an unregistered ops and
   // try different variants before actually defining the op.
   allowUnknownOperations();
 }
 
+Operation *ShapeDialect::materializeConstant(OpBuilder &builder,
+                                             Attribute value, Type type,
+                                             Location loc) {
+  if (auto shapeType = type.dyn_cast<ShapeType>()) {
+    return builder.create<ConstShapeOp>(loc, type,
+                                        value.cast<DenseIntElementsAttr>());
+  }
+  return nullptr;
+}
+
 /// Parse a type registered to this dialect.
 Type ShapeDialect::parseType(DialectAsmParser &parser) const {
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return Type();
 
   if (keyword == "component")
     return ComponentType::get(getContext());
case ShapeTypes::ValueShape:
     os << "value_shape";
     return;
   default:
     llvm_unreachable("unexpected 'shape' type kind");
   }
 }
 
 //===----------------------------------------------------------------------===//
-// Constant*Op
+// ShapeOfOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
+  auto type = getOperand().getType().cast<ShapedType>();
+  if (!type.hasStaticShape())
+    return nullptr;
+  Builder builder(getContext());
+  return builder.getI64TensorAttr(type.getShape());
+}
+
+//===----------------------------------------------------------------------===//
+// ConstShapeOp
 //===----------------------------------------------------------------------===//
 
-static void print(OpAsmPrinter &p, ConstantOp &op) {
-  p << "shape.constant ";
-  p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"value"});
+static void print(OpAsmPrinter &p, ConstShapeOp &op) {
+  p << "shape.const_shape ";
+  p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"shape"});
 
-  if (op.getAttrs().size() > 1)
-    p << ' ';
-  p.printAttributeWithoutType(op.value());
-  p << " : " << op.getType();
+  // We piggy-back on ArrayAttr printing.
+  SmallVector<Attribute, 6> extents;
+  for (auto extent : op.shape().getValues<int64_t>()) {
+    extents.push_back(
+        IntegerAttr::get(IntegerType::get(64, op.getContext()), extent));
+  }
+  p.printAttributeWithoutType(ArrayAttr::get(extents, op.getContext()));
 }
 
-static ParseResult parseConstantOp(OpAsmParser &parser,
+static ParseResult parseConstShapeOp(OpAsmParser &parser,
                                      OperationState &result) {
-  Attribute valueAttr;
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
-  Type i64Type = parser.getBuilder().getIntegerType(64);
-  if (parser.parseAttribute(valueAttr, i64Type, "value", result.attributes))
+  // We piggy-back on ArrayAttr parsing, though we don't internally store the
+  // shape as an ArrayAttr.
+  // TODO: Implement custom parser and maybe make syntax a bit more concise.
+  Attribute extentsRaw;
+  SmallVector<NamedAttribute, 6> dummy;
+  if (parser.parseAttribute(extentsRaw, "dummy", dummy))
     return failure();
-
-  Type type;
-  if (parser.parseColonType(type))
+  auto extentsArray = extentsRaw.dyn_cast<ArrayAttr>();
+  if (!extentsArray)
+    return failure();
+  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.types.push_back(ShapeType::get(builder.getContext()));
+  return success();
+}
+
+OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shape(); }
 
-  // Add the attribute type to the list.
-  return parser.addTypeToList(type, result.types);
+LogicalResult ConstShapeOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    ArrayRef<NamedAttribute> attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(ShapeType::get(context));
+  return success();
 }
 
-static LogicalResult verify(ConstantOp &op) { return success(); }
+//===----------------------------------------------------------------------===//
+// ConstSizeOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ConstSizeOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    ArrayRef<NamedAttribute> attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(SizeType::get(context));
+  return success();
+}
 
 //===----------------------------------------------------------------------===//
 // SplitAtOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult SplitAtOp::inferReturnTypes(
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     ArrayRef<NamedAttribute> attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   auto shapeType = ShapeType::get(context);
   inferredReturnTypes.push_back(shapeType);
   inferredReturnTypes.push_back(shapeType);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ConcatOp
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/Shape/canonicalize.mlir

+// RUN: mlir-opt -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
+}

mlir/test/Dialect/Shape/ops.mlir

-// RUN: mlir-opt -split-input-file %s | FileCheck %s --dump-input-on-failure
+// RUN: mlir-opt -split-input-file %s | mlir-opt | FileCheck %s --dump-input-on-failure
 
 // CHECK-LABEL: shape_num_elements
 func @shape_num_elements(%shape : !shape.shape) -> !shape.size {
-  %0 = shape.constant 0 : !shape.size
+  %0 = shape.const_size 0
   %1 = "shape.reduce"(%shape, %0) ( {
     ^bb0(%index: i32, %dim: !shape.size, %lci: !shape.size):
       %acc = "shape.add"(%lci, %dim) : (!shape.size, !shape.size) -> !shape.size
       "shape.yield"(%acc) : (!shape.size) -> ()
     }) : (!shape.shape, !shape.size) -> (!shape.size)
   return %1 : !shape.size
 }
 
 func @test_shape_num_elements_unknown() {
   %0 = "shape.unknown_shape"() : () -> !shape.shape
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %2 = "shape.print"(%1) : (!shape.size) -> !shape.size
   return
 }
 
 func @test_shape_num_elements_fixed() {
-  %0 = "shape.constant"() { value = [1, 57, 92] }: () -> !shape.shape
+  %0 = shape.const_shape [1, 57, 92]
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %3 = "shape.print"(%1) : (!shape.size) -> !shape.size
   return
 }
 
 func @test_broadcastable_fixed() {
-  %0 = "shape.constant"() { value = [10, 1, 57, 92] }: () -> !shape.shape
-  %1 = "shape.constant"() { value = [4, 57, 92] }: () -> !shape.shape
+  %0 = shape.const_shape [10, 1, 57, 92]
+  %1 = shape.const_shape [4, 57, 92]
   %2 = "shape.broadcastable"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
 func @test_shape_any_fixed() {
-  %0 = "shape.constant"() { value = [4, 57, 92] }: () -> !shape.shape
-  %1 = "shape.constant"() { value = [4, 57, 92] }: () -> !shape.shape
+  %0 = shape.const_shape [4, 57, 92]
+  %1 = shape.const_shape [4, 57, 92]
   %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
 func @test_shape_any_unknown() {
-  %0 = "shape.constant"() { value = [4, -1, 92] }: () -> !shape.shape
-  %1 = "shape.constant"() { value = [-1, 57, 92] }: () -> !shape.shape
+  %0 = shape.const_shape [4, -1, 92]
+  %1 = shape.const_shape [-1, 57, 92]
   %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
 func @test_shape_any_fixed_mismatch() {
-  %0 = "shape.constant"() { value = [4, 57, 92] }: () -> !shape.shape
-  %1 = "shape.constant"() { value = [2, 57, 92] }: () -> !shape.shape
+  %0 = shape.const_shape [4, 57, 92]
+  %1 = shape.const_shape [2, 57, 92]
   %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
+
+func @test_parse_const_shape() {
+  %0 = shape.const_shape []
+  %1 = shape.const_shape [1, 2, 3]
+  return
+}