================
@@ -464,6 +464,129 @@ module attributes {transform.with_named_sequence} {
// -----
+// Check that we can lower unpack with dynamic dimensions in the input and
destination.
+// CHECK-LABEL: func.func @unpack_with_dynamic_input_dest(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?x8x16xf32>, %[[ARG1:.*]]:
tensor<?x?xf32>)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[DIM00:.*]] = tensor.dim %[[ARG0]], %[[C0]]
+// CHECK-DAG: %[[DIM01:.*]] = tensor.dim %[[ARG0]], %[[C1]]
+// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[DIM00]], %[[DIM01]]) :
tensor<?x8x?x16xf32>
+// CHECK: %[[TRAN:.*]] = linalg.transpose
+// CHECK-SAME: ins(%[[ARG0]] : tensor<?x?x8x16xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<?x8x?x16xf32>)
+// CHECK-SAME: permutation = [0, 2, 1, 3]
+// CHECK: %[[CLP:.*]] = tensor.collapse_shape %[[TRAN]] {{\[}}[0, 1], [2,
3]]
+// CHECK-SAME: : tensor<?x8x?x16xf32> into tensor<?x?xf32>
+// CHECK: %[[DIM10:.*]] = tensor.dim %[[ARG1]], %[[C0]] : tensor<?x?xf32>
+// CHECK: %[[DIM11:.*]] = tensor.dim %[[ARG1]], %[[C1]] : tensor<?x?xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[CLP]][0, 0]
[%[[DIM10]], %[[DIM11]]] [1, 1]
+// CHECK-SAME: : tensor<?x?xf32> to tensor<?x?xf32>
+// CHECK: linalg.copy ins(%[[SLICE]] : tensor<?x?xf32>)
+// CHECK-SAME: outs(%[[ARG1]] : tensor<?x?xf32>)
+func.func @unpack_with_dynamic_input_dest(%arg0: tensor<?x?x8x16xf32>, %arg1:
tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %unpack = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8,
16] into %arg1 : tensor<?x?x8x16xf32> -> tensor<?x?xf32>
+ return %unpack : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%module_op: !transform.any_op
{transform.readonly}) {
+ %unpack = transform.structured.match ops{["tensor.unpack"]} in %module_op
+ : (!transform.any_op) -> !transform.op<"tensor.unpack">
+ transform.structured.lower_unpack %unpack :
(!transform.op<"tensor.unpack">)
+ -> (!transform.op<"tensor.empty">,
+ !transform.op<"linalg.transpose">,
+ !transform.op<"tensor.collapse_shape">,
+ !transform.op<"tensor.extract_slice">)
+ transform.yield
+ }
+}
+
+// -----
+
+// Check that we can lower unpack with dynamic dimensions in the input,
destination, inner_tiles.
+// CHECK-LABEL: func.func @unpack_fully_dynamic(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?x?x?xf32>, %[[ARG1:.*]]:
tensor<?x?xf32>, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+// CHECK-DAG: %[[DIM00:.*]] = tensor.dim %[[ARG0]], %[[C0]]
+// CHECK-DAG: %[[DIM01:.*]] = tensor.dim %[[ARG0]], %[[C1]]
+// CHECK-DAG: %[[DIM02:.*]] = tensor.dim %[[ARG0]], %[[C2]]
+// CHECK-DAG: %[[DIM03:.*]] = tensor.dim %[[ARG0]], %[[C3]]
+// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[DIM00]], %[[DIM02]],
%[[DIM01]], %[[DIM03]]) : tensor<?x?x?x?xf32>
+// CHECK: %[[TRAN:.*]] = linalg.transpose
+// CHECK-SAME: ins(%[[ARG0]] : tensor<?x?x?x?xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<?x?x?x?xf32>)
+// CHECK-SAME: permutation = [0, 2, 1, 3]
+// CHECK: %[[CLP:.*]] = tensor.collapse_shape %[[TRAN]] {{\[}}[0, 1], [2,
3]]
+// CHECK-SAME: : tensor<?x?x?x?xf32> into tensor<?x?xf32>
+// CHECK: %[[DIM10:.*]] = tensor.dim %[[ARG1]], %[[C0]] : tensor<?x?xf32>
+// CHECK: %[[DIM11:.*]] = tensor.dim %[[ARG1]], %[[C1]] : tensor<?x?xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[CLP]][0, 0]
[%[[DIM10]], %[[DIM11]]] [1, 1]
+// CHECK-SAME: : tensor<?x?xf32> to tensor<?x?xf32>
+// CHECK: linalg.copy ins(%[[SLICE]] : tensor<?x?xf32>)
+// CHECK-SAME: outs(%[[ARG1]] : tensor<?x?xf32>)
+func.func @unpack_fully_dynamic(%source: tensor<?x?x?x?xf32>, %dest:
tensor<?x?xf32>, %tile_n : index, %tile_m : index) -> tensor<?x?xf32> {
+ %0 = tensor.unpack %source inner_dims_pos = [0, 1] inner_tiles = [%tile_n,
%tile_m] into %dest : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%module_op: !transform.any_op
{transform.readonly}) {
+ %unpack = transform.structured.match ops{["tensor.unpack"]} in %module_op
+ : (!transform.any_op) -> !transform.op<"tensor.unpack">
+ transform.structured.lower_unpack %unpack :
(!transform.op<"tensor.unpack">)
+ -> (!transform.op<"tensor.empty">,
+ !transform.op<"linalg.transpose">,
+ !transform.op<"tensor.collapse_shape">,
+ !transform.op<"tensor.extract_slice">)
+ transform.yield
+ }
+}
+
+// -----
+
+// Check that we can lower unpack "as unpad" with dynamic dims
----------------
chelini wrote:
nit. '.' at the end of the sentence.
https://github.com/llvm/llvm-project/pull/75494
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
