1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
|
// RUN: mlir-opt %s --transform-interpreter -allow-unregistered-dialect --split-input-file | FileCheck %s
func.func @cast_to_dynamic(%arg0: tensor<10x13xf32>, %arg1: tensor<3x13xf32>) -> tensor<13x13xf32> {
%0 = tensor.concat dim(0) %arg0, %arg1 : (tensor<10x13xf32>, tensor<3x13xf32>) -> tensor<13x13xf32>
func.return %0 : tensor<13x13xf32>
}
func.func private @concat_replacement(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
%funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
%f:2 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%concat = transform.structured.match ops{["tensor.concat"]} in %f#0 : (!transform.any_op) -> !transform.any_op
%ins = transform.get_operand %concat[all] : (!transform.any_op) -> !transform.any_value
%out = transform.get_result %concat[all] : (!transform.any_op) -> !transform.any_value
transform.func.cast_and_call %f#1(%ins) -> %out before %concat {
transform.type_conversion.tensor.cast_shape_dynamic_dims
} : (!transform.any_op, !transform.any_value,
!transform.any_value, !transform.any_op) -> !transform.any_op
transform.apply_dce to %f#0 : !transform.any_op
transform.yield
}
}
// CHECK-LABEL: func.func @cast_to_dynamic
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<10x13xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<3x13xf32>
// CHECK-DAG: %[[CAST0:.+]] = tensor.cast %[[ARG0]] : tensor<10x13xf32> to tensor<?x?xf32>
// CHECK-DAG: %[[CAST1:.+]] = tensor.cast %[[ARG1]] : tensor<3x13xf32> to tensor<?x?xf32>
// CHECK: %[[CALL:.+]] = call @concat_replacement(%[[CAST0]], %[[CAST1]])
// CHECK: %[[CAST_RES:.+]] = tensor.cast %[[CALL]] : tensor<?x?xf32> to tensor<13x13xf32>
// CHECK: return %[[CAST_RES]] : tensor<13x13xf32>
// -----
func.func @cast_to_static(%arg0: tensor<?x?xf32>) -> tensor<?xf32> {
%0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<?x?xf32> into tensor<?xf32>
func.return %0 : tensor<?xf32>
}
func.func private @collapse_replacement(%arg0: tensor<4x5xf32>) -> tensor<20xf32>
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
%funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
%f:2 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%concat = transform.structured.match ops{["tensor.collapse_shape"]} in %f#0 : (!transform.any_op) -> !transform.any_op
%ins = transform.get_operand %concat[all] : (!transform.any_op) -> !transform.any_value
%out = transform.get_result %concat[all] : (!transform.any_op) -> !transform.any_value
transform.func.cast_and_call %f#1(%ins) -> %out before %concat {
transform.type_conversion.tensor.cast_shape_dynamic_dims ignore_dynamic_info
} : (!transform.any_op, !transform.any_value,
!transform.any_value, !transform.any_op) -> !transform.any_op
transform.apply_dce to %f#0 : !transform.any_op
transform.yield
}
}
// CHECK-LABEL: func.func @cast_to_static
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK-DAG: %[[CAST_IN:.+]] = tensor.cast %[[ARG0]] : tensor<?x?xf32> to tensor<4x5xf32>
// CHECK: %[[CALL:.+]] = call @collapse_replacement(%[[CAST_IN]])
// CHECK: %[[CAST_RES:.+]] = tensor.cast %[[CALL]] : tensor<20xf32> to tensor<?xf32>
// CHECK: return %[[CAST_RES]] : tensor<?xf32>
|
