[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 ↗

(On Diff #268456)

Same

339 ↗

(On Diff #268456)

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", [NoSideEffect]> {
   let arguments = (ins Variadic<Index>:$extents);
   let results = (outs Shape_ShapeType:$shape);
 
   let assemblyFormat = "$extents attr-dict";
 
   let hasFolder = 1;
 }
 
-def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", []> {
+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 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]> {
   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.

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 {

mlir/test/Dialect/Shape/ops.mlir

func @const_size() {
   // CHECK: %c1 = shape.const_size 1
   // CHECK: %c2 = shape.const_size 2
   // CHECK: %c2_0 = shape.const_size 2
   %0 = shape.const_size 1
   %1 = shape.const_size 2
   %2 = shape.const_size 2
   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
+}