586 lines · plain
1// RUN: mlir-opt %s | mlir-opt | FileCheck %s2 3// CHECK: shard.grid @grid04shard.grid @grid0(shape = 2x2x4)5 6// CHECK: shard.grid @grid1(shape = 4x?)7shard.grid @grid1(shape = 4x?)8 9// CHECK: shard.grid @grid2(shape = ?x4)10shard.grid @grid2(shape = ?x4)11 12// CHECK: shard.grid @grid3(shape = ?x?)13shard.grid @grid3(shape = ?x?)14 15shard.grid @grid4(shape = 3)16 17// CHECK: shard.grid @grid5(shape = ?)18shard.grid @grid5(shape = ?)19 20// CHECK-LABEL: func @grid_shard_op_fully_replicated21// CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>22func.func @grid_shard_op_fully_replicated(%arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {23 // CHECK-NEXT: %[[S:.*]] = shard.sharding @grid0 split_axes = {{\[\[}}]] : !shard.sharding24 %s = shard.sharding @grid0 split_axes = [[]] : !shard.sharding25 // CHECK-NEXT: shard.shard %[[ARG]] to %[[S]] : tensor<4x8xf32>26 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>27 return %0 : tensor<4x8xf32>28}29 30// CHECK-LABEL: func @grid_shard_op_1st_dim31// CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>32func.func @grid_shard_op_1st_dim(%arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {33 // CHECK-NEXT: %[[S:.*]] = shard.sharding @grid0 split_axes = {{\[\[}}0]] : !shard.sharding34 %s = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding35 36 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>37 return %0 : tensor<4x8xf32>38}39 40// CHECK-LABEL: func @grid_shard_op_2nd_dim41// CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>42func.func @grid_shard_op_2nd_dim(%arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {43 // CHECK-NEXT: %[[S:.*]] = shard.sharding @grid1 split_axes = {{\[\[}}], [0]] : !shard.sharding44 %s = shard.sharding @grid1 split_axes = [[], [0]] : !shard.sharding45 // CHECK-NEXT: shard.shard %[[ARG]] to %[[S]] : tensor<4x8xf32>46 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>47 return %0 : tensor<4x8xf32>48}49 50// CHECK-LABEL: func @grid_shard_op_1st_and_3rd_dim51func.func @grid_shard_op_1st_and_3rd_dim(52 // CHECK-SAME: %[[ARG:.*]]: tensor<4x8x16xf32>53 %arg0 : tensor<4x8x16xf32>) -> tensor<4x8x16xf32> {54 // CHECK-NEXT: %[[S:.*]] = shard.sharding @grid3 split_axes = {{\[\[}}0], [], [1]] : !shard.sharding55 %s = shard.sharding @grid3 split_axes = [[0], [], [1]] : !shard.sharding56 // CHECK-NEXT: shard.shard %[[ARG]] to %[[S]] : tensor<4x8x16xf32>57 %0 = shard.shard %arg0 to %s : tensor<4x8x16xf32>58 return %0 : tensor<4x8x16xf32>59}60 61// CHECK-LABEL: func @grid_shard_op_two_users62// CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>63func.func @grid_shard_op_two_users(%arg0 : tensor<4x8xf32>) ->64 (tensor<4x8xf32>, tensor<4x8xf32>) {65 // CHECK-NEXT: %[[V0:.*]] = shard.sharding @grid0 split_axes = {{\[\[}}0]] : !shard.sharding66 %s0 = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding67 %0 = shard.shard %arg0 to %s0 : tensor<4x8xf32>68 // CHECK-DAG: shard.sharding @grid0 split_axes = {{\[\[}}1]] : !shard.sharding69 %s1 = shard.sharding @grid0 split_axes = [[1]] : !shard.sharding70 %1 = shard.shard %0 to %s1 annotate_for_users : tensor<4x8xf32>71 // CHECK-DAG: shard.sharding @grid0 split_axes = {{\[\[}}2]] : !shard.sharding72 %s2 = shard.sharding @grid0 split_axes = [[2]] : !shard.sharding73 %2 = shard.shard %0 to %s2 annotate_for_users : tensor<4x8xf32>74 return %1, %2 : tensor<4x8xf32>, tensor<4x8xf32>75}76 77// CHECK-LABEL: func @grid_shard_halo_sizes78func.func @grid_shard_halo_sizes() -> () {79 // CHECK: %[[C3:.*]] = arith.constant 3 : i6480 %c3 = arith.constant 3 : i6481 // CHECK: shard.sharding @grid4 split_axes = {{\[\[}}0]] halo_sizes = [1, 4] : !shard.sharding82 %sharding1 = shard.sharding @grid4 split_axes = [[0]] halo_sizes = [1, 4] : !shard.sharding83 // CHECK: shard.sharding @grid4 split_axes = {{\[\[}}0]] halo_sizes = [4, %[[C3]]] : !shard.sharding84 %sharding2 = shard.sharding @grid4 split_axes = [[0]] halo_sizes = [4, %c3] : !shard.sharding85 return86}87 88// CHECK-LABEL: func @grid_shard_dims_sizes89func.func @grid_shard_dims_sizes() -> () {90 // CHECK: %[[C3:.*]] = arith.constant 3 : i6491 %c3 = arith.constant 3 : i6492 // CHECK: shard.sharding @grid4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 1, 4, 6] : !shard.sharding93 %sharding1 = shard.sharding @grid4 split_axes = [[0]] sharded_dims_offsets = [0, 1, 4, 6] : !shard.sharding94 // CHECK: shard.sharding @grid4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 2, %[[C3]], 5] : !shard.sharding95 %sharding2 = shard.sharding @grid4 split_axes = [[0]] sharded_dims_offsets = [0, 2, %c3, 5] : !shard.sharding96 return97}98 99// CHECK-LABEL: func @grid_shard_shape100func.func @grid_shard_shape() {101 // CHECK: %[[C3:.*]] = arith.constant 3 : index102 %c3 = arith.constant 3 : index103 // CHECK-NEXT: %[[S:.*]] = shard.sharding @grid0 split_axes = {{\[\[}}]] : !shard.sharding104 %s = shard.sharding @grid0 split_axes = [[]] : !shard.sharding105 // CHECK-NEXT: shard.shard_shape dims = [8, %[[C3]]106 // CHECK-SAME: ] sharding = %[[S]] device = [%[[C3]]107 // CHECK-SAME: ] : index, index108 %shp:2 = shard.shard_shape dims = [8, %c3] sharding = %s device = [%c3] : index, index109 // CHECK-NEXT: shard.shard_shape dims = [8, 4] sharding = %[[S]] device = [3] : index, index110 %shp1:2 = shard.shard_shape dims = [8, 4] sharding = %s device = [3] : index, index111 return112}113 114// CHECK-LABEL: func @grid_get_sharding115// CHECK-SAME: %[[ARG:.*]]: tensor<4x8xf32>116func.func @grid_get_sharding(%arg0 : tensor<4x8xf32>) -> !shard.sharding {117 // CHECK-NEXT: shard.get_sharding %[[ARG]] : tensor<4x8xf32> -> !shard.sharding118 %0 = shard.get_sharding %arg0 : tensor<4x8xf32> -> !shard.sharding119 return %0 : !shard.sharding120}121 122// CHECK-LABEL: func @grid_shape123func.func @grid_shape() -> (index, index) {124 // CHECK: %[[RES:.*]]:2 = shard.grid_shape @grid0 axes = [0, 1] : index, index125 %0:2 = shard.grid_shape @grid0 axes = [0, 1] : index, index126 // CHECK: return %[[RES]]#0, %[[RES]]#1 : index, index127 return %0#0, %0#1 : index, index128}129 130// CHECK-LABEL: func @grid_shape_default_axes131func.func @grid_shape_default_axes() -> (index, index, index) {132 // CHECK: %[[RES:.*]]:3 = shard.grid_shape @grid0 : index, index, index133 %0:3 = shard.grid_shape @grid0 : index, index, index134 // CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2 : index, index, index135 return %0#0, %0#1, %0#2 : index, index, index136}137 138// CHECK-LABEL: func @grid_shape_empty_axes139func.func @grid_shape_empty_axes() -> (index, index, index) {140 // CHECK: %[[RES:.*]]:3 = shard.grid_shape @grid0 : index, index, index141 %0:3 = shard.grid_shape @grid0 axes = [] : index, index, index142 // CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2 : index, index, index143 return %0#0, %0#1, %0#2 : index, index, index144}145 146// CHECK-LABEL: func @process_multi_index147func.func @process_multi_index() -> (index, index) {148 // CHECK: %[[RES:.*]]:2 = shard.process_multi_index on @grid0 axes = [0, 1] : index, index149 %0:2 = shard.process_multi_index on @grid0 axes = [0, 1] : index, index150 // CHECK: return %[[RES]]#0, %[[RES]]#1 : index, index151 return %0#0, %0#1 : index, index152}153 154// CHECK-LABEL: func @process_multi_index_default_axes155func.func @process_multi_index_default_axes() -> (index, index, index) {156 // CHECK: %[[RES:.*]]:3 = shard.process_multi_index on @grid0 : index, index, index157 %0:3 = shard.process_multi_index on @grid0 : index, index, index158 // CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2 : index, index, index159 return %0#0, %0#1, %0#2 : index, index, index160}161 162// CHECK-LABEL: func @process_multi_index_empty_axes163func.func @process_multi_index_empty_axes() -> (index, index, index) {164 // CHECK: %[[RES:.*]]:3 = shard.process_multi_index on @grid0 : index, index, index165 %0:3 = shard.process_multi_index on @grid0 axes = [] : index, index, index166 // CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2 : index, index, index167 return %0#0, %0#1, %0#2 : index, index, index168}169 170// CHECK-LABEL: func @process_linear_index171func.func @process_linear_index() -> index {172 // CHECK: %[[RES:.*]] = shard.process_linear_index on @grid0 : index173 %0 = shard.process_linear_index on @grid0 : index174 // CHECK: return %[[RES]] : index175 return %0 : index176}177 178// CHECK-LABEL: func @all_reduce179func.func @all_reduce(180 // CHECK-SAME: %[[ARG:.*]]: tensor<3x4xf32>181 %arg0 : tensor<3x4xf32>) -> tensor<3x4xf64> {182 // CHECK-NEXT: shard.all_reduce %[[ARG]] on @grid0 grid_axes = [1, 0] reduction = max183 // CHECK-SAME: : tensor<3x4xf32> -> tensor<3x4xf64>184 %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [1, 0] reduction = max185 : tensor<3x4xf32> -> tensor<3x4xf64>186 return %0 : tensor<3x4xf64>187}188 189// CHECK-LABEL: func @all_gather190func.func @all_gather(191 // CHECK-SAME: %[[ARG:.*]]: tensor<3x4xf32>192 %arg0 : tensor<3x4xf32>) -> tensor<3x16xf32> {193 // CHECK-NEXT: shard.all_gather %[[ARG]] on @grid0 grid_axes = [2] gather_axis = 1194 // CHECK-SAME: : tensor<3x4xf32> -> tensor<3x16xf32>195 %0 = shard.all_gather %arg0 on @grid0 grid_axes = [2] gather_axis = 1196 : tensor<3x4xf32> -> tensor<3x16xf32>197 return %0 : tensor<3x16xf32>198}199 200// CHECK-LABEL: func @all_gather_dynamic_dims_in_tensor201func.func @all_gather_dynamic_dims_in_tensor(202 // CHECK-SAME: %[[ARG:.*]]: tensor<?x?xf32>203 %arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {204 // CHECK-NEXT: shard.all_gather %[[ARG]] on @grid0 grid_axes = [2] gather_axis = 1205 // CHECK-SAME: : tensor<?x?xf32> -> tensor<?x?xf32>206 %0 = shard.all_gather %arg0 on @grid0 grid_axes = [2] gather_axis = 1207 : tensor<?x?xf32> -> tensor<?x?xf32>208 return %0 : tensor<?x?xf32>209}210 211// CHECK-LABEL: func @all_gather_dynamic_dims_in_grid212func.func @all_gather_dynamic_dims_in_grid(213 // CHECK-SAME: %[[ARG:.*]]: tensor<5x6xf32>214 %arg0 : tensor<5x6xf32>) -> tensor<5x?xf32> {215 // CHECK-NEXT: shard.all_gather %[[ARG]] on @grid3 grid_axes = [1] gather_axis = 1216 // CHECK-SAME: : tensor<5x6xf32> -> tensor<5x?xf32>217 %0 = shard.all_gather %arg0 on @grid3 grid_axes = [1] gather_axis = 1218 : tensor<5x6xf32> -> tensor<5x?xf32>219 return %0 : tensor<5x?xf32>220}221 222// CHECK-LABEL: func @all_slice_static_dimensions223func.func @all_slice_static_dimensions(224 // CHECK-SAME: %[[ARG:.*]]: tensor<3x4xf32>225 %arg0 : tensor<3x4xf32>) -> tensor<3x1xf32> {226 // CHECK-NEXT: shard.all_slice %[[ARG]]227 // CHECK-SAME: on @grid0 grid_axes = [2] slice_axis = 1228 // CHECK-SAME: : tensor<3x4xf32> -> tensor<3x1xf32>229 %0 = shard.all_slice %arg0 on @grid0 grid_axes = [2] slice_axis = 1230 : tensor<3x4xf32> -> tensor<3x1xf32>231 return %0 : tensor<3x1xf32>232}233 234// CHECK-LABEL: func @all_slice_dynamic_dimensions235func.func @all_slice_dynamic_dimensions(236 // CHECK-SAME: %[[ARG:.*]]: tensor<?xf32>237 %arg0 : tensor<?xf32>) -> tensor<?xf32> {238 // CHECK-NEXT: shard.all_slice %[[ARG]]239 // CHECK-SAME: on @grid3 grid_axes = [0, 1] slice_axis = 0240 // CHECK-SAME: : tensor<?xf32> -> tensor<?xf32>241 %0 = shard.all_slice %arg0 on @grid3 grid_axes = [0, 1] slice_axis = 0242 : tensor<?xf32> -> tensor<?xf32>243 return %0 : tensor<?xf32>244}245 246// CHECK-LABEL: func @all_to_all247func.func @all_to_all(248 // CHECK-SAME: %[[ARG:.*]]: tensor<3x6xi8>249 %arg0 : tensor<3x6xi8>) -> tensor<3x6xi8> {250 // CHECK-NEXT: shard.all_to_all %[[ARG]]251 // CHECK-SAME: on @grid4 split_axis = 1 concat_axis = 0252 // CHECK-SAME: : tensor<3x6xi8> -> tensor<3x6xi8>253 %0 = shard.all_to_all %arg0 on @grid4254 split_axis = 1 concat_axis = 0255 : tensor<3x6xi8> -> tensor<3x6xi8>256 return %0 : tensor<3x6xi8>257}258 259// CHECK-LABEL: func @all_to_all_dynamic_dims_in_result260func.func @all_to_all_dynamic_dims_in_result(261 // CHECK-SAME: %[[ARG:.*]]: tensor<3x6xi8>262 %arg0 : tensor<3x6xi8>) -> tensor<3x?xi8> {263 // CHECK-NEXT: shard.all_to_all %[[ARG]]264 // CHECK-SAME: on @grid4 split_axis = 1 concat_axis = 0265 // CHECK-SAME: : tensor<3x6xi8> -> tensor<3x?xi8>266 %0 = shard.all_to_all %arg0 on @grid4267 split_axis = 1 concat_axis = 0268 : tensor<3x6xi8> -> tensor<3x?xi8>269 return %0 : tensor<3x?xi8>270}271 272// CHECK-LABEL: func @all_to_all_same_split_concat_dim_with_dynamic_device_group_size273func.func @all_to_all_same_split_concat_dim_with_dynamic_device_group_size(274 // CHECK-SAME: %[[ARG:.*]]: tensor<3xi8>275 %arg0 : tensor<3xi8>) -> tensor<3xi8> {276 // CHECK-NEXT: shard.all_to_all %[[ARG]]277 // CHECK-SAME: @grid4 split_axis = 0 concat_axis = 0278 // CHECK-SAME: : tensor<3xi8> -> tensor<3xi8>279 %0 = shard.all_to_all %arg0 on @grid4280 split_axis = 0 concat_axis = 0281 : tensor<3xi8> -> tensor<3xi8>282 return %0 : tensor<3xi8>283}284 285// CHECK-LABEL: func @all_to_all_non_divisible_split_axis_size286func.func @all_to_all_non_divisible_split_axis_size(287 // CHECK-SAME: %[[ARG:.*]]: tensor<2x3xi8>288 %arg0 : tensor<2x3xi8>) -> tensor<?x12xi8> {289 // CHECK-NEXT: shard.all_to_all %[[ARG]]290 // CHECK-SAME: @grid0 grid_axes = [0, 1] split_axis = 0 concat_axis = 1291 // CHECK-SAME: : tensor<2x3xi8> -> tensor<?x12xi8>292 %0 = shard.all_to_all %arg0 on @grid0 grid_axes = [0, 1]293 split_axis = 0 concat_axis = 1294 : tensor<2x3xi8> -> tensor<?x12xi8>295 return %0 : tensor<?x12xi8>296}297 298// CHECK-LABEL: func @broadcast_static_root299func.func @broadcast_static_root(300 // CHECK-SAME: %[[ARG:.*]]: tensor<3x6xi8>301 %arg0 : tensor<3x6xi8>) -> tensor<3x6xi8> {302 // CHECK-NEXT: shard.broadcast %[[ARG]]303 // CHECK-SAME: on @grid0 grid_axes = [0, 2]304 // CHECK-SAME: root = [0, 1]305 // CHECK-SAME: : (tensor<3x6xi8>) -> tensor<3x6xi8>306 %0 = shard.broadcast %arg0 on @grid0 grid_axes = [0, 2]307 root = [0, 1]308 : (tensor<3x6xi8>) -> tensor<3x6xi8>309 return %0 : tensor<3x6xi8>310}311 312// CHECK-LABEL: func @broadcast_dynamic_root313func.func @broadcast_dynamic_root(314 // CHECK-SAME: %[[ARG0:.*]]: tensor<3x6xi8>315 %arg0 : tensor<3x6xi8>,316 // CHECK-SAME: %[[ARG1:.*]]: index317 %arg1 : index318 ) -> tensor<3x6xi8> {319 // CHECK-NEXT: shard.broadcast %[[ARG0]]320 // CHECK-SAME: on @grid0 grid_axes = [0, 2]321 // CHECK-SAME: root = [1, %[[ARG1]]]322 // CHECK-SAME: : (tensor<3x6xi8>, index) -> tensor<3x6xi8>323 %0 = shard.broadcast %arg0 on @grid0 grid_axes = [0, 2]324 root = [1, %arg1]325 : (tensor<3x6xi8>, index) -> tensor<3x6xi8>326 return %0 : tensor<3x6xi8>327}328 329// CHECK-LABEL: func @gather_static_root330func.func @gather_static_root(331 // CHECK-SAME: %[[ARG:.*]]: tensor<3x6xi8>332 %arg0 : tensor<3x6xi8>) -> tensor<24x6xi8> {333 // CHECK-NEXT: shard.gather %[[ARG]]334 // CHECK-SAME: on @grid0 grid_axes = [0, 2]335 // CHECK-SAME: gather_axis = 0336 // CHECK-SAME: root = [0, 1]337 // CHECK-SAME: : (tensor<3x6xi8>) -> tensor<24x6xi8>338 %0 = shard.gather %arg0 on @grid0 grid_axes = [0, 2]339 gather_axis = 0340 root = [0, 1]341 : (tensor<3x6xi8>) -> tensor<24x6xi8>342 return %0 : tensor<24x6xi8>343}344 345// CHECK-LABEL: func @gather_dynamic_root346func.func @gather_dynamic_root(347 // CHECK-SAME: %[[ARG0:.*]]: tensor<3x6xi8>348 %arg0 : tensor<3x6xi8>,349 // CHECK-SAME: %[[ARG1:.*]]: index350 %arg1 : index351 ) -> tensor<24x6xi8> {352 // CHECK-NEXT: shard.gather %[[ARG0]]353 // CHECK-SAME: on @grid0 grid_axes = [0, 2]354 // CHECK-SAME: gather_axis = 0355 // CHECK-SAME: root = [1, %[[ARG1]]]356 // CHECK-SAME: : (tensor<3x6xi8>, index) -> tensor<24x6xi8>357 %0 = shard.gather %arg0 on @grid0 grid_axes = [0, 2]358 gather_axis = 0359 root = [1, %arg1]360 : (tensor<3x6xi8>, index) -> tensor<24x6xi8>361 return %0 : tensor<24x6xi8>362}363 364// CHECK-LABEL: func @receive_static_source365func.func @receive_static_source(366 // CHECK-SAME: %[[ARG:.*]]: tensor<2xi8>367 %arg0 : tensor<2xi8>) -> tensor<2xi8> {368 // CHECK-NEXT: shard.recv %[[ARG]]369 // CHECK-SAME: on @grid0 grid_axes = [0, 2]370 // CHECK-SAME: source = [0, 1]371 // CHECK-SAME: : (tensor<2xi8>) -> tensor<2xi8>372 %0 = shard.recv %arg0 on @grid0 grid_axes = [0, 2]373 source = [0, 1]374 : (tensor<2xi8>) -> tensor<2xi8>375 return %0 : tensor<2xi8>376}377 378// CHECK-LABEL: func @receive_dynamic_source379func.func @receive_dynamic_source(380 // CHECK-SAME: %[[ARG0:.*]]: tensor<2xi8>381 %arg0 : tensor<2xi8>,382 // CHECK-SAME: %[[ARG1:.*]]: index383 %arg1 : index384 ) -> tensor<2xi8> {385 // CHECK-NEXT: shard.recv %[[ARG0]]386 // CHECK-SAME: on @grid0 grid_axes = [0, 2]387 // CHECK-SAME: source = [1, %[[ARG1]]]388 // CHECK-SAME: : (tensor<2xi8>, index) -> tensor<2xi8>389 %0 = shard.recv %arg0 on @grid0 grid_axes = [0, 2]390 source = [1, %arg1]391 : (tensor<2xi8>, index) -> tensor<2xi8>392 return %0 : tensor<2xi8>393}394 395// CHECK-LABEL: func @receive_no_source396func.func @receive_no_source(397 // CHECK-SAME: %[[ARG:.*]]: tensor<2xi8>398 %arg0 : tensor<2xi8>) -> tensor<2xi8> {399 // CHECK-NEXT: shard.recv %[[ARG]]400 // CHECK-NOT: source401 %0 = shard.recv %arg0 on @grid0 grid_axes = [0, 2]402 : (tensor<2xi8>) -> tensor<2xi8>403 return %0 : tensor<2xi8>404}405 406// CHECK-LABEL: func @reduce_static_root407func.func @reduce_static_root(408 // CHECK-SAME: %[[ARG:.*]]: tensor<2xi8>409 %arg0 : tensor<2xi8>) -> tensor<2xi8> {410 // CHECK-NEXT: shard.reduce %[[ARG]]411 // CHECK-SAME: on @grid0 grid_axes = [0, 2]412 // CHECK-SAME: root = [0, 1]413 // CHECK-SAME: : (tensor<2xi8>) -> tensor<2xi8>414 %0 = shard.reduce %arg0 on @grid0 grid_axes = [0, 2]415 root = [0, 1]416 : (tensor<2xi8>) -> tensor<2xi8>417 return %0 : tensor<2xi8>418}419 420// CHECK-LABEL: func @reduce_dynamic_root421func.func @reduce_dynamic_root(422 // CHECK-SAME: %[[ARG0:.*]]: tensor<2xi8>423 %arg0 : tensor<2xi8>,424 // CHECK-SAME: %[[ARG1:.*]]: index425 %arg1 : index426 ) -> tensor<2xi8> {427 // CHECK-NEXT: shard.reduce %[[ARG0]]428 // CHECK-SAME: on @grid0 grid_axes = [0, 2]429 // CHECK-SAME: root = [1, %[[ARG1]]]430 // CHECK-SAME: : (tensor<2xi8>, index) -> tensor<2xi8>431 %0 = shard.reduce %arg0 on @grid0 grid_axes = [0, 2]432 root = [1, %arg1]433 : (tensor<2xi8>, index) -> tensor<2xi8>434 return %0 : tensor<2xi8>435}436 437// CHECK-LABEL: func @reduce_different_return_element_type438func.func @reduce_different_return_element_type(439 // CHECK-SAME: %[[ARG:.*]]: tensor<2xi8>440 %arg0 : tensor<2xi8>) -> tensor<2xi16> {441 // CHECK-NEXT: shard.reduce %[[ARG]]442 // CHECK-SAME: on @grid0 grid_axes = [0, 2]443 // CHECK-SAME: root = [0, 1]444 // CHECK-SAME: : (tensor<2xi8>) -> tensor<2xi16>445 %0 = shard.reduce %arg0 on @grid0 grid_axes = [0, 2]446 root = [0, 1]447 : (tensor<2xi8>) -> tensor<2xi16>448 return %0 : tensor<2xi16>449}450 451// CHECK-LABEL: func @reduce_scatter_static_dimensions452func.func @reduce_scatter_static_dimensions(453 // CHECK-SAME: %[[ARG:.*]]: tensor<3x4xf32>454 %arg0 : tensor<3x4xf32>) -> tensor<3x1xf64> {455 // CHECK-NEXT: shard.reduce_scatter %[[ARG]]456 // CHECK-SAME: on @grid0 grid_axes = [2] reduction = max scatter_axis = 1457 // CHECK-SAME: : tensor<3x4xf32> -> tensor<3x1xf64>458 %0 = shard.reduce_scatter %arg0 on @grid0 grid_axes = [2]459 reduction = max scatter_axis = 1460 : tensor<3x4xf32> -> tensor<3x1xf64>461 return %0 : tensor<3x1xf64>462}463 464// CHECK-LABEL: func @reduce_scatter_dynamic_dimensions465func.func @reduce_scatter_dynamic_dimensions(466 // CHECK-SAME: %[[ARG:.*]]: tensor<?xf32>467 %arg0 : tensor<?xf32>) -> tensor<?xf64> {468 // CHECK-NEXT: shard.reduce_scatter %[[ARG]]469 // CHECK-SAME: on @grid3 grid_axes = [0, 1] scatter_axis = 0470 // CHECK-SAME: : tensor<?xf32> -> tensor<?xf64>471 %0 = shard.reduce_scatter %arg0 on @grid3 grid_axes = [0, 1] scatter_axis = 0472 : tensor<?xf32> -> tensor<?xf64>473 return %0 : tensor<?xf64>474}475 476// CHECK-LABEL: func @scatter_static_dimensions477func.func @scatter_static_dimensions(478 // CHECK-SAME: %[[ARG:.*]]: tensor<3x4xf32>479 %arg0 : tensor<3x4xf32>) -> tensor<3x1xf32> {480 // CHECK-NEXT: shard.scatter %[[ARG]]481 // CHECK-SAME: on @grid0 grid_axes = [2]482 // CHECK-SAME: scatter_axis = 1 root = [1]483 // CHECK-SAME: : (tensor<3x4xf32>) -> tensor<3x1xf32>484 %0 = shard.scatter %arg0 on @grid0 grid_axes = [2]485 scatter_axis = 1 root = [1]486 : (tensor<3x4xf32>) -> tensor<3x1xf32>487 return %0 : tensor<3x1xf32>488}489 490// CHECK-LABEL: func @scatter_dynamic_dimensions491func.func @scatter_dynamic_dimensions(492 // CHECK-SAME: %[[ARG:.*]]: tensor<?xf32>493 %arg0 : tensor<?xf32>) -> tensor<?xf32> {494 // CHECK-NEXT: shard.scatter %[[ARG]]495 // CHECK-SAME: on @grid3 grid_axes = [0, 1]496 // CHECK-SAME: scatter_axis = 0 root = [1, 2]497 // CHECK-SAME: : (tensor<?xf32>) -> tensor<?xf32>498 %0 = shard.scatter %arg0 on @grid3 grid_axes = [0, 1]499 scatter_axis = 0 root = [1, 2]500 : (tensor<?xf32>) -> tensor<?xf32>501 return %0 : tensor<?xf32>502}503 504// CHECK-LABEL: func @scatter_dynamic_root505func.func @scatter_dynamic_root(506 // CHECK-SAME: %[[ARG0:.*]]: tensor<8xi8>507 %arg0 : tensor<8xi8>,508 // CHECK-SAME: %[[ARG1:.*]]: index509 %arg1 : index510 ) -> tensor<1xi8> {511 // CHECK-NEXT: shard.scatter %[[ARG0]]512 // CHECK-SAME: on @grid0 grid_axes = [0, 2]513 // CHECK-SAME: scatter_axis = 0514 // CHECK-SAME: root = [1, %[[ARG1]]]515 // CHECK-SAME: : (tensor<8xi8>, index) -> tensor<1xi8>516 %0 = shard.scatter %arg0 on @grid0 grid_axes = [0, 2]517 scatter_axis = 0518 root = [1, %arg1]519 : (tensor<8xi8>, index) -> tensor<1xi8>520 return %0 : tensor<1xi8>521}522 523// CHECK-LABEL: func @send_static_destination524func.func @send_static_destination(525 // CHECK-SAME: %[[ARG:.*]]: tensor<2xi8>526 %arg0 : tensor<2xi8>) -> tensor<2xi8> {527 // CHECK-NEXT: shard.send %[[ARG]]528 // CHECK-SAME: on @grid0 grid_axes = [0, 2]529 // CHECK-SAME: destination = [0, 1]530 // CHECK-SAME: : (tensor<2xi8>) -> tensor<2xi8>531 %0 = shard.send %arg0 on @grid0 grid_axes = [0, 2]532 destination = [0, 1]533 : (tensor<2xi8>) -> tensor<2xi8>534 return %0 : tensor<2xi8>535}536 537// CHECK-LABEL: func @send_dynamic_destination538func.func @send_dynamic_destination(539 // CHECK-SAME: %[[ARG0:.*]]: tensor<2xi8>540 %arg0 : tensor<2xi8>,541 // CHECK-SAME: %[[ARG1:.*]]: index542 %arg1 : index543 ) -> tensor<2xi8> {544 // CHECK-NEXT: shard.send %[[ARG0]]545 // CHECK-SAME: on @grid0 grid_axes = [0, 2]546 // CHECK-SAME: destination = [1, %[[ARG1]]]547 // CHECK-SAME: : (tensor<2xi8>, index) -> tensor<2xi8>548 %0 = shard.send %arg0 on @grid0 grid_axes = [0, 2]549 destination = [1, %arg1]550 : (tensor<2xi8>, index) -> tensor<2xi8>551 return %0 : tensor<2xi8>552}553 554// CHECK-LABEL: func @shift555func.func @shift(556 // CHECK-SAME: %[[ARG:.*]]: tensor<2xi8>557 %arg0 : tensor<2xi8>) -> tensor<2xi8> {558 // CHECK-NEXT: shard.shift %[[ARG]]559 // CHECK-SAME: on @grid0 grid_axes = [0, 2]560 // CHECK-SAME: shift_axis = 2 offset = -2 rotate561 // CHECK-SAME: : tensor<2xi8> -> tensor<2xi8>562 %0 = shard.shift %arg0 on @grid0 grid_axes = [0, 2]563 shift_axis = 2 offset = -2 rotate564 : tensor<2xi8> -> tensor<2xi8>565 return %0 : tensor<2xi8>566}567 568// CHECK-LABEL: func @update_halo569func.func @update_halo(570 // CHECK-SAME: %[[ARG:.*]]: memref<12x12xi8>571 %arg0 : memref<12x12xi8>) {572 // CHECK-NEXT: %[[C2:.*]] = arith.constant 2 : i64573 // CHECK-NEXT: %[[UH1:.*]] = shard.update_halo %[[ARG]] on @grid0574 // CHECK-SAME: split_axes = {{\[\[}}0]]575 // CHECK-SAME: halo_sizes = [2, %c2_i64] : memref<12x12xi8>576 %c2 = arith.constant 2 : i64577 %uh1 = shard.update_halo %arg0 on @grid0 split_axes = [[0]]578 halo_sizes = [2, %c2] : memref<12x12xi8>579 // CHECK-NEXT: %[[UH2:.*]] = shard.update_halo %[[UH1]] on @grid0580 // CHECK-SAME: split_axes = {{\[\[}}0], [1]]581 // CHECK-SAME: halo_sizes = [2, 2, %[[C2]], 2] : memref<12x12xi8>582 %uh2 = shard.update_halo %uh1 on @grid0 split_axes = [[0], [1]]583 halo_sizes = [2, 2, %c2, 2] : memref<12x12xi8>584 return585}586