255 lines · plain
1// RUN: mlir-opt -transform-interpreter --split-input-file -canonicalize %s | FileCheck %s2 3// CHECK-LABEL: func @tensor_from_elements_0d(4// CHECK-SAME: %[[arg0:.*]]: index5// CHECK: %[[empty:.*]] = tensor.empty() : tensor<index>6// CHECK: %[[insert:.*]] = tensor.insert %[[arg0]] into %[[empty]][]7// CHECK: return %[[insert]]8func.func @tensor_from_elements_0d(%arg0: index) -> tensor<index> {9 %0 = tensor.from_elements %arg0 : tensor<index>10 return %0 : tensor<index>11}12 13module attributes {transform.with_named_sequence} {14 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {15 %0 = transform.structured.match ops{["tensor.from_elements"]} in %arg116 : (!transform.any_op) -> !transform.any_op17 transform.structured.rewrite_in_destination_passing_style %018 : (!transform.any_op) -> !transform.any_op19 transform.yield20 }21}22 23// -----24 25// CHECK-LABEL: func @tensor_from_elements_1d(26// CHECK-SAME: %[[arg0:.*]]: index, %[[arg1:.*]]: index27// CHECK-DAG: %[[empty:.*]] = tensor.empty() : tensor<2xindex>28// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index29// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index30// CHECK: %[[insert:.*]] = tensor.insert %[[arg0]] into %[[empty]][%[[c0]]]31// CHECK: %[[insert2:.*]] = tensor.insert %[[arg1]] into %[[insert]][%[[c1]]]32// CHECK: return %[[insert2]]33func.func @tensor_from_elements_1d(%arg0: index, %arg1: index) -> tensor<2xindex> {34 %0 = tensor.from_elements %arg0, %arg1 : tensor<2xindex>35 return %0 : tensor<2xindex>36}37 38module attributes {transform.with_named_sequence} {39 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {40 %0 = transform.structured.match ops{["tensor.from_elements"]} in %arg141 : (!transform.any_op) -> !transform.any_op42 transform.structured.rewrite_in_destination_passing_style %043 : (!transform.any_op) -> !transform.any_op44 transform.yield45 }46}47 48// -----49 50// CHECK-LABEL: func @tensor_from_elements_2d(51// CHECK-SAME: %[[arg0:.*]]: index, %[[arg1:.*]]: index52// CHECK-DAG: %[[empty:.*]] = tensor.empty() : tensor<3x2xindex>53// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index54// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index55// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index56// CHECK: %[[insert0:.*]] = tensor.insert %[[arg0]] into %[[empty]][%[[c0]], %[[c0]]]57// CHECK: %[[insert1:.*]] = tensor.insert %[[arg1]] into %[[insert0]][%[[c0]], %[[c1]]]58// CHECK: %[[insert2:.*]] = tensor.insert %[[arg0]] into %[[insert1]][%[[c1]], %[[c0]]]59// CHECK: %[[insert3:.*]] = tensor.insert %[[arg1]] into %[[insert2]][%[[c1]], %[[c1]]]60// CHECK: %[[insert4:.*]] = tensor.insert %[[arg0]] into %[[insert3]][%[[c2]], %[[c0]]]61// CHECK: %[[insert5:.*]] = tensor.insert %[[arg1]] into %[[insert4]][%[[c2]], %[[c1]]]62// CHECK: return %[[insert5]]63func.func @tensor_from_elements_2d(%arg0: index, %arg1: index) -> tensor<3x2xindex> {64 %0 = tensor.from_elements %arg0, %arg1, %arg0, %arg1, %arg0, %arg165 : tensor<3x2xindex>66 return %0 : tensor<3x2xindex>67}68 69module attributes {transform.with_named_sequence} {70 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {71 %0 = transform.structured.match ops{["tensor.from_elements"]} in %arg172 : (!transform.any_op) -> !transform.any_op73 transform.structured.rewrite_in_destination_passing_style %074 : (!transform.any_op) -> !transform.any_op75 transform.yield76 }77}78 79// -----80 81// CHECK: #[[$map:.*]] = affine_map<(d0, d1) -> (d0, d1)>82// CHECK-LABEL: func @tensor_generate(83// CHECK-SAME: %[[s1:.*]]: index, %[[s2:.*]]: index84// CHECK: %[[empty:.*]] = tensor.empty(%[[s1]], %[[s2]]) : tensor<?x?xindex>85// CHECK: %[[generic:.*]] = linalg.generic86// CHECK-SAME: {indexing_maps = [#[[$map]]], iterator_types = ["parallel", "parallel"]}87// CHECK-SAME: outs(%[[empty]] : tensor<?x?xindex>) {88// CHECK: %[[i0:.*]] = linalg.index 089// CHECK: %[[i1:.*]] = linalg.index 190// CHECK: %[[added:.*]] = arith.addi %[[i0]], %[[i1]]91// CHECK: linalg.yield %[[added]]92// CHECK: }93// CHECK: return %[[generic]]94func.func @tensor_generate(%s1: index, %s2: index) -> tensor<?x?xindex> {95 %0 = tensor.generate %s1, %s2 {96 ^bb0(%arg0: index, %arg1: index):97 %1 = arith.addi %arg0, %arg1 : index98 tensor.yield %1 : index99 } : tensor<?x?xindex>100 return %0 : tensor<?x?xindex>101}102 103module attributes {transform.with_named_sequence} {104 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {105 %0 = transform.structured.match ops{["tensor.generate"]} in %arg1106 : (!transform.any_op) -> !transform.any_op107 transform.structured.rewrite_in_destination_passing_style %0108 : (!transform.any_op) -> !transform.any_op109 transform.yield110 }111}112 113// -----114 115// CHECK: #[[$map:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)>116// CHECK: #[[$map1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 10)>117// CHECK: #[[$map2:.+]] = affine_map<(d0, d1) -> (d0, d1)>118// CHECK-LABEL: func @tensor_pad(119// CHECK-SAME: %[[t1:.*]]: tensor<?x10xindex>, %[[l2:.*]]: index, %[[h1:.*]]: index, %[[h2:.*]]: index120// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index121// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]122// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[h1]]]123// CHECK-DAG: %[[size1:.*]] = affine.apply #[[$map1]]()[%[[l2]], %[[h2]]]124// CHECK: %[[empty:.*]] = tensor.empty(%[[size0]], %[[size1]]) : tensor<?x?xindex>125// CHECK: %[[generic:.*]] = linalg.generic126// CHECK-SAME: {indexing_maps = [#[[$map2]]], iterator_types = ["parallel", "parallel"]}127// CHECK-SAME: outs(%[[empty]] : tensor<?x?xindex>) {128// CHECK: %[[i0:.*]] = linalg.index 0129// CHECK: %[[i1:.*]] = linalg.index 1130// CHECK: %[[mul:.*]] = arith.muli %[[i0]], %[[i1]]131// CHECK: linalg.yield %[[mul]]132// CHECK: }133// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]134// CHECK: %[[inserted:.*]] = tensor.insert_slice %[[t1]] into %[[generic]][5, %[[l2]]] [%[[dim0]], 10] [1, 1] : tensor<?x10xindex> into tensor<?x?xindex>135// CHECK: return %[[inserted]]136func.func @tensor_pad(%t1: tensor<?x10xindex>, %l2: index, %h1: index,137 %h2: index) -> tensor<?x?xindex> {138 %0 = tensor.pad %t1 low[5, %l2] high[%h1, %h2] {139 ^bb0(%arg0: index, %arg1: index):140 %m = arith.muli %arg0, %arg1 : index141 tensor.yield %m : index142 } : tensor<?x10xindex> to tensor<?x?xindex>143 return %0 : tensor<?x?xindex>144}145 146module attributes {transform.with_named_sequence} {147 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {148 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1149 : (!transform.any_op) -> !transform.any_op150 transform.structured.rewrite_in_destination_passing_style %0151 : (!transform.any_op) -> !transform.any_op152 transform.yield153 }154}155 156// -----157 158// CHECK: #[[$map:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)>159// CHECK: #[[$map1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 10)>160// CHECK-LABEL: func @tensor_pad_constant(161// CHECK-SAME: %[[t1:.*]]: tensor<?x10xindex>, %[[l2:.*]]: index, %[[h1:.*]]: index, %[[h2:.*]]: index162// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index163// CHECK-DAG: %[[c50:.*]] = arith.constant 50 : index164// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]165// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[h1]]]166// CHECK-DAG: %[[size1:.*]] = affine.apply #[[$map1]]()[%[[l2]], %[[h2]]]167// CHECK: %[[empty:.*]] = tensor.empty(%[[size0]], %[[size1]]) : tensor<?x?xindex>168// CHECK: %[[filled:.*]] = linalg.fill ins(%[[c50]] : index) outs(%[[empty]] : tensor<?x?xindex>)169// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]170// CHECK: %[[inserted:.*]] = tensor.insert_slice %[[t1]] into %[[filled]][5, %[[l2]]] [%[[dim0]], 10] [1, 1] : tensor<?x10xindex> into tensor<?x?xindex>171// CHECK: return %[[inserted]]172func.func @tensor_pad_constant(%t1: tensor<?x10xindex>, %l2: index, %h1: index,173 %h2: index) -> tensor<?x?xindex> {174 %0 = tensor.pad %t1 low[5, %l2] high[%h1, %h2] {175 ^bb0(%arg0: index, %arg1: index):176 %c = arith.constant 50 : index177 tensor.yield %c : index178 } : tensor<?x10xindex> to tensor<?x?xindex>179 return %0 : tensor<?x?xindex>180}181 182module attributes {transform.with_named_sequence} {183 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {184 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1185 : (!transform.any_op) -> !transform.any_op186 transform.structured.rewrite_in_destination_passing_style %0187 : (!transform.any_op) -> !transform.any_op188 transform.yield189 }190}191 192// -----193 194// CHECK: #[[$map:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)>195// CHECK: #[[$map1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 10)>196// CHECK-LABEL: func @tensor_pad_invariant(197// CHECK-SAME: %[[t1:.*]]: tensor<?x10xindex>, %[[l2:.*]]: index, %[[h1:.*]]: index, %[[h2:.*]]: index, %[[padding:.*]]: index198// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index199// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]200// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[h1]]]201// CHECK-DAG: %[[size1:.*]] = affine.apply #[[$map1]]()[%[[l2]], %[[h2]]]202// CHECK: %[[empty:.*]] = tensor.empty(%[[size0]], %[[size1]]) : tensor<?x?xindex>203// CHECK: %[[filled:.*]] = linalg.fill ins(%[[padding]] : index) outs(%[[empty]] : tensor<?x?xindex>)204// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]205// CHECK: %[[inserted:.*]] = tensor.insert_slice %[[t1]] into %[[filled]][5, %[[l2]]] [%[[dim0]], 10] [1, 1] : tensor<?x10xindex> into tensor<?x?xindex>206// CHECK: return %[[inserted]]207func.func @tensor_pad_invariant(%t1: tensor<?x10xindex>, %l2: index, %h1: index,208 %h2: index, %padding: index) -> tensor<?x?xindex> {209 %0 = tensor.pad %t1 low[5, %l2] high[%h1, %h2] {210 ^bb0(%arg0: index, %arg1: index):211 tensor.yield %padding : index212 } : tensor<?x10xindex> to tensor<?x?xindex>213 return %0 : tensor<?x?xindex>214}215 216module attributes {transform.with_named_sequence} {217 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {218 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1219 : (!transform.any_op) -> !transform.any_op220 transform.structured.rewrite_in_destination_passing_style %0221 : (!transform.any_op) -> !transform.any_op222 transform.yield223 }224}225 226// -----227 228// CHECK-LABEL: func @tensor_pad_nofold(229// CHECK-SAME: %[[t1:.*]]: tensor<?x?xindex>, %[[padding:.*]]: index230// CHECK-NOT: linalg.fill231// CHECK-NOT: generic232// CHECK-NOT: insert_slice233// CHECK: %[[alloc_tensor:.*]] = bufferization.alloc_tensor(%{{.*}}) : tensor<?x?xindex>234// CHECK: %[[copied:.*]] = linalg.copy ins(%[[t1]] : tensor<?x?xindex>) outs(%[[alloc_tensor]] : tensor<?x?xindex>) -> tensor<?x?xindex>235// CHECK: return %[[copied]]236func.func @tensor_pad_nofold(%t1: tensor<?x?xindex>, %padding: index)237 -> tensor<?x?xindex> {238 %c0 = arith.constant 0 : index239 %0 = tensor.pad %t1 nofold low[0, %c0] high[%c0, 0] {240 ^bb0(%arg0: index, %arg1: index):241 tensor.yield %padding : index242 } : tensor<?x?xindex> to tensor<?x?xindex>243 return %0: tensor<?x?xindex>244}245 246module attributes {transform.with_named_sequence} {247 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {248 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1249 : (!transform.any_op) -> !transform.any_op250 transform.structured.rewrite_in_destination_passing_style %0251 : (!transform.any_op) -> !transform.any_op252 transform.yield253 }254}255