1// RUN: mlir-opt -transform-interpreter --split-input-file -canonicalize %s | FileCheck %s 2 3// CHECK-LABEL: func @tensor_from_elements_0d( 4// CHECK-SAME: %[[arg0:.*]]: index 5// 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 %arg1 16 : (!transform.any_op) -> !transform.any_op 17 transform.structured.rewrite_in_destination_passing_style %0 18 : (!transform.any_op) -> !transform.any_op 19 transform.yield 20 } 21} 22 23// ----- 24 25// CHECK-LABEL: func @tensor_from_elements_1d( 26// CHECK-SAME: %[[arg0:.*]]: index, %[[arg1:.*]]: index 27// CHECK-DAG: %[[empty:.*]] = tensor.empty() : tensor<2xindex> 28// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index 29// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index 30// 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 %arg1 41 : (!transform.any_op) -> !transform.any_op 42 transform.structured.rewrite_in_destination_passing_style %0 43 : (!transform.any_op) -> !transform.any_op 44 transform.yield 45 } 46} 47 48// ----- 49 50// CHECK-LABEL: func @tensor_from_elements_2d( 51// CHECK-SAME: %[[arg0:.*]]: index, %[[arg1:.*]]: index 52// CHECK-DAG: %[[empty:.*]] = tensor.empty() : tensor<3x2xindex> 53// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index 54// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index 55// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index 56// 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, %arg1 65 : 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 %arg1 72 : (!transform.any_op) -> !transform.any_op 73 transform.structured.rewrite_in_destination_passing_style %0 74 : (!transform.any_op) -> !transform.any_op 75 transform.yield 76 } 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:.*]]: index 84// CHECK: %[[empty:.*]] = tensor.empty(%[[s1]], %[[s2]]) : tensor<?x?xindex> 85// CHECK: %[[generic:.*]] = linalg.generic 86// CHECK-SAME: {indexing_maps = [#[[$map]]], iterator_types = ["parallel", "parallel"]} 87// CHECK-SAME: outs(%[[empty]] : tensor<?x?xindex>) { 88// CHECK: %[[i0:.*]] = linalg.index 0 89// CHECK: %[[i1:.*]] = linalg.index 1 90// 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 : index 98 tensor.yield %1 : index 99 } : 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 %arg1 106 : (!transform.any_op) -> !transform.any_op 107 transform.structured.rewrite_in_destination_passing_style %0 108 : (!transform.any_op) -> !transform.any_op 109 transform.yield 110 } 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:.*]]: index 120// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index 121// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]] 122// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[h1]], %[[dim0]]] 123// CHECK-DAG: %[[size1:.*]] = affine.apply #[[$map1]]()[%[[l2]], %[[h2]]] 124// CHECK: %[[empty:.*]] = tensor.empty(%[[size0]], %[[size1]]) : tensor<?x?xindex> 125// CHECK: %[[generic:.*]] = linalg.generic 126// CHECK-SAME: {indexing_maps = [#[[$map2]]], iterator_types = ["parallel", "parallel"]} 127// CHECK-SAME: outs(%[[empty]] : tensor<?x?xindex>) { 128// CHECK: %[[i0:.*]] = linalg.index 0 129// CHECK: %[[i1:.*]] = linalg.index 1 130// 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 : index 141 tensor.yield %m : index 142 } : 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 %arg1 149 : (!transform.any_op) -> !transform.any_op 150 transform.structured.rewrite_in_destination_passing_style %0 151 : (!transform.any_op) -> !transform.any_op 152 transform.yield 153 } 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:.*]]: index 162// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index 163// CHECK-DAG: %[[c50:.*]] = arith.constant 50 : index 164// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]] 165// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[h1]], %[[dim0]]] 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 : index 177 tensor.yield %c : index 178 } : 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 %arg1 185 : (!transform.any_op) -> !transform.any_op 186 transform.structured.rewrite_in_destination_passing_style %0 187 : (!transform.any_op) -> !transform.any_op 188 transform.yield 189 } 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:.*]]: index 198// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index 199// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]] 200// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[h1]], %[[dim0]]] 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 : index 212 } : 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 %arg1 219 : (!transform.any_op) -> !transform.any_op 220 transform.structured.rewrite_in_destination_passing_style %0 221 : (!transform.any_op) -> !transform.any_op 222 transform.yield 223 } 224} 225 226// ----- 227 228// CHECK-LABEL: func @tensor_pad_nofold( 229// CHECK-SAME: %[[t1:.*]]: tensor<?x?xindex>, %[[padding:.*]]: index 230// CHECK-NOT: linalg.fill 231// CHECK-NOT: generic 232// CHECK-NOT: insert_slice 233// 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 : index 239 %0 = tensor.pad %t1 nofold low[0, %c0] high[%c0, 0] { 240 ^bb0(%arg0: index, %arg1: index): 241 tensor.yield %padding : index 242 } : 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 %arg1 249 : (!transform.any_op) -> !transform.any_op 250 transform.structured.rewrite_in_destination_passing_style %0 251 : (!transform.any_op) -> !transform.any_op 252 transform.yield 253 } 254} 255