[MLIR] Add missing traits and assembly format to `shape.from/to_extent_tensor`

[frgossen]

Summary

Add NoSideEffect trait to shape.to_extent_tensor and
shape.from_extent_tensor and defined custom assembly format for the
operations.
Infer return type for shape.from_extent_tensor.

[jpienaar]

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

160

This should no longer be needed

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

303

Same

339

Same

[silvas]

Looks great! (modulo jacques' comments)

[frgossen]

Apply suggestions

[frgossen]

Rebase

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

def Shape_FromExtentsOp : Shape_Op<"from_extents", [
   let arguments = (ins Variadic<Index>:$extents);
   let results = (outs Shape_ShapeType:$shape);
 
   let assemblyFormat = "attr-dict $extents";
 
   let hasFolder = 1;
 }
 
-def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", []> {
+def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", [
+    NoSideEffect,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   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 results = (outs Shape_ShapeType:$result);
+
+  let assemblyFormat = "attr-dict $input `:` type($input)";
 }
 
-def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", []> {
+def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
   let summary = "Creates a dimension tensor from a shape";
-  // TODO: Think more about the error situation. Perhaps factor out the
-  // error detection into a separate op so downstream consumers can control
-  // their error behavior. Then this op would assume that the input has
-  // been properly checked to not be an error (and could thus be a
-  // NoSideEffect op).
   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.
 
-    If the shape represents an error, then this op currently aborts the program.
+    If the shape represents an error, this op's behavior is undefined.
   }];
 
   let arguments = (ins Shape_ShapeType:$input);
   let results = (outs IndexTensor:$result);
 
+  let assemblyFormat = "attr-dict $input `:` type($result)";
+
   let hasFolder = 1;
 }
 
 def Shape_GetExtentOp : Shape_Op<"get_extent",
     [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let summary = "Gets the specified extent from a shape";
   let description = [{
     Gets the extent indexed by `dim` from `shape`.

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

LogicalResult IndexToSizeOp::inferReturnTypes(
   inferredReturnTypes.push_back(SizeType::get(context));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // FromExtentsOp
 //===----------------------------------------------------------------------===//
 
-LogicalResult FromExtentsOp::inferReturnTypes(
-    MLIRContext *context, Optional<Location> location, ValueRange operands,
-    DictionaryAttr attributes, RegionRange regions,
-    SmallVectorImpl<Type> &inferredReturnTypes) {
-  inferredReturnTypes.push_back(ShapeType::get(context));
-  return success();
-}
-
 OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
   if (llvm::any_of(operands, [](Attribute a) { return !a; }))
     return nullptr;
   SmallVector<int64_t, 6> extents;
   for (auto attr : operands)
     extents.push_back(attr.cast<IntegerAttr>().getInt());
   Builder builder(getContext());
   return builder.getIndexTensorAttr(extents);
 }
 
+LogicalResult FromExtentsOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(ShapeType::get(context));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // GetExtentOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 GetExtentOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                               ValueRange operands, DictionaryAttr attributes,
                               RegionRange regions,
OpFoldResult GetExtentOp::fold(ArrayRef<Attribute> operands) {
   uint64_t dimToGet = dim().getLimitedValue();
   // TODO: Constant fold this to some kind of constant error.
   if (dimToGet >= (uint64_t)elements.getNumElements())
     return nullptr;
   return elements.getValue({dimToGet});
 }
 
 //===----------------------------------------------------------------------===//
+// FromExtentTensorOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult FromExtentTensorOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(ShapeType::get(context));
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // NumElementsOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
 
   // Fold only when argument constant.
   Attribute shape = operands[0];
   if (!shape)

mlir/test/Dialect/Shape/canonicalize.mlir

 }
 
 // -----
 // 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>
+  %0 = shape.to_extent_tensor %cs : tensor<2xindex>
   return %0 : tensor<2xindex>
 }
 
 // -----
 // Basic case.
 // CHECK-LABEL: func @f()
 func @f() -> !shape.shape {
   // CHECK: shape.const_shape [3, 5, 11]

mlir/test/Dialect/Shape/ops.mlir

func @test_constraints() {
   %w1 = shape.cstr_eq %0, %1
   %w3 = shape.assuming_all %w0, %w1
   shape.assuming %w3 -> !shape.shape {
     %2 = shape.any %0, %1
     shape.assuming_yield %2 : !shape.shape
   }
   return
 }
+
+func @test_to_extent_tensor(%arg: !shape.shape) -> tensor<3xindex> {
+  %0 = shape.to_extent_tensor %arg : tensor<3xindex>
+  return %0 : tensor<3xindex>
+}
+
+func @test_from_extent_tensor(%arg: tensor<?xindex>) -> !shape.shape {
+  %0 = shape.from_extent_tensor %arg : tensor<?xindex>
+  return %0 : !shape.shape
+}
+