From c519460745ece21238b0ee6c0f9b921366701308 Mon Sep 17 00:00:00 2001
From: Jing Pu <jingpu@google.com>
Date: Wed, 24 Feb 2021 14:35:23 -0800
Subject: [PATCH] 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
---
 mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td |  4 ++--
 mlir/test/Dialect/Shape/canonicalize.mlir      | 16 ++++++++++++++--
 2 files changed, 16 insertions(+), 4 deletions(-)

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
 }