1// RUN: mlir-opt -split-input-file -verify-diagnostics \ 2// RUN: -transform-interpreter -canonicalize \ 3// RUN: -allow-unregistered-dialect -split-input-file %s | FileCheck %s 4 5// CHECK: #[[$map:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)> 6// CHECK: #[[$map1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 10)> 7// CHECK-LABEL: func @tensor_pad_constant( 8// CHECK-SAME: %[[t:.*]]: tensor<?x10xindex>, %[[l2:.*]]: index, %[[h1:.*]]: index, %[[h2:.*]]: index 9// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index 10// CHECK-DAG: %[[c50:.*]] = arith.constant 50 : index 11// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t]], %[[c0]] 12// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[h1]], %[[dim0]]] 13// CHECK-DAG: %[[size1:.*]] = affine.apply #[[$map1]]()[%[[l2]], %[[h2]]] 14// CHECK: %[[alloc:.*]] = memref.alloc(%[[size0]], %[[size1]]) : memref<?x?xindex> 15// CHECK: linalg.fill ins(%[[c50]] : index) outs(%[[alloc]] : memref<?x?xindex>) 16// CHECK: %[[dim0:.*]] = tensor.dim %[[t]], %[[c0]] 17// CHECK: %[[subview:.*]] = memref.subview %[[alloc]][5, %[[l2]]] [%[[dim0]], 10] [1, 1] 18// CHECK: bufferization.materialize_in_destination %[[t]] in writable %[[subview]] 19// CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable : memref<?x?xindex> 20// CHECK: memref.dealloc %[[alloc]] 21// CHECK: return %[[r]] 22func.func @tensor_pad_constant(%t: tensor<?x10xindex>, %l2: index, %h1: index, 23 %h2: index) -> tensor<?x?xindex> { 24 %0 = tensor.pad %t low[5, %l2] high[%h1, %h2] { 25 ^bb0(%arg0: index, %arg1: index): 26 %c = arith.constant 50 : index 27 tensor.yield %c : index 28 } : tensor<?x10xindex> to tensor<?x?xindex> 29 return %0 : tensor<?x?xindex> 30} 31 32module attributes {transform.with_named_sequence} { 33 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 34 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op 35 %2, %new = transform.structured.bufferize_to_allocation %0 {emit_dealloc} : !transform.any_op 36 37 // Ensure that one linalg.fill was generated. 38 %fill_op = transform.select "linalg.fill" in %new : (!transform.any_op) -> !transform.any_op 39 %p = transform.num_associations %fill_op : (!transform.any_op) -> !transform.param<i64> 40 // expected-remark @below{{1}} 41 transform.debug.emit_param_as_remark %p : !transform.param<i64> 42 43 // Ensure that one linalg.copy was generated. 44 %mat = transform.select "bufferization.materialize_in_destination" in %new : (!transform.any_op) -> !transform.any_op 45 %p2 = transform.num_associations %mat : (!transform.any_op) -> !transform.param<i64> 46 // expected-remark @below{{1}} 47 transform.debug.emit_param_as_remark %p2 : !transform.param<i64> 48 transform.yield 49 } 50} 51 52// ----- 53 54// CHECK-LABEL: func @tensor_pad_constant_with_custom_copy( 55// CHECK-NOT: bufferization.materialize_in_destination 56// CHECK-NOT: memref.copy 57// CHECK: memref.alloca 58// CHECK: linalg.copy 59func.func @tensor_pad_constant_with_custom_copy( 60 %t: tensor<?x10xindex>, %l2: index, %h1: index, %h2: index) 61 -> tensor<?x?xindex> 62{ 63 %0 = tensor.pad %t low[5, %l2] high[%h1, %h2] { 64 ^bb0(%arg0: index, %arg1: index): 65 %c = arith.constant 50 : index 66 tensor.yield %c : index 67 } : tensor<?x10xindex> to tensor<?x?xindex> 68 return %0 : tensor<?x?xindex> 69} 70 71module attributes {transform.with_named_sequence} { 72 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.consumed}) { 73 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op 74 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 3, alloc_op = "memref.alloca", memcpy_op = "linalg.copy", emit_dealloc}: !transform.any_op 75 76 // Ensure that one linalg.fill was generated. 77 %fill_op = transform.select "linalg.fill" in %new : (!transform.any_op) -> !transform.any_op 78 %p = transform.num_associations %fill_op : (!transform.any_op) -> !transform.param<i64> 79 // expected-remark @below{{1}} 80 transform.debug.emit_param_as_remark %p : !transform.param<i64> 81 82 // Ensure that one linalg.copy was generated. 83 %linalg_copy = transform.select "linalg.copy" in %new : (!transform.any_op) -> !transform.any_op 84 %p2 = transform.num_associations %linalg_copy : (!transform.any_op) -> !transform.param<i64> 85 // expected-remark @below{{1}} 86 transform.debug.emit_param_as_remark %p2 : !transform.param<i64> 87 88 // Ensure that one memref.alloca was generated. 89 %alloca = transform.select "memref.alloca" in %new : (!transform.any_op) -> !transform.any_op 90 %p3 = transform.num_associations %alloca : (!transform.any_op) -> !transform.param<i64> 91 // expected-remark @below{{1}} 92 transform.debug.emit_param_as_remark %p3 : !transform.param<i64> 93 94 // Make sure that One-Shot Bufferize can bufferize the rest. 95 %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op 96 transform.yield 97 } 98} 99 100// ----- 101 102// CHECK-LABEL: func @tensor_pad_constant( 103// CHECK-SAME: %[[t:.*]]: tensor<?x10xindex> 104// CHECK: %[[src:.*]] = bufferization.to_memref %[[t]] 105// CHECK: %[[alloc:.*]] = memref.alloc 106// CHECK: %[[subview:.*]] = memref.subview %[[alloc]] 107// CHECK: memref.copy %[[src]], %[[subview]] 108// CHECK: bufferization.to_tensor %[[alloc]] restrict writable 109func.func @tensor_pad_constant(%t: tensor<?x10xindex>, %l2: index, %h1: index, 110 %h2: index) -> tensor<?x?xindex> { 111 %0 = tensor.pad %t low[5, %l2] high[%h1, %h2] { 112 ^bb0(%arg0: index, %arg1: index): 113 %c = arith.constant 50 : index 114 tensor.yield %c : index 115 } : tensor<?x10xindex> to tensor<?x?xindex> 116 return %0 : tensor<?x?xindex> 117} 118 119module attributes {transform.with_named_sequence} { 120 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.consumed}) { 121 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op 122 %2, %new = transform.structured.bufferize_to_allocation %0 {emit_dealloc} : !transform.any_op 123 // Make sure that One-Shot Bufferize can bufferize the rest. 124 %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op 125 transform.yield 126 } 127} 128 129// ----- 130 131// CHECK-LABEL: func @tensor_insert( 132// CHECK-SAME: %[[t:.*]]: tensor<?x10xindex> 133// CHECK: %[[m:.*]] = bufferization.to_memref %[[t]] 134// CHECK: %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4> 135// CHECK: memref.copy %[[m]], %[[alloc]] 136// CHECK: memref.store %{{.*}}, %[[alloc]] 137// CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable 138// CHECK: memref.dealloc %[[alloc]] 139// CHECK: return %[[r]] 140func.func @tensor_insert(%t: tensor<?x10xindex>, %idx: index, %v: index) -> tensor<?x10xindex> { 141 %r = tensor.insert %v into %t[%idx, %idx] : tensor<?x10xindex> 142 return %r : tensor<?x10xindex> 143} 144 145module attributes {transform.with_named_sequence} { 146 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.consumed}) { 147 %0 = transform.structured.match ops{["tensor.insert"]} in %arg1 : (!transform.any_op) -> !transform.any_op 148 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, emit_dealloc} : !transform.any_op 149 // Make sure that One-Shot Bufferize can bufferize the rest. 150 %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op 151 transform.yield 152 } 153} 154 155// ----- 156 157// CHECK-LABEL: func @tensor_insert_into_empty( 158// CHECK: %[[alloc:.*]] = memref.alloc() : memref<10xindex, 4> 159// CHECK-NOT: memref.copy 160// CHECK: memref.store %{{.*}}, %[[alloc]] 161// CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable 162// CHECK: memref.dealloc %[[alloc]] 163// CHECK: return %[[r]] 164func.func @tensor_insert_into_empty(%idx: index, %v: index) -> tensor<10xindex> { 165 %e = tensor.empty() : tensor<10xindex> 166 %r = tensor.insert %v into %e[%idx] : tensor<10xindex> 167 return %r : tensor<10xindex> 168} 169 170module attributes {transform.with_named_sequence} { 171 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.consumed}) { 172 %0 = transform.structured.match ops{["tensor.insert"]} in %arg1 : (!transform.any_op) -> !transform.any_op 173 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, emit_dealloc} : !transform.any_op 174 // Make sure that One-Shot Bufferize can bufferize the rest. 175 %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op 176 transform.yield 177 } 178} 179 180// ----- 181 182func.func @tensor_extract(%t: tensor<?x10xindex>, %idx: index) -> index { 183 // expected-note @below{{target payload op}} 184 %r = tensor.extract %t[%idx, %idx] : tensor<?x10xindex> 185 return %r : index 186} 187 188module attributes {transform.with_named_sequence} { 189 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 190 %0 = transform.structured.match ops{["tensor.extract"]} in %arg1 : (!transform.any_op) -> !transform.any_op 191 // expected-error @below{{failed to bufferize operation}} 192 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, emit_dealloc} : !transform.any_op 193 transform.yield 194 } 195} 196 197// ----- 198 199// CHECK-LABEL: func @vector_mask( 200// CHECK-SAME: %[[t:.*]]: tensor<?xf32>, 201// CHECK: %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?xf32, 4> 202// CHECK: bufferization.materialize_in_destination %[[t]] in writable %[[alloc]] 203// CHECK: vector.mask %{{.*}} { vector.transfer_write %{{.*}}, %[[alloc]] 204// CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable 205// CHECK: memref.dealloc %[[alloc]] 206// CHECK: return %[[r]] 207func.func @vector_mask(%t: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vector<16xi1>) -> tensor<?xf32> { 208 %r = vector.mask %m0 { vector.transfer_write %val, %t[%idx] : vector<16xf32>, tensor<?xf32> } : vector<16xi1> -> tensor<?xf32> 209 return %r : tensor<?xf32> 210} 211 212module attributes {transform.with_named_sequence} { 213 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 214 %0 = transform.structured.match ops{["vector.mask"]} in %arg1 : (!transform.any_op) -> !transform.any_op 215 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, emit_dealloc} : !transform.any_op 216 transform.yield 217 } 218} 219 220// ----- 221 222// CHECK-LABEL: func @tensor_insert_destination( 223// CHECK-SAME: %[[t:.*]]: tensor<?x10xindex> 224// CHECK: %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4> 225// CHECK: bufferization.materialize_in_destination %[[t]] in writable %[[alloc]] 226// CHECK: %[[t2:.*]] = bufferization.to_tensor %[[alloc]] restrict writable 227// CHECK: %[[inserted:.*]] = tensor.insert %{{.*}} into %[[t2]] 228// CHECK: memref.dealloc %[[alloc]] 229// CHECK: return %[[inserted]] 230func.func @tensor_insert_destination(%t: tensor<?x10xindex>, %idx: index, %v: index) -> tensor<?x10xindex> { 231 %r = tensor.insert %v into %t[%idx, %idx] : tensor<?x10xindex> 232 return %r : tensor<?x10xindex> 233} 234 235module attributes {transform.with_named_sequence} { 236 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 237 %0 = transform.structured.match ops{["tensor.insert"]} in %arg1 : (!transform.any_op) -> !transform.any_op 238 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, bufferize_destination_only, emit_dealloc} : !transform.any_op 239 transform.yield 240 } 241} 242 243// ----- 244 245// CHECK-LABEL: func @scf_for_destination( 246// CHECK-SAME: %[[t:.*]]: tensor<?x10xindex> 247// CHECK: %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4> 248// CHECK: bufferization.materialize_in_destination %[[t]] in writable %[[alloc]] 249// CHECK: %[[t2:.*]] = bufferization.to_tensor %[[alloc]] restrict writable 250// CHECK: %[[for:.*]] = scf.for {{.*}} iter_args(%{{.*}} = %[[t2]]) 251// CHECK: memref.dealloc %[[alloc]] 252// CHECK: return %[[for]] 253func.func @scf_for_destination(%t: tensor<?x10xindex>, %lb: index, %ub: index, %step: index) -> tensor<?x10xindex> { 254 %r = scf.for %iv = %lb to %ub step %step iter_args(%a = %t) -> tensor<?x10xindex> { 255 %b = "test.foo"(%a) : (tensor<?x10xindex>) -> (tensor<?x10xindex>) 256 scf.yield %b : tensor<?x10xindex> 257 } 258 return %r : tensor<?x10xindex> 259} 260 261module attributes {transform.with_named_sequence} { 262 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 263 %0 = transform.structured.match ops{["scf.for"]} in %arg1 : (!transform.any_op) -> !transform.any_op 264 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, bufferize_destination_only, emit_dealloc} : !transform.any_op 265 transform.yield 266 } 267} 268 269// ----- 270 271// CHECK-LABEL: func @tensor_insert_destination_no_dealloc 272// CHECK-NOT: dealloc 273func.func @tensor_insert_destination_no_dealloc(%t: tensor<?x10xindex>, %idx: index, %v: index) -> tensor<?x10xindex> { 274 %r = tensor.insert %v into %t[%idx, %idx] : tensor<?x10xindex> 275 return %r : tensor<?x10xindex> 276} 277 278module attributes {transform.with_named_sequence} { 279 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { 280 %0 = transform.structured.match ops{["tensor.insert"]} in %arg1 : (!transform.any_op) -> !transform.any_op 281 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, bufferize_destination_only} : !transform.any_op 282 transform.yield 283 } 284} 285