[Mlir-commits] [mlir] [DRAFT] Generalize expand_shape to take shape as explicit input (PR #69267)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Thu Oct 26 13:06:56 PDT 2023


================
@@ -1050,39 +1050,83 @@ def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
 
     ```mlir
     // Dimension expansion i -> (i', j') and (k) -> (k')
-    %b = tensor.expand_shape %a [[0, 1], [2]]
-        : tensor<?x?xf32> into tensor<?x?x?xf32>
+    %b = tensor.expand_shape %a [[0, 1], [2]] [%sz0, %sz1, 32]
+        : tensor<?x32xf32> into tensor<?x?x32xf32>
     ```
   }];
+
+  let arguments = (ins AnyTensor:$src, IndexListArrayAttr:$reassociation,
+                       Variadic<Index>:$output_shape,
+                       DenseI64ArrayAttr:$static_output_shape);
+
+  let assemblyFormat = [{
+    $src $reassociation `output_shape`
+    custom<DynamicIndexList>($output_shape, $static_output_shape) attr-dict `:`
+    type($src) `into` type($result)
+  }];
+
   let builders = [
     // Builders using ReassociationIndices.
+    OpBuilder<(ins "Type":$resultType, "Value":$src,
+      "ArrayRef<ReassociationIndices>":$reassociation),
+    [{
+      SmallVector<OpFoldResult> inputShape = getMixedSizes($_builder, $_state, src);
+      SmallVector<OpFoldResult> outputShape;
+      assert(failed(inferOutputShape($_builder, $_state, resultType, reassociation, 
----------------
MaheshRavishankar wrote:

Please take these out of `asserts` . This will be ignored during release builds. SO
```
auto status = inferOutputShape(...)
(void) status
assert(succeeded(status) && ...);
```

https://github.com/llvm/llvm-project/pull/69267


More information about the Mlir-commits mailing list