285 lines · plain
1// RUN: mlir-opt -split-input-file -verify-diagnostics \2// RUN: -transform-interpreter -canonicalize \3// RUN: -allow-unregistered-dialect -split-input-file %s | FileCheck %s4 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:.*]]: index9// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index10// CHECK-DAG: %[[c50:.*]] = arith.constant 50 : index11// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t]], %[[c0]]12// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[h1]]]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 : index27 tensor.yield %c : index28 } : 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_op35 %2, %new = transform.structured.bufferize_to_allocation %0 {emit_dealloc} : !transform.any_op36 37 // Ensure that one linalg.fill was generated.38 %fill_op = transform.select "linalg.fill" in %new : (!transform.any_op) -> !transform.any_op39 %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_op45 %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.yield49 }50}51 52// -----53 54// CHECK-LABEL: func @tensor_pad_constant_with_custom_copy(55// CHECK-NOT: bufferization.materialize_in_destination56// CHECK-NOT: memref.copy57// CHECK: memref.alloca58// CHECK: linalg.copy59func.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 : index66 tensor.yield %c : index67 } : 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_op74 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 3, alloc_op = "memref.alloca", memcpy_op = "linalg.copy", emit_dealloc}: !transform.any_op75 76 // Ensure that one linalg.fill was generated.77 %fill_op = transform.select "linalg.fill" in %new : (!transform.any_op) -> !transform.any_op78 %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_op84 %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_op90 %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_op96 transform.yield97 }98}99 100// -----101 102// CHECK-LABEL: func @tensor_pad_constant(103// CHECK-SAME: %[[t:.*]]: tensor<?x10xindex>104// CHECK: %[[src:.*]] = bufferization.to_buffer %[[t]]105// CHECK: %[[alloc:.*]] = memref.alloc106// CHECK: %[[subview:.*]] = memref.subview %[[alloc]]107// CHECK: memref.copy %[[src]], %[[subview]]108// CHECK: bufferization.to_tensor %[[alloc]] restrict writable109func.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 : index114 tensor.yield %c : index115 } : 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_op122 %2, %new = transform.structured.bufferize_to_allocation %0 {emit_dealloc} : !transform.any_op123 // Make sure that One-Shot Bufferize can bufferize the rest.124 %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op125 transform.yield126 }127}128 129// -----130 131// CHECK-LABEL: func @tensor_insert(132// CHECK-SAME: %[[t:.*]]: tensor<?x10xindex>133// CHECK: %[[m:.*]] = bufferization.to_buffer %[[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 writable138// 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_op148 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, emit_dealloc} : !transform.any_op149 // Make sure that One-Shot Bufferize can bufferize the rest.150 %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op151 transform.yield152 }153}154 155// -----156 157// CHECK-LABEL: func @tensor_insert_into_empty(158// CHECK: %[[alloc:.*]] = memref.alloc() : memref<10xindex, 4>159// CHECK-NOT: memref.copy160// CHECK: memref.store %{{.*}}, %[[alloc]]161// CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable162// 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_op173 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, emit_dealloc} : !transform.any_op174 // Make sure that One-Shot Bufferize can bufferize the rest.175 %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op176 transform.yield177 }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 : index186}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_op191 // expected-error @below{{failed to bufferize operation}}192 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, emit_dealloc} : !transform.any_op193 transform.yield194 }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 writable205// 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_op215 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, emit_dealloc} : !transform.any_op216 transform.yield217 }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 writable227// 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_op238 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, bufferize_destination_only, emit_dealloc} : !transform.any_op239 transform.yield240 }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 writable250// 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_op264 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, bufferize_destination_only, emit_dealloc} : !transform.any_op265 transform.yield266 }267}268 269// -----270 271// CHECK-LABEL: func @tensor_insert_destination_no_dealloc272// CHECK-NOT: dealloc273func.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_op281 %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, bufferize_destination_only} : !transform.any_op282 transform.yield283 }284}285