[Mlir-commits] [mlir] c519460 - Allow !shape.size type operands in "shape.from_extents" op.

[Sean Silva]

Author: Jing Pu
Date: 2021-02-24T14:50:07-08:00
New Revision: c519460745ece21238b0ee6c0f9b921366701308

URL: https://github.com/llvm/llvm-project/commit/c519460745ece21238b0ee6c0f9b921366701308
DIFF: https://github.com/llvm/llvm-project/commit/c519460745ece21238b0ee6c0f9b921366701308.diff

LOG: Allow !shape.size type operands in "shape.from_extents" op.

This expands the op to support error propagation and also makes it symmetric with  "shape.get_extent" op.

Reviewed By: silvas

Differential Revision: https://reviews.llvm.org/D97261

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
    mlir/test/Dialect/Shape/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index ecd7f3ae0c95..756d0ec706b8 100644
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -198,10 +198,10 @@ def Shape_FromExtentsOp : Shape_Op<"from_extents", [NoSideEffect]> {
     %s1 = shape.from_extents
     ```
   }];
-  let arguments = (ins Variadic<Index>:$extents);
+  let arguments = (ins Variadic<Shape_SizeOrIndexType>:$extents);
   let results = (outs Shape_ShapeType:$shape);
 
-  let assemblyFormat = "$extents attr-dict";
+  let assemblyFormat = "$extents attr-dict `:` type($extents)";
 
   let hasFolder = 1;
 }

diff  --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index d5bf3f711985..ea895015b1e5 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -173,7 +173,19 @@ func @f() -> !shape.shape {
   %e0 = constant 3 : index
   %e1 = constant 5 : index
   %e2 = constant 11 : index
-  %ret = shape.from_extents %e0, %e1, %e2
+  %ret = shape.from_extents %e0, %e1, %e2 : index, index, index
+  return %ret : !shape.shape
+}
+
+// -----
+
+// fold_const_size
+// CHECK-LABEL: func @fold_const_size()
+func @fold_const_size() -> !shape.shape {
+  // CHECK: shape.const_shape [3, 5] : !shape.shape
+  %e0 = shape.const_size 3
+  %e1 = shape.const_size 5
+  %ret = shape.from_extents %e0, %e1 : !shape.size, !shape.size
   return %ret : !shape.shape
 }
 
@@ -183,7 +195,7 @@ func @f() -> !shape.shape {
 func @no_fold(%arg0: index) -> !shape.shape {
   // CHECK-NOT: shape.const_shape
   %e0 = constant 3 : index
-  %ret = shape.from_extents %e0, %arg0
+  %ret = shape.from_extents %e0, %arg0 : index, index
   return %ret : !shape.shape
 }