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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
|
// RUN: mlir-opt -test-transform-dialect-interpreter %s --split-input-file --allow-unregistered-dialect | FileCheck %s
// CHECK-LABEL: func @insert_slice_to_copy
// CHECK-SAME: %[[I:.*]]: tensor<2x3xf32>
// CHECK-SAME: %[[O:.*]]: tensor<?x?xf32>,
// CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[ST0:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[ST1:[0-9a-zA-Z]+]]: index)
func.func @insert_slice_to_copy(
%I : tensor<2x3xf32>, %O : tensor<?x?xf32>,
%off0 : index, %off1 : index,
%sz0 : index, %sz1 : index,
%st0 : index, %st1 : index) -> tensor<?x?xf32> {
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]
// CHECK-SAME: : tensor<?x?xf32> to tensor<2x3xf32>
// CHECK: linalg.copy ins(%[[I]] : tensor<2x3xf32>) outs(%[[EXTRACTED_SLICE]] : tensor<2x3xf32>) -> tensor<2x3xf32>
// CHECK: tensor.insert_slice %{{.*}} into %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]
// CHECK-SAME: : tensor<2x3xf32> into tensor<?x?xf32>
%0 = tensor.insert_slice %I into %O[%off0, %off1] [2, 3] [%st0, %st1]
: tensor<2x3xf32> into tensor<?x?xf32>
return %0 : tensor<?x?xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op
transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">
}
// -----
// CHECK-LABEL: func @insert_slice_to_copy
// CHECK-SAME: %[[I:[0-9a-zA-Z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[O:[0-9a-zA-Z]+]]: tensor<?x?xf32>,
// CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[ST0:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[ST1:[0-9a-zA-Z]+]]: index)
func.func @insert_slice_to_copy(
%I : tensor<?x?xf32>, %O : tensor<?x?xf32>,
%off0 : index, %off1 : index,
%sz0 : index, %sz1 : index,
%st0 : index, %st1 : index) -> tensor<?x?xf32> {
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[O]][%[[OFF0]], %[[OFF1]]] [%[[SZ0]], %[[SZ1]]] [1, 1]
// CHECK-SAME: : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: linalg.copy ins(%[[I]] : tensor<?x?xf32>) outs(%[[EXTRACTED_SLICE]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: tensor.insert_slice %{{.*}} into %[[O]][%[[OFF0]], %[[OFF1]]] [%[[SZ0]], %[[SZ1]]] [1, 1]
// CHECK-SAME: : tensor<?x?xf32> into tensor<?x?xf32>
%0 = tensor.insert_slice %I into %O[%off0, %off1] [%sz0, %sz1] [1, 1]
: tensor<?x?xf32> into tensor<?x?xf32>
return %0 : tensor<?x?xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op
transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">
}
// -----
// CHECK-LABEL: func @insert_slice_to_copy
// CHECK-SAME: %[[I:.*]]: tensor<2x3xf32>
// CHECK-SAME: %[[O:.*]]: tensor<?x?xf32>,
// CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[ST0:[0-9a-zA-Z]+]]: index,
// CHECK-SAME: %[[ST1:[0-9a-zA-Z]+]]: index)
func.func @insert_slice_to_copy(
%I : tensor<2x3xf32>, %O : tensor<?x?xf32>,
%off0 : index, %off1 : index,
%sz0 : index, %sz1 : index,
%st0 : index, %st1 : index) -> tensor<?x?xf32> {
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]
// CHECK-SAME: : tensor<?x?xf32> to tensor<2x3xf32>
// CHECK: linalg.copy ins(%[[I]] : tensor<2x3xf32>) outs(%[[EXTRACTED_SLICE]] : tensor<2x3xf32>) -> tensor<2x3xf32>
// CHECK-NOT: linalg.copy
// CHECK: tensor.insert_slice %{{.*}} into %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]
// CHECK-SAME: : tensor<2x3xf32> into tensor<?x?xf32>
%extracted_slice = tensor.extract_slice %O[%off0, %off1] [2, 3] [%st0, %st1]
: tensor<?x?xf32> to tensor<2x3xf32>
%0 = linalg.copy ins(%I : tensor<2x3xf32>) outs(%extracted_slice
: tensor<2x3xf32>) -> tensor<2x3xf32>
%inserted_slice = tensor.insert_slice %0 into %O[%off0, %off1] [2, 3] [%st0, %st1]
: tensor<2x3xf32> into tensor<?x?xf32>
return %inserted_slice : tensor<?x?xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op
transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">
}
// -----
// CHECK-LABEL: func @parallel_insert_slice_to_copy
func.func @parallel_insert_slice_to_copy(%out : tensor<?x?xf32>, %sz0: index, %sz1: index) {
%0 = scf.forall (%arg0, %arg1) in (27, 8) shared_outs(%arg2 = %out) -> (tensor<?x?xf32>) {
%t = "make_me_a_tensor"() : () -> (tensor<?x?xf32> )
// CHECK: tensor.extract_slice
// CHECK: linalg.copy
// CHECK: scf.forall.in_parallel
// CHECK: tensor.parallel_insert_slice
scf.forall.in_parallel {
tensor.parallel_insert_slice %t into %arg2[0, 0] [%sz0, %sz1] [1, 1]
: tensor<?x?xf32> into tensor<?x?xf32>
}
}
return
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg1
: (!transform.any_op) -> !transform.any_op
%1 = transform.structured.insert_slice_to_copy %0
: (!transform.any_op) -> !transform.any_op
transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">
}
|
