[MLIR] Minor cleanup for Shape dialect.

[jurahul]

Summary

• Remove some unused types from the Shape dialect
• Fix from_extent_tensor to only allow 1D index tensors
• Fix assuming_yield to only allow shape.assuming as the parent op.
• Fix some documentation typos and reword some things.

[silvas]

mlir/test/Dialect/Shape/invalid.mlir

160 ↗

(On Diff #310396)

we generally don't explicitly test things verified by ODS predicates

[jurahul]

mlir/test/Dialect/Shape/invalid.mlir

160 ↗

(On Diff #310396)

Makes sense. I'll remove this test.

[jurahul]

Remove test case

[jpienaar]

Thanks. Yes the element type we have no tests for yet and can reintroduce it when we do.

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

333

While here, the first one is a must, the second is a may. E.g., for the first one where we need to represent errors we can't use index, but the second one we could use size still. I think in the code we already allow this to allow local folding without needing to do full propagation.

[jurahul]

Address feedback

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

333

Done. I also added similar language for shape.add as well.

mlir/include/mlir/Dialect/Shape/IR/Shape.h

 namespace mlir {
 class PatternRewriter;
 
 namespace shape {
 
 /// Alias type for extent tensors.
 RankedTensorType getExtentTensorType(MLIRContext *ctx);
 
-/// The component type corresponding to shape, element type and attribute.
-class ComponentType : public Type::TypeBase<ComponentType, Type, TypeStorage> {
-public:
-  using Base::Base;
-};
-
-/// The element type of the shaped type.
-class ElementType : public Type::TypeBase<ElementType, Type, TypeStorage> {
-public:
-  using Base::Base;
-};
-
 /// The shape descriptor type represents rank and dimension sizes.
 class ShapeType : public Type::TypeBase<ShapeType, Type, TypeStorage> {
 public:
   using Base::Base;
 };
 
 /// The type of a single dimension.
 class SizeType : public Type::TypeBase<SizeType, Type, TypeStorage> {

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

let description = [{
     concatting etc. on how to combine them).
   }];
 
   let cppNamespace = "::mlir::shape";
 
   let hasConstantMaterializer = 1;
 }
 
-def Shape_ComponentType : DialectType<ShapeDialect,
-    CPred<"$_self.isa<::mlir::shape::ComponentType>()">, "component type">,
-    BuildableType<"$_builder.getType<::mlir::shape::ComponentType>()"> {
-  let typeDescription = [{
-    `shape.component_type` represents the tuple of shape, element type and
-    attribute.
-  }];
-}
-
-def Shape_ElementType : DialectType<ShapeDialect,
-    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">,
     BuildableType<"$_builder.getType<::mlir::shape::ShapeType>()"> {
   let typeDescription = [{
-    `shape.type` represents either an unranked shape, a ranked shape with
+    `shape.shape` 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

 
 def Shape_ValueShapeType : DialectType<ShapeDialect,
     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`. That is, the shape of
-    the value conforms to the shape of the ValueShape, so that if we have
-    ValueShape `(value, shape)` then `join(shape_of(value), shape)` would be
-    error free and in particular it means that if both are statically known,
+    tuple of a value (potentially unknown) and `shape.shape`. 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`. That is, the
+    shape of the value conforms to the shape of the ValueShape, so that if we
+    have ValueShape `(value, shape)` then `join(shape_of(value), shape)` would
+    be error free and in particular it means that if both are statically known,
     then they are equal.
   }];
 }
 
 def Shape_ExtentTensorType :
     1DTensorOf<[Index]>,
     BuildableType<"::mlir::RankedTensorType::get({ShapedType::kDynamicSize}, "
                   "$_builder.getType<::mlir::IndexType>())"> {
   let typeDescription = [{
     The extent tensor is a tensor of rank one with arbitrarily many index
-    elements. Like `!shape.shape`, it is used to represent shapes with the
-    difference that it is guaranteed to be error-free.
+    elements (tensor<?xindex>). Like `!shape.shape`, it is used to represent
+    shapes with the difference that it is guaranteed to be error-free.
   }];
 }
 
 def Shape_ShapeOrSizeType : AnyTypeOf<[Shape_SizeType, Shape_ShapeType],
   "shape or size">;
 
 def Shape_ShapeOrExtentTensorType : AnyTypeOf<[Shape_ShapeType,
                                                Shape_ExtentTensorType],

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

class Shape_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<ShapeDialect, mnemonic, traits>;
 
 def Shape_AddOp : Shape_Op<"add", [Commutative, NoSideEffect]> {
   let summary = "Addition of sizes and indices";
   let description = [{
     Adds two sizes or indices. If either operand is an error it will be
     propagated to the result. The operands can be of type `size` or `index`. If
     at least one of the operands can hold an error, i.e. if it is of type `size`,
-    then also the result must be of type `size`.
+    the result must be of type `size`. If error propagation is not possible
+    because both operands are of type `index` then the result may be of type
+    `size` or `index`.
   }];
 
   let arguments = (ins Shape_SizeOrIndexType:$lhs, Shape_SizeOrIndexType:$rhs);
   let results = (outs Shape_SizeOrIndexType:$result);
 
   let assemblyFormat = [{
     $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
   }];

 def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", [NoSideEffect]> {
   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.
   }];
 
-  let arguments = (ins IndexTensor:$input);
+  let arguments = (ins 1DTensorOf<[Index]>:$input);
   let results = (outs Shape_ShapeType:$result);
 
   let assemblyFormat = "$input attr-dict `:` type($input)";
 }
 
 def Shape_IsBroadcastableOp : Shape_Op<"is_broadcastable", [Commutative]> {
   let summary = "Determines if 2 shapes can be successfully broadcasted";
   let description = [{
def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
 
   let hasFolder = 1;
 }
 
 def Shape_GetExtentOp : Shape_Op<"get_extent", [NoSideEffect]> {
   let summary = "Gets the specified extent from a shape or extent tensor";
   let description = [{
     Gets the extent indexed by `dim` from the `shape` operand. If the shape is
-    an error then it returns an error size.
+    an error then it returns an invalid size.
   }];
   let arguments = (ins Shape_ShapeOrExtentTensorType:$shape,
                        Shape_SizeOrIndexType:$dim);
   let results = (outs Shape_SizeOrIndexType:$extent);
   let assemblyFormat = "$shape `,` $dim attr-dict `:` type($shape) `,` type($dim) `->` "
                        "type($extent)";
 
   let builders = [
def Shape_IndexToSizeOp : Shape_Op<"index_to_size", [NoSideEffect]> {
 
   let assemblyFormat = "$arg attr-dict";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 def Shape_JoinOp : Shape_Op<"join", [Commutative]> {
-  let summary = "Returns the least general shape.size of its operands";
+  let summary = "Returns the least general shape.shape 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.shape` and propagate the
     most restrictive information, returning an invalid shape if there are
     contradictory requirements. E.g., using pseudo code
 
     ```

 }
 
 def Shape_MulOp : Shape_Op<"mul", [Commutative, NoSideEffect]> {
   let summary = "Multiplication of sizes and indices";
   let description = [{
     Multiplies two sizes or indices. If either operand is an error it will be
     propagated to the result. The operands can be of type `size` or `index`. If
     at least one of the operands can hold an error, i.e. if it is of type `size`,
-    then also the result must be of type `size`. If error propagation is not
-    possible because both operands are of type `index` then the result must also
-    be of type `index`.
+    the result must be of type `size`. If error propagation is not possible
+    because both operands are of type `index` then the result may be of type
+    `size` or `index`.
   }];
 
   let arguments = (ins Shape_SizeOrIndexType:$lhs, Shape_SizeOrIndexType:$rhs);
   let results = (outs Shape_SizeOrIndexType:$result);
 
   let assemblyFormat = [{
     $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
   }];
def Shape_NumElementsOp : Shape_Op<"num_elements", [NoSideEffect]> {
   let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
 }
 
 def Shape_ReduceOp : Shape_Op<"reduce",
     [SingleBlockImplicitTerminator<"YieldOp">]> {
   let summary = "Returns an expression reduced over a shape or extent tensor";
   let description = [{
     An operation that takes as input a shape or extent tensor, and a number of
-    initial values. This operation has a region/function that is applied
-    repeatedly for every extent of the input. Starting with the initial values,
-    the individual extents are then aggregated as defined by the associated
-    region.
+    initial values. This operation has a region that is applied repeatedly for
+    every extent of the input. Starting with the initial values, the individual
+    extents are then aggregated as defined by the associated region.
 
     Conceptually this op performs the following reduction:
 
     ```
     res[] = init;
     for (int i = 0, i < shape.rank(); i++) {
-      res = fn(i, shape[i], res[0], ..., res[n]);
+      res = reduce(i, shape[i], res[0], ..., res[n]);
     }
     ```
 
-    Where `fn` is provided by the user and the result of the reduce op is the
-    last computed output of the reduce function. As an example, computing the
-    number of elements can be defined as follows:
+    Where `reduce` represents the region attached and the result of the reduce
+    op is the last computed output of the reduce region. As an example, the
+    number of elements can be computed as follows:
 
     ```mlir
     func @reduce(%shape : !shape.shape, %init : !shape.size) -> !shape.size {
       %num_elements = shape.reduce(%shape, %init) -> !shape.size  {
         ^bb0(%index: index, %dim: !shape.size, %acc: !shape.size):
           %updated_acc = "shape.mul"(%acc, %dim) :
             (!shape.size, !shape.size) -> !shape.size
           shape.yield %updated_acc : !shape.size
let extraClassDeclaration = [{
     // This does no checks on the inputs to the AssumingOp.
     static void inlineRegionIntoParent(AssumingOp &op, PatternRewriter &rewriter);
   }];
 
   let hasCanonicalizer = 1;
 }
 
 def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
-                                     [NoSideEffect, ReturnLike, Terminator]> {
+       [NoSideEffect, ReturnLike, Terminator, HasParent<"AssumingOp">]> {
   let summary = "Yield operation";
   let description = [{
-    This yield operation represents a return operation within the assert_and_exec
-    region. The operation takes variable number of operands and produces no
-    results. The operand number and types must match the return signature of
-    the region that contains the operation.
+    This yield operation represents a return operation within the
+    `shape.assuming` operation region. The operation takes variable number of
+    operands and produces no results. The operand number and types must match
+    the number and types of parent `shape.assuming` results.
   }];
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
   let builders = [OpBuilderDAG<(ins), [{ /* nothing to do */ }]>];
 
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
def Shape_ConstWitnessOp : Shape_Op<"const_witness", [ConstantLike, NoSideEffect]> {
   let summary = "An operation that returns a statically known witness value";
   let description = [{
   This operation represents a statically known witness result. This can be
   often used to canonicalize/fold constraint and assuming code that will always
   pass.
 
   ```mlir
   %0 = shape.const_shape [1,2,3]
-  %1 = shape.const_shape [1, 2, 3]
+  %1 = shape.const_shape [1,2,3]
   %w0 = shape.cstr_eq(%0, %1) // Can be folded to "const_witness true"
   %w1 = shape.const_witness true
   %w2 = shape.assuming_all(%w0, %w2) // Can be folded to "const_witness true"
   ```
   }];
   let arguments = (ins BoolAttr:$passing);
   let results = (outs Shape_WitnessType:$result);
 

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

 #include "ShapeCanonicalization.inc"
 }
 
 RankedTensorType shape::getExtentTensorType(MLIRContext *ctx) {
   return RankedTensorType::get({ShapedType::kDynamicSize}, IndexType::get(ctx));
 }
 
 static bool isErrorPropagationPossible(TypeRange operandTypes) {
-  for (Type ty : operandTypes)
-    if (ty.isa<SizeType>() || ty.isa<ShapeType>() || ty.isa<ValueShapeType>())
-      return true;
-  return false;
+  return llvm::any_of(operandTypes, [](Type ty) {
+    return ty.isa<SizeType, ShapeType, ValueShapeType>();
+  });
 }
 
 static LogicalResult verifySizeOrIndexOp(Operation *op) {
   assert(op != nullptr && op->getNumResults() == 1);
   Type resultTy = op->getResultTypes().front();
   if (isErrorPropagationPossible(op->getOperandTypes())) {
     if (!resultTy.isa<SizeType>())
       return op->emitOpError()

 };
 } // namespace
 
 void ShapeDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
       >();
-  addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType,
-           WitnessType>();
+  addTypes<ShapeType, SizeType, ValueShapeType, WitnessType>();
   addInterfaces<ShapeInlinerInterface>();
   // 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,

 }
 
 /// 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());
-  if (keyword == "element")
-    return ElementType::get(getContext());
   if (keyword == "shape")
     return ShapeType::get(getContext());
   if (keyword == "size")
     return SizeType::get(getContext());
   if (keyword == "value_shape")
     return ValueShapeType::get(getContext());
   if (keyword == "witness")
     return WitnessType::get(getContext());
 
   parser.emitError(parser.getNameLoc(), "unknown shape type: ") << keyword;
   return Type();
 }
 
 /// Print a type registered to this dialect.
 void ShapeDialect::printType(Type type, DialectAsmPrinter &os) const {
   TypeSwitch<Type>(type)
-      .Case<ComponentType>([&](Type) { os << "component"; })
-      .Case<ElementType>([&](Type) { os << "element"; })
       .Case<ShapeType>([&](Type) { os << "shape"; })
       .Case<SizeType>([&](Type) { os << "size"; })
       .Case<ValueShapeType>([&](Type) { os << "value_shape"; })
       .Case<WitnessType>([&](Type) { os << "witness"; })
       .Default([](Type) { llvm_unreachable("unexpected 'shape' type kind"); });
 }
 
 //===----------------------------------------------------------------------===//