add output_shape to memref/tensor (#89)

Author: makslevental

mlir/extras/dialects/ext/memref.py

@@ -210,7 +210,13 @@ def expand_shape(
 
     return MemRef(
         memref.expand_shape(
-            T.memref(*result_shape, inp.dtype), inp, reassoc_list, loc=loc, ip=ip
+            T.memref(*result_shape, inp.dtype),
+            inp,
+            reassoc_list,
+            output_shape=[],
+            static_output_shape=result_shape,
+            loc=loc,
+            ip=ip,
         )
     )
 

mlir/extras/dialects/ext/tensor.py

@@ -325,6 +325,8 @@ def expand_dims(
             RankedTensorType.get(result_shape, inp.dtype),
             inp,
             reassoc_list,
+            output_shape=[],
+            static_output_shape=result_shape,
             loc=loc,
             ip=ip,
         )

tests/test_memref.py

@@ -343,27 +343,27 @@ def test_none_indices(ctx: MLIRContext):
         """\
     module {
       %alloc = memref.alloc() : memref<10x22x333x4444xi32>
-      %expand_shape = memref.expand_shape %alloc [[0, 1], [2], [3], [4]] : memref<10x22x333x4444xi32> into memref<1x10x22x333x4444xi32>
+      %expand_shape = memref.expand_shape %alloc [[0, 1], [2], [3], [4]] output_shape [1, 10, 22, 333, 4444] : memref<10x22x333x4444xi32> into memref<1x10x22x333x4444xi32>
       %subview = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_0 = memref.expand_shape %subview [[0, 1], [2], [3], [4]] : memref<10x22x333x4444xi32> into memref<10x1x22x333x4444xi32>
+      %expand_shape_0 = memref.expand_shape %subview [[0, 1], [2], [3], [4]] output_shape [10, 1, 22, 333, 4444] : memref<10x22x333x4444xi32> into memref<10x1x22x333x4444xi32>
       %subview_1 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_2 = memref.expand_shape %subview_1 [[0, 1, 2], [3], [4], [5]] : memref<10x22x333x4444xi32> into memref<1x10x1x22x333x4444xi32>
+      %expand_shape_2 = memref.expand_shape %subview_1 [[0, 1, 2], [3], [4], [5]] output_shape [1, 10, 1, 22, 333, 4444] : memref<10x22x333x4444xi32> into memref<1x10x1x22x333x4444xi32>
       %subview_3 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_4 = memref.expand_shape %subview_3 [[0], [1, 2], [3], [4]] : memref<10x22x333x4444xi32> into memref<10x22x1x333x4444xi32>
+      %expand_shape_4 = memref.expand_shape %subview_3 [[0], [1, 2], [3], [4]] output_shape [10, 22, 1, 333, 4444] : memref<10x22x333x4444xi32> into memref<10x22x1x333x4444xi32>
       %subview_5 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_6 = memref.expand_shape %subview_5 [[0], [1], [2, 3], [4]] : memref<10x22x333x4444xi32> into memref<10x22x333x1x4444xi32>
+      %expand_shape_6 = memref.expand_shape %subview_5 [[0], [1], [2, 3], [4]] output_shape [10, 22, 333, 1, 4444] : memref<10x22x333x4444xi32> into memref<10x22x333x1x4444xi32>
       %subview_7 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_8 = memref.expand_shape %subview_7 [[0], [1], [2], [3, 4]] : memref<10x22x333x4444xi32> into memref<10x22x333x4444x1xi32>
+      %expand_shape_8 = memref.expand_shape %subview_7 [[0], [1], [2], [3, 4]] output_shape [10, 22, 333, 4444, 1] : memref<10x22x333x4444xi32> into memref<10x22x333x4444x1xi32>
       %subview_9 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_10 = memref.expand_shape %subview_9 [[0], [1], [2], [3, 4]] : memref<10x22x333x4444xi32> into memref<10x22x333x4444x1xi32>
+      %expand_shape_10 = memref.expand_shape %subview_9 [[0], [1], [2], [3, 4]] output_shape [10, 22, 333, 4444, 1] : memref<10x22x333x4444xi32> into memref<10x22x333x4444x1xi32>
       %subview_11 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_12 = memref.expand_shape %subview_11 [[0, 1], [2], [3], [4, 5]] : memref<10x22x333x4444xi32> into memref<10x1x22x333x4444x1xi32>
+      %expand_shape_12 = memref.expand_shape %subview_11 [[0, 1], [2], [3], [4, 5]] output_shape [10, 1, 22, 333, 4444, 1] : memref<10x22x333x4444xi32> into memref<10x1x22x333x4444x1xi32>
       %subview_13 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_14 = memref.expand_shape %subview_13 [[0, 1], [2, 3], [4], [5, 6]] : memref<10x22x333x4444xi32> into memref<10x1x22x1x333x4444x1xi32>
+      %expand_shape_14 = memref.expand_shape %subview_13 [[0, 1], [2, 3], [4], [5, 6]] output_shape [10, 1, 22, 1, 333, 4444, 1] : memref<10x22x333x4444xi32> into memref<10x1x22x1x333x4444x1xi32>
       %subview_15 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_16 = memref.expand_shape %subview_15 [[0, 1], [2, 3], [4, 5], [6, 7]] : memref<10x22x333x4444xi32> into memref<10x1x22x1x333x1x4444x1xi32>
+      %expand_shape_16 = memref.expand_shape %subview_15 [[0, 1], [2, 3], [4, 5], [6, 7]] output_shape [10, 1, 22, 1, 333, 1, 4444, 1] : memref<10x22x333x4444xi32> into memref<10x1x22x1x333x1x4444x1xi32>
       %subview_17 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
-      %expand_shape_18 = memref.expand_shape %subview_17 [[0, 1, 2], [3, 4], [5, 6], [7, 8]] : memref<10x22x333x4444xi32> into memref<1x10x1x22x1x333x1x4444x1xi32>
+      %expand_shape_18 = memref.expand_shape %subview_17 [[0, 1, 2], [3, 4], [5, 6], [7, 8]] output_shape [1, 10, 1, 22, 1, 333, 1, 4444, 1] : memref<10x22x333x4444xi32> into memref<1x10x1x22x1x333x1x4444x1xi32>
       %subview_19 = memref.subview %alloc[0, 0, 0, 0] [10, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x333x4444xi32>
     }
     """

tests/test_tensor.py

@@ -223,17 +223,17 @@ def test_none_indices(ctx: MLIRContext):
         """\
     module {
       %0 = tensor.empty() : tensor<10x22x333x4444xi32>
-      %expanded = tensor.expand_shape %0 [[0, 1], [2], [3], [4]] : tensor<10x22x333x4444xi32> into tensor<1x10x22x333x4444xi32>
-      %expanded_0 = tensor.expand_shape %0 [[0, 1], [2], [3], [4]] : tensor<10x22x333x4444xi32> into tensor<10x1x22x333x4444xi32>
-      %expanded_1 = tensor.expand_shape %0 [[0, 1, 2], [3], [4], [5]] : tensor<10x22x333x4444xi32> into tensor<1x10x1x22x333x4444xi32>
-      %expanded_2 = tensor.expand_shape %0 [[0], [1, 2], [3], [4]] : tensor<10x22x333x4444xi32> into tensor<10x22x1x333x4444xi32>
-      %expanded_3 = tensor.expand_shape %0 [[0], [1], [2, 3], [4]] : tensor<10x22x333x4444xi32> into tensor<10x22x333x1x4444xi32>
-      %expanded_4 = tensor.expand_shape %0 [[0], [1], [2], [3, 4]] : tensor<10x22x333x4444xi32> into tensor<10x22x333x4444x1xi32>
-      %expanded_5 = tensor.expand_shape %0 [[0], [1], [2], [3, 4]] : tensor<10x22x333x4444xi32> into tensor<10x22x333x4444x1xi32>
-      %expanded_6 = tensor.expand_shape %0 [[0, 1], [2], [3], [4, 5]] : tensor<10x22x333x4444xi32> into tensor<10x1x22x333x4444x1xi32>
-      %expanded_7 = tensor.expand_shape %0 [[0, 1], [2, 3], [4], [5, 6]] : tensor<10x22x333x4444xi32> into tensor<10x1x22x1x333x4444x1xi32>
-      %expanded_8 = tensor.expand_shape %0 [[0, 1], [2, 3], [4, 5], [6, 7]] : tensor<10x22x333x4444xi32> into tensor<10x1x22x1x333x1x4444x1xi32>
-      %expanded_9 = tensor.expand_shape %0 [[0, 1, 2], [3, 4], [5, 6], [7, 8]] : tensor<10x22x333x4444xi32> into tensor<1x10x1x22x1x333x1x4444x1xi32>
+      %expanded = tensor.expand_shape %0 [[0, 1], [2], [3], [4]] output_shape [1, 10, 22, 333, 4444] : tensor<10x22x333x4444xi32> into tensor<1x10x22x333x4444xi32>
+      %expanded_0 = tensor.expand_shape %0 [[0, 1], [2], [3], [4]] output_shape [10, 1, 22, 333, 4444] : tensor<10x22x333x4444xi32> into tensor<10x1x22x333x4444xi32>
+      %expanded_1 = tensor.expand_shape %0 [[0, 1, 2], [3], [4], [5]] output_shape [1, 10, 1, 22, 333, 4444] : tensor<10x22x333x4444xi32> into tensor<1x10x1x22x333x4444xi32>
+      %expanded_2 = tensor.expand_shape %0 [[0], [1, 2], [3], [4]] output_shape [10, 22, 1, 333, 4444] : tensor<10x22x333x4444xi32> into tensor<10x22x1x333x4444xi32>
+      %expanded_3 = tensor.expand_shape %0 [[0], [1], [2, 3], [4]] output_shape [10, 22, 333, 1, 4444] : tensor<10x22x333x4444xi32> into tensor<10x22x333x1x4444xi32>
+      %expanded_4 = tensor.expand_shape %0 [[0], [1], [2], [3, 4]] output_shape [10, 22, 333, 4444, 1] : tensor<10x22x333x4444xi32> into tensor<10x22x333x4444x1xi32>
+      %expanded_5 = tensor.expand_shape %0 [[0], [1], [2], [3, 4]] output_shape [10, 22, 333, 4444, 1] : tensor<10x22x333x4444xi32> into tensor<10x22x333x4444x1xi32>
+      %expanded_6 = tensor.expand_shape %0 [[0, 1], [2], [3], [4, 5]] output_shape [10, 1, 22, 333, 4444, 1] : tensor<10x22x333x4444xi32> into tensor<10x1x22x333x4444x1xi32>
+      %expanded_7 = tensor.expand_shape %0 [[0, 1], [2, 3], [4], [5, 6]] output_shape [10, 1, 22, 1, 333, 4444, 1] : tensor<10x22x333x4444xi32> into tensor<10x1x22x1x333x4444x1xi32>
+      %expanded_8 = tensor.expand_shape %0 [[0, 1], [2, 3], [4, 5], [6, 7]] output_shape [10, 1, 22, 1, 333, 1, 4444, 1] : tensor<10x22x333x4444xi32> into tensor<10x1x22x1x333x1x4444x1xi32>
+      %expanded_9 = tensor.expand_shape %0 [[0, 1, 2], [3, 4], [5, 6], [7, 8]] output_shape [1, 10, 1, 22, 1, 333, 1, 4444, 1] : tensor<10x22x333x4444xi32> into tensor<1x10x1x22x1x333x1x4444x1xi32>
     }
     """
     )