1// RUN: mlir-opt --test-transform-dialect-interpreter %s -split-input-file -verify-diagnostics | FileCheck %s
2
3// Test One-Shot Bufferize.
4
5transform.with_pdl_patterns {
6^bb0(%arg0: !pdl.operation):
7  sequence %arg0 {
8    ^bb0(%arg1: !pdl.operation):
9      %0 = pdl_match @pdl_target in %arg1
10      transform.bufferization.one_shot_bufferize %0
11          {target_is_module = false}
12  }
13
14  pdl.pattern @pdl_target : benefit(1) {
15    %0 = operation "func.func"
16    rewrite %0 with "transform.dialect"
17  }
18}
19
20// CHECK-LABEL: func @test_function(
21//  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
22func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
23  %c0 = arith.constant 0 : index
24
25  // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
26  // CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]
27  // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
28  // CHECK: memref.copy %[[A_memref]], %[[alloc]]
29  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
30  // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
31  %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
32
33  // CHECK: memref.dealloc %[[alloc]]
34  // CHECK: return %[[res_tensor]]
35  return %0 : tensor<?xf32>
36}
37
38// -----
39
40// Test analysis of One-Shot Bufferize only.
41
42transform.with_pdl_patterns {
43^bb0(%arg0: !pdl.operation):
44  sequence %arg0 {
45    ^bb0(%arg1: !pdl.operation):
46      %0 = pdl_match @pdl_target in %arg1
47      transform.bufferization.one_shot_bufferize %0
48          {target_is_module = false, test_analysis_only = true}
49  }
50
51  pdl.pattern @pdl_target : benefit(1) {
52    %0 = operation "func.func"
53    rewrite %0 with "transform.dialect"
54  }
55}
56
57// CHECK-LABEL: func @test_function_analysis(
58//  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
59func.func @test_function_analysis(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
60  %c0 = arith.constant 0 : index
61  // CHECK: vector.transfer_write
62  // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false", "none"]}
63  // CHECK-SAME: tensor<?xf32>
64  %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
65  return %0 : tensor<?xf32>
66}
67
68// -----
69
70// Test One-Shot Bufferize transform failure with an unknown op. This would be
71// allowed with `allow_unknown_ops`.
72
73transform.with_pdl_patterns {
74^bb0(%arg0: !pdl.operation):
75  sequence %arg0 {
76    ^bb0(%arg1: !pdl.operation):
77      %0 = pdl_match @pdl_target in %arg1
78      // expected-error @+1 {{bufferization failed}}
79      transform.bufferization.one_shot_bufferize %0 {target_is_module = false}
80  }
81
82  pdl.pattern @pdl_target : benefit(1) {
83    %0 = operation "func.func"
84    rewrite %0 with "transform.dialect"
85  }
86}
87
88func.func @test_unknown_op_failure() -> (tensor<?xf32>) {
89  // expected-error @+1 {{op was not bufferized}}
90  %0 = "test.dummy_op"() : () -> (tensor<?xf32>)
91  return %0 : tensor<?xf32>
92}
93
94// -----
95
96// Test One-Shot Bufferize transform failure with a module op.
97
98transform.with_pdl_patterns {
99^bb0(%arg0: !pdl.operation):
100  sequence %arg0 {
101    ^bb0(%arg1: !pdl.operation):
102      // %arg1 is the module
103      transform.bufferization.one_shot_bufferize %arg1
104  }
105}
106
107module {
108  // CHECK-LABEL: func @test_function(
109  //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
110  func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
111    %c0 = arith.constant 0 : index
112
113    // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
114    // CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]
115    // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
116    // CHECK: memref.copy %[[A_memref]], %[[alloc]]
117    // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
118    // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]
119    %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
120
121    // CHECK: memref.dealloc %[[alloc]]
122    // CHECK: return %[[res_tensor]]
123    return %0 : tensor<?xf32>
124  }
125}