/
transform-ops.mlir
125 lines (103 loc) · 3.79 KB
/
transform-ops.mlir
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
// RUN: mlir-opt --test-transform-dialect-interpreter %s -split-input-file -verify-diagnostics | FileCheck %s
// Test One-Shot Bufferize.
transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation):
sequence %arg0 {
^bb0(%arg1: !pdl.operation):
%0 = pdl_match @pdl_target in %arg1
transform.bufferization.one_shot_bufferize %0
{target_is_module = false}
}
pdl.pattern @pdl_target : benefit(1) {
%0 = operation "func.func"
rewrite %0 with "transform.dialect"
}
}
// CHECK-LABEL: func @test_function(
// CHECK-SAME: %[[A:.*]]: tensor<?xf32>
func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
%c0 = arith.constant 0 : index
// CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
// CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
// CHECK: memref.copy %[[A_memref]], %[[alloc]]
// CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
// CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
%0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
// CHECK: memref.dealloc %[[alloc]]
// CHECK: return %[[res_tensor]]
return %0 : tensor<?xf32>
}
// -----
// Test analysis of One-Shot Bufferize only.
transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation):
sequence %arg0 {
^bb0(%arg1: !pdl.operation):
%0 = pdl_match @pdl_target in %arg1
transform.bufferization.one_shot_bufferize %0
{target_is_module = false, test_analysis_only = true}
}
pdl.pattern @pdl_target : benefit(1) {
%0 = operation "func.func"
rewrite %0 with "transform.dialect"
}
}
// CHECK-LABEL: func @test_function_analysis(
// CHECK-SAME: %[[A:.*]]: tensor<?xf32>
func.func @test_function_analysis(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
%c0 = arith.constant 0 : index
// CHECK: vector.transfer_write
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "false", "none"]}
// CHECK-SAME: tensor<?xf32>
%0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
return %0 : tensor<?xf32>
}
// -----
// Test One-Shot Bufferize transform failure with an unknown op. This would be
// allowed with `allow_unknown_ops`.
transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation):
sequence %arg0 {
^bb0(%arg1: !pdl.operation):
%0 = pdl_match @pdl_target in %arg1
// expected-error @+1 {{bufferization failed}}
transform.bufferization.one_shot_bufferize %0 {target_is_module = false}
}
pdl.pattern @pdl_target : benefit(1) {
%0 = operation "func.func"
rewrite %0 with "transform.dialect"
}
}
func.func @test_unknown_op_failure() -> (tensor<?xf32>) {
// expected-error @+1 {{op was not bufferized}}
%0 = "test.dummy_op"() : () -> (tensor<?xf32>)
return %0 : tensor<?xf32>
}
// -----
// Test One-Shot Bufferize transform failure with a module op.
transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation):
sequence %arg0 {
^bb0(%arg1: !pdl.operation):
// %arg1 is the module
transform.bufferization.one_shot_bufferize %arg1
}
}
module {
// CHECK-LABEL: func @test_function(
// CHECK-SAME: %[[A:.*]]: tensor<?xf32>
func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
%c0 = arith.constant 0 : index
// CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
// CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
// CHECK: memref.copy %[[A_memref]], %[[alloc]]
// CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
// CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
%0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
// CHECK: memref.dealloc %[[alloc]]
// CHECK: return %[[res_tensor]]
return %0 : tensor<?xf32>
}
}