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