76 lines · plain
1// RUN: mlir-opt %s --convert-shard-to-mpi -canonicalize | FileCheck %s2 3module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 24> } {4 5 // CHECK: shard.grid @grid06 shard.grid @grid0(shape = 3x4x5)7 8 // Notice: comm_world_rank/linear index 24 is multiindex [1, 0, 4] in @grid09 10 // all shards are equal11 // CHECK-LABEL: func.func @shard_shape_equal() -> (index, index, index) {12 func.func @shard_shape_equal() -> (index, index, index) {13 %sharding = shard.sharding @grid0 split_axes = [[0], [1], [2]] : !shard.sharding14 %0:3 = shard.process_multi_index on @grid0 : index, index, index15 %c9 = arith.constant 9 : index16 %c12 = arith.constant 12 : index17 // CHECK: [[vc3:%.*]] = arith.constant 3 : index18 %1:3 = shard.shard_shape dims = [%c9, %c12, 15] sharding = %sharding device = [%0#0, %0#1, %0#2] : index, index, index19 // CHECK: return [[vc3]], [[vc3]], [[vc3]] : index, index, index20 return %1#0, %1#1, %1#2 : index, index, index21 }22 23 // last shard in last dim gets an extra element24 // CHECK-LABEL: func.func @shard_shape_odd_1() -> (index, index, index) {25 func.func @shard_shape_odd_1() -> (index, index, index) {26 %sharding = shard.sharding @grid0 split_axes = [[0], [1], [2]] : !shard.sharding27 %0:3 = shard.process_multi_index on @grid0 : index, index, index28 %c9 = arith.constant 9 : index29 %c12 = arith.constant 12 : index30 // CHECK-DAG: [[vc3:%.*]] = arith.constant 3 : index31 // CHECK-DAG: [[vc4:%.*]] = arith.constant 4 : index32 %1:3 = shard.shard_shape dims = [%c9, %c12, 16] sharding = %sharding device = [%0#0, %0#1, %0#2] : index, index, index33 // CHECK: return [[vc3]], [[vc3]], [[vc4]] : index, index, index34 return %1#0, %1#1, %1#2 : index, index, index35 }36 37 // In the second dimension the shard sizes are now [3 4 4 4]38 // CHECK-LABEL: func.func @shard_shape_odd_2() -> (index, index, index) {39 func.func @shard_shape_odd_2() -> (index, index, index) {40 %sharding = shard.sharding @grid0 split_axes = [[0], [1], [2]] : !shard.sharding41 %0:3 = shard.process_multi_index on @grid0 : index, index, index42 %c9 = arith.constant 9 : index43 // CHECK: [[vc3:%.*]] = arith.constant 3 : index44 %1:3 = shard.shard_shape dims = [%c9, 15, 15] sharding = %sharding device = [%0#0, %0#1, %0#2] : index, index, index45 // CHECK: return [[vc3]], [[vc3]], [[vc3]] : index, index, index46 return %1#0, %1#1, %1#2 : index, index, index47 }48 49 // In the first dimension the shard sizes are now [3 4 4]50 // CHECK-LABEL: func.func @shard_shape_odd_3() -> (index, index, index) {51 func.func @shard_shape_odd_3() -> (index, index, index) {52 %sharding = shard.sharding @grid0 split_axes = [[0], [1], [2]] : !shard.sharding53 %0:3 = shard.process_multi_index on @grid0 : index, index, index54 // CHECK-DAG: [[vc3:%.*]] = arith.constant 3 : index55 // CHECK-DAG: [[vc4:%.*]] = arith.constant 4 : index56 %1:3 = shard.shard_shape dims = [11, 12, 15] sharding = %sharding device = [%0#0, %0#1, %0#2] : index, index, index57 // CHECK: return [[vc4]], [[vc3]], [[vc3]] : index, index, index58 return %1#0, %1#1, %1#2 : index, index, index59 }60 61 // extract from sharded_dims_offsets62 // CHECK-LABEL: func.func @shard_shape_sharded_dims_offs() -> (index, index, index) {63 func.func @shard_shape_sharded_dims_offs() -> (index, index, index) {64 %sharding = shard.sharding @grid0 split_axes = [[0], [1], [2]]65 sharded_dims_offsets = [0, 1, 4, 9, 0, 2, 6, 12, 12, 0, 3, 6, 9, 12, 15]: !shard.sharding66 %0:3 = shard.process_multi_index on @grid0 : index, index, index67 %c9 = arith.constant 9 : index68 %c12 = arith.constant 12 : index69 // CHECK: [[vc3:%.*]] = arith.constant 3 : index70 // CHECK: [[vc2:%.*]] = arith.constant 2 : index71 %1:3 = shard.shard_shape dims = [%c9, %c12, 15] sharding = %sharding device = [%0#0, %0#1, %0#2] : index, index, index72 // CHECK: return [[vc3]], [[vc2]], [[vc3]] : index, index, index73 return %1#0, %1#1, %1#2 : index, index, index74 }75}76