249 lines · plain
1// RUN: mlir-opt --canonicalize %s | FileCheck %s2 3shard.grid @grid0(shape = 2x4)4 5// CHECK-LABEL: func @all_reduce_empty_grid_axes6func.func @all_reduce_empty_grid_axes(7// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>8 %arg0 : tensor<4xf32>) -> tensor<4xf32> {9// CHECK-NOT: shard.all_reduce10 %0 = shard.all_reduce %arg0 on @grid011 grid_axes = []12 : tensor<4xf32> -> tensor<4xf32>13// CHECK: return %[[ARG]]14 return %0 : tensor<4xf32>15}16 17// CHECK-LABEL: func @all_reduce_empty_grid_axes_different_return_type18func.func @all_reduce_empty_grid_axes_different_return_type(19 %arg0 : tensor<4xf32>) -> tensor<4xf64> {20// CHECK: shard.all_reduce21 %0 = shard.all_reduce %arg0 on @grid022// CHECK-NOT: grid_axes23 grid_axes = []24 : tensor<4xf32> -> tensor<4xf64>25 return %0 : tensor<4xf64>26}27 28// CHECK-LABEL: func @all_reduce_default_reduction29func.func @all_reduce_default_reduction(30 %arg0 : tensor<4xf32>) -> tensor<4xf64> {31 %0 = shard.all_reduce %arg0 on @grid032 grid_axes = [0]33// CHECK-NOT: reduction34 reduction = sum35 : tensor<4xf32> -> tensor<4xf64>36 return %0 : tensor<4xf64>37}38 39// CHECK-LABEL: func @all_to_all_empty_grid_axes40func.func @all_to_all_empty_grid_axes(41// CHECK-SAME: %[[ARG:.*]]: tensor<8xf32>42 %arg0 : tensor<8xf32>) -> tensor<8xf32> {43// CHECK-NOT: shard.all_to_all44 %0 = shard.all_to_all %arg0 on @grid045 grid_axes = []46 split_axis = 047 concat_axis = 048 : tensor<8xf32> -> tensor<8xf32>49// CHECK: return %[[ARG]]50 return %0 : tensor<8xf32>51}52 53// CHECK-LABEL: func @all_gather_empty_grid_axes54func.func @all_gather_empty_grid_axes(55// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>56 %arg0 : tensor<4xf32>) -> tensor<4xf32> {57// CHECK-NOT: shard.all_gather58 %0 = shard.all_gather %arg0 on @grid059 grid_axes = []60 gather_axis = 061 : tensor<4xf32> -> tensor<4xf32>62// CHECK: return %[[ARG]]63 return %0 : tensor<4xf32>64}65 66// CHECK-LABEL: func @all_slice_empty_grid_axes67func.func @all_slice_empty_grid_axes(68// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>69 %arg0 : tensor<4xf32>) -> tensor<4xf32> {70// CHECK-NOT: shard.scatter71 %0 = shard.all_slice %arg0 on @grid072 grid_axes = []73 slice_axis = 074 : tensor<4xf32> -> tensor<4xf32>75// CHECK: return %[[ARG]]76 return %0 : tensor<4xf32>77}78 79// CHECK-LABEL: func @broadcast_empty_grid_axes80func.func @broadcast_empty_grid_axes(81// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>82 %arg0 : tensor<4xf32>) -> tensor<4xf32> {83// CHECK-NOT: shard.broadcast84 %0 = shard.broadcast %arg0 on @grid085 grid_axes = []86 root = []87 : (tensor<4xf32>) -> tensor<4xf32>88// CHECK: return %[[ARG]]89 return %0 : tensor<4xf32>90}91 92// CHECK-LABEL: func @gather_empty_grid_axes93func.func @gather_empty_grid_axes(94// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>95 %arg0 : tensor<4xf32>) -> tensor<4xf32> {96// CHECK-NOT: shard.gather97 %0 = shard.gather %arg0 on @grid098 grid_axes = []99 gather_axis = 0100 root = []101 : (tensor<4xf32>) -> tensor<4xf32>102// CHECK: return %[[ARG]]103 return %0 : tensor<4xf32>104}105 106// CHECK-LABEL: func @receive_empty_grid_axes107func.func @receive_empty_grid_axes(108// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>109 %arg0 : tensor<4xf32>) -> tensor<4xf32> {110// CHECK-NOT: shard.recv111 %0 = shard.recv %arg0 on @grid0112 grid_axes = []113 : (tensor<4xf32>) -> tensor<4xf32>114// CHECK: return %[[ARG]]115 return %0 : tensor<4xf32>116}117 118// CHECK-LABEL: func @reduce_empty_grid_axes119func.func @reduce_empty_grid_axes(120// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>121 %arg0 : tensor<4xf32>) -> tensor<4xf32> {122// CHECK-NOT: shard.reduce123 %0 = shard.reduce %arg0 on @grid0124 grid_axes = []125 root = []126 : (tensor<4xf32>) -> tensor<4xf32>127// CHECK: return %[[ARG]]128 return %0 : tensor<4xf32>129}130 131// CHECK-LABEL: func @reduce_scatter_empty_grid_axes132func.func @reduce_scatter_empty_grid_axes(133// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>134 %arg0 : tensor<4xf32>) -> tensor<4xf32> {135// CHECK-NOT: shard.reduce_scatter136 %0 = shard.reduce_scatter %arg0 on @grid0137 grid_axes = []138 scatter_axis = 0139 : tensor<4xf32> -> tensor<4xf32>140// CHECK: return %[[ARG]]141 return %0 : tensor<4xf32>142}143 144// CHECK-LABEL: func @reduce_scatter_empty_grid_axes_different_return_type145func.func @reduce_scatter_empty_grid_axes_different_return_type(146 %arg0 : tensor<4xf32>) -> tensor<4xf64> {147// CHECK: shard.reduce_scatter148 %0 = shard.reduce_scatter %arg0 on @grid0149// CHECK-NOT: grid_axes150 grid_axes = []151 scatter_axis = 0152 : tensor<4xf32> -> tensor<4xf64>153 return %0 : tensor<4xf64>154}155 156// CHECK-LABEL: func @reduce_scatter_default_reduction157func.func @reduce_scatter_default_reduction(158 %arg0 : tensor<4xf32>) -> tensor<2xf64> {159 %0 = shard.reduce_scatter %arg0 on @grid0160 grid_axes = [0]161// CHECK-NOT: reduction162 reduction = sum163 scatter_axis = 0164 : tensor<4xf32> -> tensor<2xf64>165 return %0 : tensor<2xf64>166}167 168// CHECK-LABEL: func @scatter_empty_grid_axes169func.func @scatter_empty_grid_axes(170// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>171 %arg0 : tensor<4xf32>) -> tensor<4xf32> {172// CHECK-NOT: shard.scatter173 %0 = shard.scatter %arg0 on @grid0174 grid_axes = []175 scatter_axis = 0176 root = []177 : (tensor<4xf32>) -> tensor<4xf32>178// CHECK: return %[[ARG]]179 return %0 : tensor<4xf32>180}181 182// CHECK-LABEL: func @send_empty_grid_axes183func.func @send_empty_grid_axes(184// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>185 %arg0 : tensor<4xf32>) -> tensor<4xf32> {186// CHECK-NOT: shard.send187 %0 = shard.send %arg0 on @grid0188 grid_axes = []189 destination = []190 : (tensor<4xf32>) -> tensor<4xf32>191// CHECK: return %[[ARG]]192 return %0 : tensor<4xf32>193}194 195shard.grid @grid4x4(shape = 4x4)196// CHECK-LABEL: func @test_halo_sizes197func.func @test_halo_sizes() -> !shard.sharding {198 %c2_i64 = arith.constant 2 : i64199 // CHECK shard.sharding @grid4x4 split_axes = [[0], [1]] halo_sizes = [1, 2, 2, 22] : !shard.sharding200 %sharding = shard.sharding @grid4x4 split_axes = [[0], [1]] halo_sizes = [1, %c2_i64, %c2_i64, 22] : !shard.sharding201 return %sharding : !shard.sharding202}203 204// CHECK-LABEL: func @test_shard_offs205func.func @test_shard_offs() -> !shard.sharding {206 %c2_i64 = arith.constant 2 : i64207 // CHECK shard.sharding @grid4x4 split_axes = [[0], [1]] sharded_dims_offsets = [0, 1, 2, 3, 4, 0, 2, 3, 4, 22] : !shard.sharding208 %sharding = shard.sharding @grid4x4 split_axes = [[0], [1]] sharded_dims_offsets = [0, 1, %c2_i64, 3, 4, 0, %c2_i64, 3, 4, 22] : !shard.sharding209 return %sharding : !shard.sharding210}211 212// CHECK-LABEL: func @test_duplicate_shardops213func.func @test_duplicate_shardops() -> (tensor<1024x1024xf32>, tensor<1024x1024xf32>) attributes {llvm.emit_c_interface} {214 // CHECK-NEXT: [[vcst:%.*]] = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>215 %cst_1 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>216 // CHECK-NEXT: [[vsharding:%.*]] = shard.sharding @grid4x4 split_axes = {{\[\[}}0, 1]] : !shard.sharding217 %sharding_1 = shard.sharding @grid4x4 split_axes = [[0, 1]] : !shard.sharding218 %cst_2 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>219 %sharding_2 = shard.sharding @grid4x4 split_axes = [[0, 1]] : !shard.sharding220 %sharded_2 = shard.shard %cst_2 to %sharding_2 : tensor<1024x1024xf32>221 %cst_3 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>222 %sharding_3 = shard.sharding @grid4x4 split_axes = [[0, 1]] : !shard.sharding223 %sharded_3 = shard.shard %cst_3 to %sharding_3 : tensor<1024x1024xf32>224 // CHECK-NEXT: [[vsharded:%.*]] = shard.shard [[vcst]] to [[vsharding]] : tensor<1024x1024xf32>225 %sharded_1 = shard.shard %cst_1 to %sharding_1 : tensor<1024x1024xf32>226 // CHECK-NEXT: return [[vsharded]], [[vsharded]] : tensor<1024x1024xf32>, tensor<1024x1024xf32>227 return %sharded_1, %sharded_2 : tensor<1024x1024xf32>, tensor<1024x1024xf32>228}229 230// CHECK-LABEL: func @test_duplicate_shardops_diff231func.func @test_duplicate_shardops_diff() -> (tensor<1024x1024xf32>, tensor<1024x1024xf32>) attributes {llvm.emit_c_interface} {232 // CHECK-NEXT: [[vcst:%.*]] = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>233 %cst_1 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>234 // CHECK-NEXT: [[vsharding:%.*]] = shard.sharding @grid4x4 split_axes = {{\[\[}}0]] : !shard.sharding235 %sharding_1 = shard.sharding @grid4x4 split_axes = [[0]] : !shard.sharding236 %cst_2 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>237 // CHECK-NEXT: [[vsharding_0:%.*]] = shard.sharding @grid4x4 split_axes = {{\[\[}}0, 1]] : !shard.sharding238 %sharding_2 = shard.sharding @grid4x4 split_axes = [[0, 1]] : !shard.sharding239 // CHECK-NEXT: [[vsharded:%.*]] = shard.shard [[vcst]] to [[vsharding_0]] : tensor<1024x1024xf32>240 %sharded_2 = shard.shard %cst_2 to %sharding_2 : tensor<1024x1024xf32>241 %cst_3 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>242 %sharding_3 = shard.sharding @grid4x4 split_axes = [[0]] : !shard.sharding243 %sharded_3 = shard.shard %cst_3 to %sharding_3 : tensor<1024x1024xf32>244 // CHECK-NEXT: [[vsharded_1:%.*]] = shard.shard [[vsharded]] to [[vsharding]] : tensor<1024x1024xf32>245 %sharded_1 = shard.shard %cst_1 to %sharding_1 : tensor<1024x1024xf32>246 // CHECK-NEXT: return [[vsharded_1]], [[vsharded]] : tensor<1024x1024xf32>, tensor<1024x1024xf32>247 return %sharded_1, %sharded_2 : tensor<1024x1024xf32>, tensor<1024x1024xf32>248}249