930 lines · plain
1// RUN: mlir-opt -split-input-file -verify-diagnostics %s2 3// expected-error@+1 {{rank of grid is expected to be a positive integer}}4shard.grid @grid0(shape = [])5 6// -----7 8// expected-error@+1 {{custom op 'shard.grid' Failed parsing dimension list. Did you mean an empty list? It must be denoted by "[]".}}9shard.grid @grid0(shape = -1)10 11// -----12 13shard.grid @grid0(shape = 2x4)14 15func.func @grid_axis_duplicated_different_subarray(16 %arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {17 // expected-error@+1 {{grid axis duplicated}}18 %s = shard.sharding @grid0 split_axes = [[0], [0]] : !shard.sharding19 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>20 return %0 : tensor<4x8xf32>21}22 23// -----24 25shard.grid @grid0(shape = 2x4)26 27func.func @grid_axis_duplicated_same_subarray(28 %arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {29 // expected-error@+1 {{grid axis duplicated}}30 %s = shard.sharding @grid0 split_axes = [[0, 0]] : !shard.sharding31 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>32 return %0 : tensor<4x8xf32>33}34 35// -----36 37shard.grid @grid0(shape = 2x4)38 39func.func @grid_axis_negtive_in_split_part(40 %arg0 : tensor<4x8xf32>) -> tensor<4x8xf32> {41 // expected-error@+1 {{grid axis is expected to be non-negative}}42 %s = shard.sharding @grid0 split_axes = [[-1]] : !shard.sharding43 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>44 return %0 : tensor<4x8xf32>45}46 47// -----48 49func.func @sharding_attribute_invalid_nested_symbol(%arg0 : tensor<4x8xf32>) {50 // expected-error@+1 {{custom op 'shard.sharding' invalid kind of attribute specified}}51 %s = shard.sharding @a::@b split_axes = [[0]] : !shard.sharding52 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>53 return54}55 56// -----57 58func.func @sharding_attribute_invalid_halo(%arg0 : tensor<4x8xf32>) {59 // expected-error@+1 {{halo sizes must be specified for all split axes}}60 %s = shard.sharding @grid0 split_axes = [[0], [1]] halo_sizes = [1, 2] : !shard.sharding61 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>62 return63}64 65// -----66 67func.func @sharding_attribute_invalid_sizes(%arg0 : tensor<4x8xf32>) {68 // expected-error@+1 {{halo sizes and shard offsets are mutually exclusive}}69 %s = shard.sharding @grid0 split_axes = [[0]] halo_sizes = [1, 2] sharded_dims_offsets = [0, 2, 2] : !shard.sharding70 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>71 return72}73 74// -----75 76shard.grid @grid_dyn(shape = ?x?)77func.func @sharding_dyn_grid_and_sizes(%arg0 : tensor<4x8xf32>) {78 // expected-error@+1 {{sharded dims offsets are not allowed for device grids with dynamic shape}}79 %s = shard.sharding @grid_dyn split_axes = [[0]] sharded_dims_offsets = [0, 2, 2] : !shard.sharding80 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>81 return82}83 84// -----85 86shard.grid @grid0(shape = 2x4)87func.func @sharding_sizes_count(%arg0 : tensor<4x8xf32>) {88 // expected-error@+1 {{sharded dims offsets has wrong size}}89 %s = shard.sharding @grid0 split_axes = [[0], [1]] sharded_dims_offsets = [0, 2, 4, 0, 2, 4, 6] : !shard.sharding90 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>91 return92}93 94// -----95 96shard.grid @grid0(shape = 4)97func.func @sharding_sizes_decreasing(%arg0 : tensor<4x8xf32>) {98 // expected-error@+1 {{sharded dims offsets must be non-decreasing}}99 %s = shard.sharding @grid0 split_axes = [[0]] sharded_dims_offsets = [0, 2, 3, 2] : !shard.sharding100 %0 = shard.shard %arg0 to %s : tensor<4x8xf32>101 return102}103 104// -----105 106shard.grid @grid0(shape = 2x4)107 108func.func @grid_shape_grid_axis_out_of_bounds() -> (index, index) {109 // expected-error@+1 {{0-based grid axis index 2 is out of bounds. The referenced grid "grid0" is of rank 2.}}110 %0:2 = shard.grid_shape @grid0 axes = [0, 2] : index, index111 return %0#0, %0#1 : index, index112}113 114// -----115 116shard.grid @grid0(shape = 1x2x3)117 118func.func @grid_shape_duplicate_grid_axis() -> (index, index, index) {119 // expected-error@+1 {{Grid axes contains duplicate elements.}}120 %0:3 = shard.grid_shape @grid0 axes = [0, 2, 0] : index, index, index121 return %0#0, %0#1, %0#2 : index, index, index122}123 124// -----125 126shard.grid @grid0(shape = 2x4)127 128func.func @grid_shape_wrong_number_of_results() -> (index, index) {129 // expected-error@+1 {{Unexpected number of results 2. Expected 1.}}130 %0:2 = shard.grid_shape @grid0 axes = [0] : index, index131 return %0#0, %0#1 : index, index132}133 134// -----135 136shard.grid @grid0(shape = 1x2x3)137 138func.func @grid_shape_wrong_number_of_results_empty_grid_axes() -> (index, index) {139 // expected-error@+1 {{Unexpected number of results 2. Expected 3.}}140 %0:2 = shard.grid_shape @grid0 : index, index141 return %0#0, %0#1 : index, index142}143 144// -----145 146func.func @grid_shape_invalid_grid_name() -> (index) {147 // expected-error@+1 {{Undefined required grid symbol "this_grid_symbol_does_not_exist".}}148 %0 = shard.grid_shape @this_grid_symbol_does_not_exist : index149 return %0#0 : index150}151 152// -----153 154shard.grid @grid0(shape = 2x4)155 156func.func @process_multi_index_grid_axis_out_of_bounds() -> (index, index) {157 // expected-error@+1 {{0-based grid axis index 2 is out of bounds. The referenced grid "grid0" is of rank 2.}}158 %0:2 = shard.process_multi_index on @grid0 axes = [0, 2] : index, index159 return %0#0, %0#1 : index, index160}161 162// -----163 164shard.grid @grid0(shape = 1x2x3)165 166func.func @process_multi_index_duplicate_grid_axis() -> (index, index, index) {167 // expected-error@+1 {{Grid axes contains duplicate elements.}}168 %0:3 = shard.process_multi_index on @grid0 axes = [0, 2, 0] : index, index, index169 return %0#0, %0#1, %0#2 : index, index, index170}171 172// -----173 174shard.grid @grid0(shape = 2x4)175 176func.func @process_multi_index_wrong_number_of_results() -> (index, index) {177 // expected-error@+1 {{Unexpected number of results 2. Expected 1.}}178 %0:2 = shard.process_multi_index on @grid0 axes = [0] : index, index179 return %0#0, %0#1 : index, index180}181 182// -----183 184shard.grid @grid0(shape = 1x2x3)185 186func.func @process_multi_index_wrong_number_of_results_empty_grid_axes() -> (index, index) {187 // expected-error@+1 {{Unexpected number of results 2. Expected 3.}}188 %0:2 = shard.process_multi_index on @grid0 : index, index189 return %0#0, %0#1 : index, index190}191 192// -----193 194func.func @process_multi_index_invalid_grid_name() -> (index) {195 // expected-error@+1 {{Undefined required grid symbol "this_grid_symbol_does_not_exist".}}196 %0 = shard.process_multi_index on @this_grid_symbol_does_not_exist : index197 return %0 : index198}199 200// -----201 202func.func @process_linear_index_invalid_grid_name() -> (index) {203 // expected-error@+1 {{Undefined required grid symbol "this_grid_symbol_does_not_exist".}}204 %0 = shard.process_linear_index on @this_grid_symbol_does_not_exist : index205 return %0 : index206}207 208// -----209 210func.func @all_reduce_invalid_grid_symbol(211 %arg0 : tensor<4xf32>) -> tensor<4xf64> {212 // expected-error@+1 {{Undefined required grid symbol "this_grid_symbol_does_not_exist".}}213 %0 = shard.all_reduce %arg0 on @this_grid_symbol_does_not_exist reduction = sum214 : tensor<4xf32> -> tensor<4xf64>215 return %0 : tensor<4xf64>216}217 218// -----219 220shard.grid @grid0(shape = 2x4)221 222func.func @all_reduce_invalid_grid_axis(223 %arg0 : tensor<4xf32>) -> tensor<4xf64> {224 // expected-error@+1 {{0-based grid axis index 2 is out of bounds. The referenced grid "grid0" is of rank 2.}}225 %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [2] reduction = sum226 : tensor<4xf32> -> tensor<4xf64>227 return %0 : tensor<4xf64>228}229 230// -----231 232shard.grid @grid0(shape = 2x4)233 234func.func @all_reduce_duplicate_grid_axis(235 %arg0 : tensor<4xf32>) -> tensor<4xf64> {236 // expected-error@+1 {{Grid axes contains duplicate elements.}}237 %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0, 1, 0] reduction = sum238 : tensor<4xf32> -> tensor<4xf64>239 return %0 : tensor<4xf64>240}241 242// -----243 244shard.grid @grid0(shape = 2x4)245 246func.func @all_reduce_invalid_tensor_dimension_size(247 %arg0 : tensor<4xf32>) -> tensor<5xf64> {248 // expected-error@+1 {{'shard.all_reduce' op requires the same shape for all operands and results}}249 %0 = shard.all_reduce %arg0 on @grid0 : tensor<4xf32> -> tensor<5xf64>250 return %0 : tensor<5xf64>251}252 253// -----254 255func.func @all_gather_invalid_grid_symbol(256 %arg0 : tensor<4xf32>) -> tensor<4xf32> {257 // expected-error@+1 {{Undefined required grid symbol "this_grid_symbol_does_not_exist".}}258 %0 = shard.all_gather %arg0 on @this_grid_symbol_does_not_exist gather_axis = 0259 : tensor<4xf32> -> tensor<4xf32>260 return %0 : tensor<4xf32>261}262 263// -----264 265shard.grid @grid0(shape = 2x4)266 267func.func @all_gather_invalid_grid_axis(268 %arg0 : tensor<4xf32>) -> tensor<4xf32> {269 // expected-error@+1 {{0-based grid axis index 2 is out of bounds. The referenced grid "grid0" is of rank 2.}}270 %0 = shard.all_gather %arg0 on @grid0 grid_axes = [2] gather_axis = 0271 : tensor<4xf32> -> tensor<4xf32>272 return %0 : tensor<4xf32>273}274 275// -----276 277shard.grid @grid0(shape = 2x4)278 279func.func @all_reduce_duplicate_grid_axis(280 %arg0 : tensor<4xf32>) -> tensor<4xf32> {281 // expected-error@+1 {{Grid axes contains duplicate elements.}}282 %0 = shard.all_gather %arg0 on @grid0 grid_axes = [2, 2] gather_axis = 0283 : tensor<4xf32> -> tensor<4xf32>284 return %0 : tensor<4xf32>285}286 287// -----288 289shard.grid @grid0(shape = 1)290 291func.func @all_gather_invalid_non_gather_axis_dimension_size(292 %arg0 : tensor<3x4xf32>) -> tensor<3x5xf32> {293 // expected-error@+1 {{Dimension size mismatch for result axis 1. Expected 4, but got 5.}}294 %0 = shard.all_gather %arg0 on @grid0 grid_axes = [0] gather_axis = 0295 : tensor<3x4xf32> -> tensor<3x5xf32>296 return %0 : tensor<3x5xf32>297}298 299// -----300 301shard.grid @grid0(shape = 1x2)302 303func.func @all_gather_invalid_gather_axis_dimension_size(304 %arg0 : tensor<3x4xf32>) -> tensor<3x5xf32> {305 // expected-error@+1 {{Dimension size mismatch for result axis 1. Expected 8, but got 5.}}306 %0 = shard.all_gather %arg0 on @grid0 grid_axes = [1] gather_axis = 1307 : tensor<3x4xf32> -> tensor<3x5xf32>308 return %0 : tensor<3x5xf32>309}310 311// -----312 313shard.grid @grid0(shape = 1)314 315func.func @all_gather_invalid_gather_axis_dynamic_dimension(316 %arg0 : tensor<?xf32>) -> tensor<3xf32> {317 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected dynamic, but got 3.}}318 %0 = shard.all_gather %arg0 on @grid0 gather_axis = 0319 : tensor<?xf32> -> tensor<3xf32>320 return %0 : tensor<3xf32>321}322 323// -----324 325shard.grid @grid0(shape = 1)326 327func.func @all_gather_invalid_gather_axis(328 %arg0 : tensor<3xf32>) -> tensor<3xf32> {329 // expected-error@+1 {{Gather axis 1 is out of bounds [0, 1).}}330 %0 = shard.all_gather %arg0 on @grid0 grid_axes = [0] gather_axis = 1331 : tensor<3xf32> -> tensor<3xf32>332 return %0 : tensor<3xf32>333}334 335// -----336 337shard.grid @grid0(shape = 1)338 339func.func @all_gather_invalid_negative_gather_axis(340 %arg0 : tensor<3xf32>) -> tensor<3xf32> {341 // expected-error@+1 {{Gather axis -1 is out of bounds [0, 1).}}342 %0 = shard.all_gather %arg0 on @grid0 grid_axes = [0] gather_axis = -1343 : tensor<3xf32> -> tensor<3xf32>344 return %0 : tensor<3xf32>345}346 347// -----348 349shard.grid @grid0(shape = 3)350 351func.func @all_slice_duplicate_grid_axis(352 %arg0 : tensor<?xf32>) -> tensor<?xf32> {353 // expected-error@+1 {{Grid axes contains duplicate elements.}}354 %0 = shard.all_slice %arg0 on @grid0 grid_axes = [0, 0]355 slice_axis = 0356 : tensor<?xf32> -> tensor<?xf32>357 return %0 : tensor<?xf32>358}359 360// -----361 362shard.grid @grid0(shape = 3)363 364func.func @all_slice_invalid_dynamic_dimension(365 %arg0 : tensor<?xf32>) -> tensor<2xf32> {366 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected dynamic, but got 2.}}367 %0 = shard.all_slice %arg0 on @grid0368 slice_axis = 0369 : tensor<?xf32> -> tensor<2xf32>370 return %0 : tensor<2xf32>371}372 373// -----374 375shard.grid @grid0(shape = 3)376 377func.func @all_slice_invalid_static_dimension_size(378 %arg0 : tensor<3xf32>) -> tensor<2xf32> {379 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected 1, but got 2.}}380 %0 = shard.all_slice %arg0 on @grid0 grid_axes = [0]381 slice_axis = 0382 : tensor<3xf32> -> tensor<2xf32>383 return %0 : tensor<2xf32>384}385 386// -----387 388shard.grid @grid0(shape = 3)389 390func.func @all_slice_invalid_operand_static_dimension_size(391 %arg0 : tensor<4xf32>) -> tensor<?xf32> {392 // expected-error@+1 {{Operand dimension size 4 is not divisible by collective device group size 3 for tensor axis 0.}}393 %0 = shard.all_slice %arg0 on @grid0 grid_axes = [0]394 slice_axis = 0395 : tensor<4xf32> -> tensor<?xf32>396 return %0 : tensor<?xf32>397}398 399// -----400 401func.func @all_to_all_invalid_grid_symbol(402 %arg0 : tensor<3x6xi8>) -> tensor<3x6xi8> {403 // expected-error@+1 {{Undefined required grid symbol "this_grid_symbol_does_not_exist".}}404 %0 = shard.all_to_all %arg0 on @this_grid_symbol_does_not_exist405 split_axis = 1 concat_axis = 0406 : tensor<3x6xi8> -> tensor<3x6xi8>407 return %0 : tensor<3x6xi8>408}409 410// -----411 412shard.grid @grid0(shape = 1)413 414func.func @all_to_all_duplicate_grid_axis(415 %arg0 : tensor<3x6xi8>) -> tensor<3x6xi8> {416 // expected-error@+1 {{Grid axes contains duplicate elements.}}417 %0 = shard.all_to_all %arg0 on @grid0 grid_axes = [0, 0]418 split_axis = 0 concat_axis = 0419 : tensor<3x6xi8> -> tensor<3x6xi8>420 return %0 : tensor<3x6xi8>421}422 423// -----424 425shard.grid @grid0(shape = ?x1)426 427func.func @all_to_all_invalid_non_dynamic_result_dimension_induced_by_dynamic_device_group(428 %arg0 : tensor<3x6xi8>) -> tensor<3x6xi8> {429 // expected-error@+1 {{Dimension size mismatch for result axis 1. Expected dynamic, but got 6.}}430 %0 = shard.all_to_all %arg0 on @grid0 grid_axes = [0]431 split_axis = 0 concat_axis = 1432 : tensor<3x6xi8> -> tensor<3x6xi8>433 return %0 : tensor<3x6xi8>434}435 436// -----437 438shard.grid @grid0(shape = 1x1)439 440func.func @all_to_all_invalid_non_dynamic_result_split_dimension_induced_by_dynamic_operand_dimension(441 %arg0 : tensor<?x6xi8>) -> tensor<3x?xi8> {442 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected dynamic, but got 3.}}443 %0 = shard.all_to_all %arg0 on @grid0 grid_axes = [1]444 split_axis = 0 concat_axis = 1445 : tensor<?x6xi8> -> tensor<3x?xi8>446 return %0 : tensor<3x?xi8>447}448 449// -----450 451shard.grid @grid0(shape = 1x1)452 453func.func @all_to_all_invalid_non_dynamic_result_concat_dimension_induced_by_dynamic_operand_dimension(454 %arg0 : tensor<3x?xi8>) -> tensor<?x3xi8> {455 // expected-error@+1 {{Dimension size mismatch for result axis 1. Expected dynamic, but got 3.}}456 %0 = shard.all_to_all %arg0 on @grid0 grid_axes = [1]457 split_axis = 0 concat_axis = 1458 : tensor<3x?xi8> -> tensor<?x3xi8>459 return %0 : tensor<?x3xi8>460}461 462// -----463 464shard.grid @grid0(shape = 3)465 466func.func @all_to_all_invalid_non_dynamic_result_concat_dimension_size(467 %arg0 : tensor<3x2xi8>) -> tensor<1x7xi8> {468 // expected-error@+1 {{Dimension size mismatch for result axis 1. Expected 6, but got 7.}}469 %0 = shard.all_to_all %arg0 on @grid0 grid_axes = [0]470 split_axis = 0 concat_axis = 1471 : tensor<3x2xi8> -> tensor<1x7xi8>472 return %0 : tensor<1x7xi8>473}474 475// -----476 477shard.grid @grid0(shape = 3)478 479func.func @all_to_all_invalid_non_dynamic_result_split_dimension_size(480 %arg0 : tensor<3x2xi8>) -> tensor<2x6xi8> {481 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected 1, but got 2.}}482 %0 = shard.all_to_all %arg0 on @grid0 grid_axes = [0]483 split_axis = 0 concat_axis = 1484 : tensor<3x2xi8> -> tensor<2x6xi8>485 return %0 : tensor<2x6xi8>486}487 488// -----489 490shard.grid @grid0(shape = 3x?)491 492func.func @broadcast_root_dimension_out_of_bounds(493 %arg0 : tensor<2xi8>) -> tensor<2xi8> {494 // expected-error@+1 {{Out of bounds coordinate 0 for in-group device "root". Got 3, but expected value in the range [0, 2].}}495 %0 = shard.broadcast %arg0 on @grid0 grid_axes = [0]496 root = [3]497 : (tensor<2xi8>) -> tensor<2xi8>498 return %0 : tensor<2xi8>499}500 501// -----502 503shard.grid @grid0(shape = 3x?)504 505func.func @broadcast_root_wrong_number_dimensions(506 %arg0 : tensor<2xi8>) -> tensor<2xi8> {507 // expected-error@+1 {{In-group device "root" has unexpected multi-index size 2. Expected 1.}}508 %0 = shard.broadcast %arg0 on @grid0 grid_axes = [0]509 root = [2, 2]510 : (tensor<2xi8>) -> tensor<2xi8>511 return %0 : tensor<2xi8>512}513 514// -----515 516shard.grid @grid0(shape = 3x?)517 518func.func @broadcast_different_input_and_result_type(519 %arg0 : tensor<2xi8>) -> tensor<2xi16> {520 // expected-error@+1 {{'shard.broadcast' op failed to verify that all of {input, result} have same element type}}521 %0 = shard.broadcast %arg0 on @grid0 grid_axes = [0]522 root = [2]523 : (tensor<2xi8>) -> tensor<2xi16>524 return %0 : tensor<2xi16>525}526 527// -----528 529shard.grid @grid0(shape = 1)530 531func.func @gather_wrong_return_element_type(532 %arg0 : tensor<1xf32>) -> tensor<1xi8> {533 // expected-error@+1 {{'shard.gather' op failed to verify that all of {input, result} have same element type}}534 %0 = shard.gather %arg0 on @grid0 grid_axes = [0] gather_axis = 0 root = [0]535 : (tensor<1xf32>) -> tensor<1xi8>536 return %0 : tensor<1xi8>537}538 539// -----540 541shard.grid @grid0(shape = 1)542 543func.func @gather_invalid_non_gather_axis_dimension_size(544 %arg0 : tensor<3x4xf32>) -> tensor<3x5xf32> {545 // expected-error@+1 {{Dimension size mismatch for result axis 1. Expected 4, but got 5.}}546 %0 = shard.gather %arg0 on @grid0 grid_axes = [0] gather_axis = 0 root = [0]547 : (tensor<3x4xf32>) -> tensor<3x5xf32>548 return %0 : tensor<3x5xf32>549}550 551// -----552 553shard.grid @grid0(shape = 1x2)554 555func.func @gather_invalid_gather_axis_dimension_size(556 %arg0 : tensor<3x4xf32>) -> tensor<3x5xf32> {557 // expected-error@+1 {{Dimension size mismatch for result axis 1. Expected 8, but got 5.}}558 %0 = shard.gather %arg0 on @grid0 grid_axes = [1] gather_axis = 1 root = [0]559 : (tensor<3x4xf32>) -> tensor<3x5xf32>560 return %0 : tensor<3x5xf32>561}562 563// -----564 565shard.grid @grid0(shape = 1)566 567func.func @gather_invalid_gather_axis_dynamic_dimension(568 %arg0 : tensor<?xf32>) -> tensor<3xf32> {569 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected dynamic, but got 3.}}570 %0 = shard.gather %arg0 on @grid0 gather_axis = 0 root = []571 : (tensor<?xf32>) -> tensor<3xf32>572 return %0 : tensor<3xf32>573}574 575// -----576 577shard.grid @grid0(shape = 1)578 579func.func @gather_invalid_gather_axis(580 %arg0 : tensor<3xf32>) -> tensor<3xf32> {581 // expected-error@+1 {{Gather axis 1 is out of bounds [0, 1).}}582 %0 = shard.gather %arg0 on @grid0 grid_axes = [0] gather_axis = 1 root = [0]583 : (tensor<3xf32>) -> tensor<3xf32>584 return %0 : tensor<3xf32>585}586 587// -----588 589shard.grid @grid0(shape = 1)590 591func.func @gather_invalid_negative_gather_axis(592 %arg0 : tensor<3xf32>) -> tensor<3xf32> {593 // expected-error@+1 {{Gather axis -1 is out of bounds [0, 1).}}594 %0 = shard.gather %arg0 on @grid0 grid_axes = [0] gather_axis = -1 root = [0]595 : (tensor<3xf32>) -> tensor<3xf32>596 return %0 : tensor<3xf32>597}598 599// -----600 601shard.grid @grid0(shape = 3x?)602 603func.func @gather_root_dimension_out_of_bounds(604 %arg0 : tensor<2xi8>) -> tensor<6xi8> {605 // expected-error@+1 {{Out of bounds coordinate 0 for in-group device "root". Got 3, but expected value in the range [0, 2].}}606 %0 = shard.gather %arg0 on @grid0 grid_axes = [0] gather_axis = 0607 root = [3]608 : (tensor<2xi8>) -> tensor<6xi8>609 return %0 : tensor<6xi8>610}611 612// -----613 614shard.grid @grid0(shape = 3x?)615 616func.func @gather_root_wrong_number_dimensions(617 %arg0 : tensor<2xi8>) -> tensor<2xi8> {618 // expected-error@+1 {{In-group device "root" has unexpected multi-index size 2. Expected 1.}}619 %0 = shard.gather %arg0 on @grid0 grid_axes = [0] gather_axis = 0620 root = [2, 2]621 : (tensor<2xi8>) -> tensor<2xi8>622 return %0 : tensor<2xi8>623}624 625// -----626 627shard.grid @grid0(shape = 3x?)628 629func.func @receive_source_dimension_out_of_bounds(630 %arg0 : tensor<2xi8>) -> tensor<2xi8> {631 // expected-error@+1 {{Out of bounds coordinate 0 for in-group device "source". Got 3, but expected value in the range [0, 2].}}632 %0 = shard.recv %arg0 on @grid0 grid_axes = [0]633 source = [3]634 : (tensor<2xi8>) -> tensor<2xi8>635 return %0 : tensor<2xi8>636}637 638// -----639 640shard.grid @grid0(shape = 3x?)641 642func.func @receive_source_wrong_number_dimensions(643 %arg0 : tensor<2xi8>) -> tensor<2xi8> {644 // expected-error@+1 {{In-group device "source" has unexpected multi-index size 2. Expected 1.}}645 %0 = shard.recv %arg0 on @grid0 grid_axes = [0]646 source = [2, 2]647 : (tensor<2xi8>) -> tensor<2xi8>648 return %0 : tensor<2xi8>649}650 651// -----652 653shard.grid @grid0(shape = 3x?)654 655func.func @receive_different_input_and_result_type(656 %arg0 : tensor<2xi8>) -> tensor<2xi16> {657 // expected-error@+1 {{'shard.recv' op failed to verify that all of {input, result} have same element type}}658 %0 = shard.recv %arg0 on @grid0 grid_axes = [0]659 source = [2]660 : (tensor<2xi8>) -> tensor<2xi16>661 return %0 : tensor<2xi16>662}663 664// -----665 666shard.grid @grid0(shape = 3x?)667 668func.func @reduce_root_dimension_out_of_bounds(669 %arg0 : tensor<2xi8>) -> tensor<2xi8> {670 // expected-error@+1 {{Out of bounds coordinate 0 for in-group device "root". Got 3, but expected value in the range [0, 2].}}671 %0 = shard.reduce %arg0 on @grid0 grid_axes = [0]672 root = [3]673 : (tensor<2xi8>) -> tensor<2xi8>674 return %0 : tensor<2xi8>675}676 677// -----678 679shard.grid @grid0(shape = 3x?)680 681func.func @reduce_root_wrong_number_dimensions(682 %arg0 : tensor<2xi8>) -> tensor<2xi8> {683 // expected-error@+1 {{In-group device "root" has unexpected multi-index size 2. Expected 1.}}684 %0 = shard.reduce %arg0 on @grid0 grid_axes = [0]685 root = [2, 2]686 : (tensor<2xi8>) -> tensor<2xi8>687 return %0 : tensor<2xi8>688}689 690// -----691 692shard.grid @grid0(shape = 3x?)693 694func.func @reduce_different_input_and_result_shape(695 %arg0 : tensor<2xi8>) -> tensor<3xi16> {696 // expected-error@+1 {{'shard.reduce' op failed to verify that all of {input, result} have same shape}}697 %0 = shard.reduce %arg0 on @grid0 grid_axes = [0]698 root = [2]699 : (tensor<2xi8>) -> tensor<3xi16>700 return %0 : tensor<3xi16>701}702 703// -----704 705shard.grid @grid0(shape = 3)706 707func.func @reduce_scatter_duplicate_grid_axis(708 %arg0 : tensor<?xf32>) -> tensor<?xf64> {709 // expected-error@+1 {{Grid axes contains duplicate elements.}}710 %0 = shard.reduce_scatter %arg0 on @grid0 grid_axes = [0, 0] scatter_axis = 0711 : tensor<?xf32> -> tensor<?xf64>712 return %0 : tensor<?xf64>713}714 715// -----716 717shard.grid @grid0(shape = 3)718 719func.func @reduce_scatter_invalid_dynamic_dimension(720 %arg0 : tensor<?xf32>) -> tensor<2xf64> {721 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected dynamic, but got 2.}}722 %0 = shard.reduce_scatter %arg0 on @grid0 scatter_axis = 0723 : tensor<?xf32> -> tensor<2xf64>724 return %0 : tensor<2xf64>725}726 727// -----728 729shard.grid @grid0(shape = 3)730 731func.func @reduce_scatter_invalid_static_dimension_size(732 %arg0 : tensor<3xf32>) -> tensor<2xf64> {733 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected 1, but got 2.}}734 %0 = shard.reduce_scatter %arg0 on @grid0 grid_axes = [0] scatter_axis = 0735 : tensor<3xf32> -> tensor<2xf64>736 return %0 : tensor<2xf64>737}738 739// -----740 741shard.grid @grid0(shape = 3)742 743func.func @reduce_scatter_invalid_operand_static_dimension_size(744 %arg0 : tensor<4xf32>) -> tensor<?xf64> {745 // expected-error@+1 {{Operand dimension size 4 is not divisible by collective device group size 3 for tensor axis 0.}}746 %0 = shard.reduce_scatter %arg0 on @grid0 grid_axes = [0] scatter_axis = 0747 : tensor<4xf32> -> tensor<?xf64>748 return %0 : tensor<?xf64>749}750 751// -----752 753shard.grid @grid0(shape = 3)754 755func.func @scatter_duplicate_grid_axis(756 %arg0 : tensor<?xf32>) -> tensor<?xf32> {757 // expected-error@+1 {{Grid axes contains duplicate elements.}}758 %0 = shard.scatter %arg0 on @grid0 grid_axes = [0, 0]759 scatter_axis = 0 root = [0, 0]760 : (tensor<?xf32>) -> tensor<?xf32>761 return %0 : tensor<?xf32>762}763 764// -----765 766shard.grid @grid0(shape = 3)767 768func.func @scatter_invalid_dynamic_dimension(769 %arg0 : tensor<?xf32>) -> tensor<2xf32> {770 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected dynamic, but got 2.}}771 %0 = shard.scatter %arg0 on @grid0772 scatter_axis = 0 root = []773 : (tensor<?xf32>) -> tensor<2xf32>774 return %0 : tensor<2xf32>775}776 777// -----778 779shard.grid @grid0(shape = 3)780 781func.func @scatter_invalid_static_dimension_size(782 %arg0 : tensor<3xf32>) -> tensor<2xf32> {783 // expected-error@+1 {{Dimension size mismatch for result axis 0. Expected 1, but got 2.}}784 %0 = shard.scatter %arg0 on @grid0 grid_axes = [0]785 scatter_axis = 0 root = [1]786 : (tensor<3xf32>) -> tensor<2xf32>787 return %0 : tensor<2xf32>788}789 790// -----791 792shard.grid @grid0(shape = 3)793 794func.func @scatter_invalid_operand_static_dimension_size(795 %arg0 : tensor<4xf32>) -> tensor<?xf32> {796 // expected-error@+1 {{Operand dimension size 4 is not divisible by collective device group size 3 for tensor axis 0.}}797 %0 = shard.scatter %arg0 on @grid0 grid_axes = [0]798 scatter_axis = 0 root = [1]799 : (tensor<4xf32>) -> tensor<?xf32>800 return %0 : tensor<?xf32>801}802 803// -----804 805shard.grid @grid0(shape = 3x?)806 807func.func @scatter_root_dimension_out_of_bounds(808 %arg0 : tensor<3xi8>) -> tensor<1xi8> {809 // expected-error@+1 {{Out of bounds coordinate 0 for in-group device "root". Got 3, but expected value in the range [0, 2].}}810 %0 = shard.scatter %arg0 on @grid0 grid_axes = [0]811 scatter_axis = 0 root = [3]812 : (tensor<3xi8>) -> tensor<1xi8>813 return %0 : tensor<1xi8>814}815 816// -----817 818shard.grid @grid0(shape = 3x?)819 820func.func @scatter_root_wrong_number_dimensions(821 %arg0 : tensor<3xi8>) -> tensor<1xi8> {822 // expected-error@+1 {{In-group device "root" has unexpected multi-index size 2. Expected 1.}}823 %0 = shard.scatter %arg0 on @grid0 grid_axes = [0]824 scatter_axis = 0 root = [2, 2]825 : (tensor<3xi8>) -> tensor<1xi8>826 return %0 : tensor<1xi8>827}828 829// -----830 831shard.grid @grid0(shape = 3x?)832 833func.func @send_destination_dimension_out_of_bounds(834 %arg0 : tensor<2xi8>) -> tensor<2xi8> {835 // expected-error@+1 {{Out of bounds coordinate 0 for in-group device "destination". Got 3, but expected value in the range [0, 2].}}836 %0 = shard.send %arg0 on @grid0 grid_axes = [0]837 destination = [3]838 : (tensor<2xi8>) -> tensor<2xi8>839 return %0 : tensor<2xi8>840}841 842// -----843 844shard.grid @grid0(shape = 3x?)845 846func.func @send_destination_wrong_number_dimensions(847 %arg0 : tensor<2xi8>) -> tensor<2xi8> {848 // expected-error@+1 {{In-group device "destination" has unexpected multi-index size 2. Expected 1.}}849 %0 = shard.send %arg0 on @grid0 grid_axes = [0]850 destination = [2, 2]851 : (tensor<2xi8>) -> tensor<2xi8>852 return %0 : tensor<2xi8>853}854 855// -----856 857shard.grid @grid0(shape = 3x?)858 859func.func @send_different_input_and_result_type(860 %arg0 : tensor<2xi8>) -> tensor<2xi16> {861 // expected-error@+1 {{'shard.send' op failed to verify that all of {input, result} have same element type}}862 %0 = shard.send %arg0 on @grid0 grid_axes = [0]863 destination = [2]864 : (tensor<2xi8>) -> tensor<2xi16>865 return %0 : tensor<2xi16>866}867 868// -----869 870func.func @shift_invalid_grid_symbol(871 %arg0 : tensor<4xi8>) -> tensor<4xi8> {872 // expected-error@+1 {{Undefined required grid symbol "this_grid_symbol_does_not_exist".}}873 %0 = shard.shift %arg0 on @this_grid_symbol_does_not_exist874 shift_axis = 0 offset = -2875 : tensor<4xi8> -> tensor<4xi8>876 return %0 : tensor<4xi8>877}878 879// -----880 881shard.grid @grid0(shape = 2x4)882 883func.func @shift_invalid_grid_axis(884 %arg0 : tensor<4xi8>) -> tensor<4xi8> {885 // expected-error@+1 {{0-based grid axis index 2 is out of bounds. The referenced grid "grid0" is of rank 2.}}886 %0 = shard.shift %arg0 on @grid0 grid_axes = [2]887 shift_axis = 2 offset = -2888 : tensor<4xi8> -> tensor<4xi8>889 return %0 : tensor<4xi8>890}891 892// -----893 894shard.grid @grid0(shape = 2x4)895 896func.func @shift_duplicate_grid_axis(897 %arg0 : tensor<4xi8>) -> tensor<4xi8> {898 // expected-error@+1 {{Grid axes contains duplicate elements.}}899 %0 = shard.shift %arg0 on @grid0 grid_axes = [0, 1, 0]900 shift_axis = 0 offset = -2901 : tensor<4xi8> -> tensor<4xi8>902 return %0 : tensor<4xi8>903}904 905// -----906 907shard.grid @grid0(shape = 2x4)908 909func.func @shift_invalid_tensor_dimension_size(910 %arg0 : tensor<4xi8>) -> tensor<5xi8> {911 // expected-error@+1 {{'shard.shift' op requires the same shape for all operands and results}}912 %0 = shard.shift %arg0 on @grid0 grid_axes = [0]913 shift_axis = 0 offset = 2914 : tensor<4xi8> -> tensor<5xi8>915 return %0 : tensor<5xi8>916}917 918// -----919 920shard.grid @grid0(shape = 2x4)921 922func.func @shift_invalid_shift_axis(923 %arg0 : tensor<4xi8>) -> tensor<4xi8> {924 // expected-error@+1 {{Invalid shift axis 1. It must be one of the grouping grid axes.}}925 %0 = shard.shift %arg0 on @grid0 grid_axes = [0]926 shift_axis = 1 offset = 2927 : tensor<4xi8> -> tensor<4xi8>928 return %0 : tensor<4xi8>929}930