239 lines · plain
1// RUN: mlir-opt --transform-interpreter %s -split-input-file -verify-diagnostics | FileCheck %s2 3// Test One-Shot Bufferize.4 5module attributes {transform.with_named_sequence} {6 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {7 %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op8 %1 = transform.bufferization.one_shot_bufferize %0 : (!transform.any_op) -> !transform.any_op9 transform.yield10 }11}12 13// CHECK-LABEL: func @test_function(14// CHECK-SAME: %[[A:.*]]: tensor<?xf32>15func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {16 %c0 = arith.constant 0 : index17 18 // CHECK: %[[A_memref:.*]] = bufferization.to_buffer %[[A]]19 // CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]20 // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])21 // CHECK: memref.copy %[[A_memref]], %[[alloc]]22 // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]23 // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]24 %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>25 26 // CHECK: return %[[res_tensor]]27 return %0 : tensor<?xf32>28}29 30// -----31 32// Emit linalg.copy instead of memref.copy.33 34module attributes {transform.with_named_sequence} {35 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {36 %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op37 %1 = transform.bufferization.one_shot_bufferize %0 {memcpy_op = "linalg.copy"} : (!transform.any_op) -> !transform.any_op38 transform.yield39 }40}41 42// CHECK-LABEL: func @test_function(43// CHECK-SAME: %[[A:.*]]: tensor<?xf32>44// CHECK-NOT: memref.copy45func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {46 %c0 = arith.constant 0 : index47 48 // CHECK: %[[A_memref:.*]] = bufferization.to_buffer %[[A]]49 // CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]50 // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])51 // CHECK: linalg.copy ins(%[[A_memref]] : memref<{{.*}}>) outs(%[[alloc]]52 // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]53 // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]54 %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>55 56 // CHECK: return %[[res_tensor]]57 return %0 : tensor<?xf32>58}59 60// -----61 62// Test analysis of One-Shot Bufferize only.63 64module attributes {transform.with_named_sequence} {65 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {66 %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op67 %1 = transform.bufferization.one_shot_bufferize %068 {test_analysis_only = true} : (!transform.any_op) -> !transform.any_op69 transform.yield70 }71}72 73// CHECK-LABEL: func @test_function_analysis(74// CHECK-SAME: %[[A:.*]]: tensor<?xf32>75func.func @test_function_analysis(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {76 %c0 = arith.constant 0 : index77 // CHECK: vector.transfer_write78 // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false", "none"]}79 // CHECK-SAME: tensor<?xf32>80 %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>81 return %0 : tensor<?xf32>82}83 84// -----85 86// Test One-Shot Bufferize transform failure with an unknown op. This would be87// allowed with `allow_unknown_ops`.88 89module attributes {transform.with_named_sequence} {90 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {91 %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op92 // expected-error @+1 {{bufferization failed}}93 %1 = transform.bufferization.one_shot_bufferize %0 : (!transform.any_op) -> !transform.any_op94 transform.yield95 }96}97 98func.func @test_unknown_op_failure() -> (tensor<?xf32>) {99 // expected-error @+1 {{op was not bufferized}}100 %0 = "test.dummy_op"() : () -> (tensor<?xf32>)101 return %0 : tensor<?xf32>102}103 104// -----105 106module attributes {transform.with_named_sequence} {107 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.consumed}) {108 // %arg1 is the module109 %0 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op110 transform.yield111 }112}113 114// CHECK-LABEL: func @test_function(115// CHECK-SAME: %[[A:.*]]: tensor<?xf32>116func.func @test_function(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {117 %c0 = arith.constant 0 : index118 119 // CHECK: %[[A_memref:.*]] = bufferization.to_buffer %[[A]]120 // CHECK: %[[dim:.*]] = memref.dim %[[A_memref]]121 // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])122 // CHECK: memref.copy %[[A_memref]], %[[alloc]]123 // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]124 // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]]125 %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>126 127 // CHECK: return %[[res_tensor]]128 return %0 : tensor<?xf32>129}130 131// -----132 133// Test we use identity layout at function boundaries.134 135module attributes {transform.with_named_sequence} {136 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.consumed}) {137 %0 = transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %arg1138 { bufferize_function_boundaries = true } : (!transform.any_op) -> !transform.any_op139 transform.yield140 }141}142 143// CHECK: func.func @matmul(144// CHECK-SAME: %[[A:.*]]: memref<12x9xf32>,145// CHECK-SAME: %[[B:.*]]: memref<9x6xf32>,146// CHECK-SAME: %[[C:.*]]: memref<12x6xf32>) -> memref<12x6xf32> {147func.func @matmul(%A: tensor<12x9xf32>, %B: tensor<9x6xf32>, %C: tensor<12x6xf32>) -> tensor<12x6xf32> {148 // CHECK: linalg.matmul ins(%[[A]], %[[B]] : memref<12x9xf32>, memref<9x6xf32>) outs(%[[C]] : memref<12x6xf32>)149 %D = linalg.matmul ins(%A, %B: tensor<12x9xf32>, tensor<9x6xf32>) outs(%C: tensor<12x6xf32>) -> tensor<12x6xf32>150 // CHECK: return %[[C]] : memref<12x6xf32>151 return %D : tensor<12x6xf32>152}153 154// -----155 156module attributes {transform.with_named_sequence} {157 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {158 %0 = transform.structured.match ops{["tensor.empty"]} in %arg1 : (!transform.any_op) -> !transform.any_op159 %1 = transform.cast %0 : !transform.any_op to !transform.op<"tensor.empty">160 transform.bufferization.empty_tensor_to_alloc_tensor %1 : (!transform.op<"tensor.empty">) -> !transform.op<"bufferization.alloc_tensor">161 transform.yield162 }163}164 165// Expect `bufferization.empty_tensor_to_alloc_tensor` to replace the tensor.empty.166func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {167 // CHECK: bufferization.alloc_tensor168 %0 = tensor.empty() : tensor<2x2xf32>169 return %0 : tensor<2x2xf32>170}171 172// -----173 174module attributes {transform.with_named_sequence} {175 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {176 %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op177 transform.bufferization.eliminate_empty_tensors %0 : !transform.any_op178 transform.yield179 }180}181 182// CHECK-LABEL: func @empty_tensor_elimination(183// CHECK: tensor.extract_slice184// CHECK: linalg.fill185// CHECK: tensor.insert_slice186func.func @empty_tensor_elimination(187 %t: tensor<10xf32>, %f: f32) -> tensor<10xf32> {188 %0 = tensor.empty() : tensor<5xf32>189 %1 = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>190 %2 = tensor.insert_slice %1 into %t [1][5][1]191 : tensor<5xf32> into tensor<10xf32>192 return %2 : tensor<10xf32>193}194 195// -----196 197module attributes {transform.with_named_sequence} {198 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {199 %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op200 transform.bufferization.buffer_loop_hoisting %0 : !transform.any_op201 transform.yield202 }203}204 205// CHECK-LABEL: func @buffer_loop_hoisting(206// CHECK: memref.alloca207// CHECK: scf.for208// CHECK: memref.store209func.func @buffer_loop_hoisting(%lb: index, %ub: index, %step: index, %f: f32, %pos: index) {210 scf.for %iv = %lb to %ub step %step {211 %0 = memref.alloca() : memref<5xf32>212 memref.store %f, %0[%pos] : memref<5xf32>213 }214 return215}216 217// -----218 219module attributes {transform.with_named_sequence} {220 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {221 %alloc_tensor = transform.structured.match ops{["bufferization.alloc_tensor"]} in %arg1222 : (!transform.any_op) -> !transform.op<"bufferization.alloc_tensor">223 %2, %new = transform.structured.bufferize_to_allocation %alloc_tensor224 {alloc_op = "memref.alloca"}225 : !transform.op<"bufferization.alloc_tensor">226 transform.yield227 }228}229 230// Expect `bufferization.bufferize_to_allocation` to create an alloc.231// CHECK-LABEL: func.func @empty_to_tensor_alloc()232func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {233 // CHECK-NEXT: %[[alloca:.*]] = memref.alloca() : memref<2x2xf32>234 // CHECK-NEXT: %[[tensor:.*]] = bufferization.to_tensor %[[alloca]] restrict writable : memref<2x2xf32>235 // CHECK-NEXT: return %[[tensor]] : tensor<2x2xf32>236 %0 = bufferization.alloc_tensor() : tensor<2x2xf32>237 return %0 : tensor<2x2xf32>238}239