318 lines · plain
1// RUN: mlir-opt \2// RUN: --pass-pipeline="builtin.module(func.func(shard-partition,test-single-fold))" \3// RUN: %s | FileCheck %s4 5shard.grid @grid_1d(shape = 2)6 7// CHECK-LABEL: func @return_sharding8func.func @return_sharding(9 // CHECK-SAME: [[ARG:%.*]]: tensor<1xf32>10 %arg0: tensor<2xf32>11// CHECK-SAME: ) -> (tensor<1xf32>, !shard.sharding) {12) -> (tensor<2xf32>, !shard.sharding) {13 %ssharded = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding14 %sharded = shard.shard %arg0 to %ssharded : tensor<2xf32>15 // CHECK-NEXT: [[vsharding:%.*]] = shard.sharding @grid_1d split_axes = {{\[\[}}0]] : !shard.sharding16 %r = shard.get_sharding %sharded : tensor<2xf32> -> !shard.sharding17 // CHECK-NEXT: return [[ARG]], [[vsharding]] : tensor<1xf32>, !shard.sharding18 return %sharded, %r : tensor<2xf32>, !shard.sharding19}20 21// CHECK-LABEL: func @full_replication22func.func @full_replication(23 // CHECK-SAME: %[[ARG:.*]]: tensor<2xi8>24 %arg0: tensor<2xi8>25// CHECK-SAME: -> tensor<2xi8> {26) -> tensor<2xi8> {27 %s0 = shard.sharding @grid_1d split_axes = [[]] : !shard.sharding28 %0 = shard.shard %arg0 to %s0 : tensor<2xi8>29 %s1 = shard.sharding @grid_1d split_axes = [[]] : !shard.sharding30 %1 = shard.shard %0 to %s1 annotate_for_users : tensor<2xi8>31 // CHECK: return %[[ARG]] : tensor<2xi8>32 return %1 : tensor<2xi8>33}34 35// CHECK-LABEL: func @sharding_triplet36func.func @sharding_triplet(37 // CHECK-SAME: %[[ARG:.*]]: tensor<1xf32>38 %arg0: tensor<2xf32>39// CHECK-SAME: ) -> tensor<2xf32> {40) -> tensor<2xf32> {41 // CHECK: %[[ALL_GATHER:.*]] = shard.all_gather %[[ARG]] on @grid_1d grid_axes = [0] gather_axis = 0 : tensor<1xf32> -> tensor<2xf32>42 %ssharded = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding43 %sharded = shard.shard %arg0 to %ssharded : tensor<2xf32>44 %ssharded_0 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding45 %sharded_0 = shard.shard %sharded to %ssharded_0 annotate_for_users : tensor<2xf32>46 %ssharded_1 = shard.sharding @grid_1d split_axes = [[]] : !shard.sharding47 %sharded_1 = shard.shard %sharded_0 to %ssharded_1 : tensor<2xf32>48 // CHECK: return %[[ALL_GATHER]] : tensor<2xf32>49 return %sharded_1 : tensor<2xf32>50}51 52 53// CHECK-LABEL: func @move_split_axis54func.func @move_split_axis(55 // CHECK-SAME: %[[ARG:.*]]: tensor<1x2xi8>56 %arg0: tensor<2x2xi8>57// CHECK-SAME: -> tensor<2x1xi8> {58) -> tensor<2x2xi8> {59 // CHECK: %[[ALL_TO_ALL:.*]] = shard.all_to_all %[[ARG]] on @grid_1d60 // CHECK-SAME: grid_axes = [0] split_axis = 1 concat_axis = 0 : tensor<1x2xi8> -> tensor<2x1xi8>61 %s0 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding62 %0 = shard.shard %arg0 to %s0 : tensor<2x2xi8>63 %s1 = shard.sharding @grid_1d split_axes = [[], [0]] : !shard.sharding64 %1 = shard.shard %0 to %s1 annotate_for_users : tensor<2x2xi8>65 // CHECK: return %[[ALL_TO_ALL]] : tensor<2x1xi8>66 return %1 : tensor<2x2xi8>67}68 69// CHECK-LABEL: func @non_tensor_value70func.func @non_tensor_value(71 // CHECK-SAME: %[[ARG:.*]]: i872 %arg0: i873// CHECK-SAME: -> i8 {74) -> i8 {75 // CHECK: %[[RES:.*]] = arith.addi %[[ARG]], %[[ARG]] : i876 %0 = arith.addi %arg0, %arg0 : i877 // CHECK: return %[[RES]] : i878 return %0 : i879}80 81// CHECK-LABEL: func @unary_elementwise82func.func @unary_elementwise(83 // CHECK-SAME: %[[ARG:.*]]: tensor<1xi8>84 %arg0: tensor<2xi8>85// CHECK-SAME: -> tensor<1xi8> {86) -> tensor<2xi8> {87 %s0 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding88 %0 = shard.shard %arg0 to %s0 : tensor<2xi8>89 %s1 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding90 %1 = shard.shard %0 to %s1 annotate_for_users : tensor<2xi8>91 // CHECK: %[[RES:.*]] = tosa.abs %[[ARG]] : (tensor<1xi8>) -> tensor<1xi8>92 %2 = tosa.abs %1 : (tensor<2xi8>) -> tensor<2xi8>93 %s3 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding94 %3 = shard.shard %2 to %s3 : tensor<2xi8>95 %s4 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding96 %4 = shard.shard %3 to %s4 annotate_for_users : tensor<2xi8>97 // CHECK: return %[[RES]] : tensor<1xi8>98 return %4 : tensor<2xi8>99}100 101// full replication -> shard axis -> abs -> shard axis -> full replication102// CHECK-LABEL: func @unary_elementwise_with_resharding103func.func @unary_elementwise_with_resharding(104 // CHECK-SAME: %[[ARG:.*]]: tensor<2xi8>105 %arg0: tensor<2xi8>106// CHECK-SAME: -> tensor<2xi8> {107) -> tensor<2xi8> {108 // CHECK: %[[SLICE:.*]] = shard.all_slice %[[ARG]] on @grid_1d grid_axes = [0] slice_axis = 0109 // CHECK-SAME: tensor<2xi8> -> tensor<1xi8>110 %s0 = shard.sharding @grid_1d split_axes = [[]] : !shard.sharding111 %0 = shard.shard %arg0 to %s0 : tensor<2xi8>112 %s1 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding113 %1 = shard.shard %0 to %s1 annotate_for_users : tensor<2xi8>114 // CHECK: %[[ABS:.*]] = tosa.abs %[[SLICE]] : (tensor<1xi8>) -> tensor<1xi8>115 %2 = tosa.abs %1 : (tensor<2xi8>) -> tensor<2xi8>116 // CHECK: %[[RES:.*]] = shard.all_gather %[[ABS]] on @grid_1d117 // CHECK-SAME: grid_axes = [0] gather_axis = 0 : tensor<1xi8> -> tensor<2xi8>118 %s3 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding119 %3 = shard.shard %2 to %s3 : tensor<2xi8>120 %s4 = shard.sharding @grid_1d split_axes = [[]] : !shard.sharding121 %4 = shard.shard %3 to %s4 annotate_for_users : tensor<2xi8>122 // CHECK: return %[[RES]] : tensor<2xi8>123 return %4 : tensor<2xi8>124}125 126// CHECK-LABEL: func @binary_elementwise127func.func @binary_elementwise(128 // CHECK-SAME: %[[ARG0:.*]]: tensor<1xi8>,129 %arg0: tensor<2xi8>,130 // CHECK-SAME: %[[ARG1:.*]]: tensor<1xi8>131 %arg1: tensor<2xi8>132// CHECK-SAME: -> tensor<1xi8> {133) -> tensor<2xi8> {134 %sarg0_sharded = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding135 %arg0_sharded = shard.shard %arg0 to %sarg0_sharded : tensor<2xi8>136 %sop_arg0 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding137 %op_arg0 = shard.shard %arg0_sharded to %sop_arg0 annotate_for_users : tensor<2xi8>138 %sarg1_sharded = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding139 %arg1_sharded = shard.shard %arg1 to %sarg1_sharded : tensor<2xi8>140 %sop_arg1 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding141 %op_arg1 = shard.shard %arg1_sharded to %sop_arg1 annotate_for_users : tensor<2xi8>142 // CHECK: %[[RES:.*]] = tosa.add %[[ARG0]], %[[ARG1]] : (tensor<1xi8>, tensor<1xi8>) -> tensor<1xi8>143 %op_res = tosa.add %op_arg0, %op_arg1 : (tensor<2xi8>, tensor<2xi8>) -> tensor<2xi8>144 %sop_res_sharded = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding145 %op_res_sharded = shard.shard %op_res to %sop_res_sharded : tensor<2xi8>146 %sres = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding147 %res = shard.shard %op_res_sharded to %sres annotate_for_users : tensor<2xi8>148 // CHECK: return %[[RES]] : tensor<1xi8>149 return %res : tensor<2xi8>150}151 152// reshard153// abs154// reshard155// abs156// reshard157// CHECK-LABEL: func @multiple_chained_ops158func.func @multiple_chained_ops(159 // CHECK-SAME: %[[ARG:.*]]: tensor<2xi8>160 %arg0: tensor<2xi8>161// CHECK-SAME: -> tensor<1xi8> {162) -> tensor<2xi8> {163 // CHECK: %[[RESHARD1:.*]] = shard.all_slice %[[ARG]] on @grid_1d grid_axes = [0] slice_axis = 0164 // CHECK-SAME: tensor<2xi8> -> tensor<1xi8>165 %s0 = shard.sharding @grid_1d split_axes = [[]] : !shard.sharding166 %0 = shard.shard %arg0 to %s0 : tensor<2xi8>167 %s1 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding168 %1 = shard.shard %0 to %s1 annotate_for_users : tensor<2xi8>169 // CHECK: %[[ABS1:.*]] = tosa.abs %[[RESHARD1]] : (tensor<1xi8>) -> tensor<1xi8>170 %2 = tosa.abs %1 : (tensor<2xi8>) -> tensor<2xi8>171 // CHECK: %[[RESHARD2:.*]] = shard.all_gather %[[ABS1]] on @grid_1d172 // CHECK-SAME: grid_axes = [0] gather_axis = 0 : tensor<1xi8> -> tensor<2xi8>173 %s3 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding174 %3 = shard.shard %2 to %s3 : tensor<2xi8>175 %s4 = shard.sharding @grid_1d split_axes = [[]] : !shard.sharding176 %4 = shard.shard %3 to %s4 annotate_for_users : tensor<2xi8>177 // CHECK: %[[ABS2:.*]] = tosa.abs %[[RESHARD2]] : (tensor<2xi8>) -> tensor<2xi8>178 %5 = tosa.abs %4 : (tensor<2xi8>) -> tensor<2xi8>179 // CHECK: %[[RESHARD3:.*]] = shard.all_slice %[[ABS2]] on @grid_1d grid_axes = [0] slice_axis = 0 :180 // CHECK-SAME: tensor<2xi8> -> tensor<1xi8>181 %s6 = shard.sharding @grid_1d split_axes = [[]] : !shard.sharding182 %6 = shard.shard %5 to %s6 : tensor<2xi8>183 %s7 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding184 %7 = shard.shard %6 to %s7 annotate_for_users : tensor<2xi8>185 // CHECK: return %[[RESHARD3]] : tensor<1xi8>186 return %7 : tensor<2xi8>187}188 189// CHECK-LABEL: func @incomplete_sharding190func.func @incomplete_sharding(191 // CHECK-SAME: %[[ARG:.*]]: tensor<4x16xf32>192 %arg0: tensor<8x16xf32>193// CHECK-SAME: -> tensor<4x16xf32> {194) -> tensor<8x16xf32> {195 %s0 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding196 %0 = shard.shard %arg0 to %s0 annotate_for_users : tensor<8x16xf32>197 // CHECK: %[[RES:.*]] = tosa.sigmoid %[[ARG]] : (tensor<4x16xf32>) -> tensor<4x16xf32>198 %1 = tosa.sigmoid %0 : (tensor<8x16xf32>) -> tensor<8x16xf32>199 %s2 = shard.sharding @grid_1d split_axes = [[0]] : !shard.sharding200 %2 = shard.shard %1 to %s2 : tensor<8x16xf32>201 // CHECK: return %[[RES]] : tensor<4x16xf32>202 return %2 : tensor<8x16xf32>203}204 205shard.grid @grid_1d_4(shape = 4)206 207// CHECK-LABEL: func @ew_chain_with_halo208func.func @ew_chain_with_halo(209 // CHECK-SAME: %[[IN1:[A-Za-z0-9_]+]]: tensor<5x16xf32>210 %arg0: tensor<8x16xf32>,211 // CHECK-SAME: %[[IN2:[A-Za-z0-9_]+]]: tensor<1xf32>212 %arg1: tensor<1xf32>,213 // CHECK-SAME: %[[IN3:[A-Za-z0-9_]+]]: tensor<1xf32>214 %arg2: tensor<1xf32>)215 // CHECK-SAME: -> tensor<5x16xf32>216 -> tensor<8x16xf32> {217 %ssharded = shard.sharding @grid_1d_4 split_axes = [[0]] halo_sizes = [2, 1] : !shard.sharding218 %sharded = shard.shard %arg0 to %ssharded annotate_for_users : tensor<8x16xf32>219 // CHECK: %[[TMP1:.*]] = tosa.tanh %[[IN1]] : (tensor<5x16xf32>) -> tensor<5x16xf32>220 %0 = tosa.tanh %sharded : (tensor<8x16xf32>) -> tensor<8x16xf32>221 %ssharded_0 = shard.sharding @grid_1d_4 split_axes = [[0]] halo_sizes = [2, 1] : !shard.sharding222 %sharded_0 = shard.shard %0 to %ssharded_0 : tensor<8x16xf32>223 %ssharded_1 = shard.sharding @grid_1d_4 split_axes = [[0]] halo_sizes = [2, 1] : !shard.sharding224 %sharded_1 = shard.shard %sharded_0 to %ssharded_1 annotate_for_users : tensor<8x16xf32>225 // CHECK-NEXT: %[[TMP2:.*]] = tosa.abs %[[TMP1]] : (tensor<5x16xf32>) -> tensor<5x16xf32>226 %1 = tosa.abs %sharded_1 : (tensor<8x16xf32>) -> tensor<8x16xf32>227 %ssharded_2 = shard.sharding @grid_1d_4 split_axes = [[0]] halo_sizes = [2, 1] : !shard.sharding228 %sharded_2 = shard.shard %1 to %ssharded_2 : tensor<8x16xf32>229 %ssharded_4 = shard.sharding @grid_1d_4 split_axes = [[0]] halo_sizes = [2, 1] : !shard.sharding230 %sharded_4 = shard.shard %sharded_2 to %ssharded_4 annotate_for_users : tensor<8x16xf32>231 // CHECK-NEXT: %[[TMP3:.*]] = tosa.negate %[[TMP2]], %[[IN2]], %[[IN3]] : (tensor<5x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5x16xf32>232 %sharding_1 = shard.sharding @grid_1d_4 split_axes = [[]] : !shard.sharding233 %zero_point_1 = shard.shard %arg1 to %sharding_1 annotate_for_users : tensor<1xf32>234 %zero_point_2 = shard.shard %arg2 to %sharding_1 annotate_for_users : tensor<1xf32>235 %2 = tosa.negate %sharded_4, %zero_point_1, %zero_point_2 : (tensor<8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<8x16xf32>236 %ssharded_5 = shard.sharding @grid_1d_4 split_axes = [[0]] halo_sizes = [2, 1] : !shard.sharding237 %sharded_5 = shard.shard %2 to %ssharded_5 : tensor<8x16xf32>238 %ssharded_6 = shard.sharding @grid_1d_4 split_axes = [[0]] halo_sizes = [2, 1] : !shard.sharding239 %sharded_6 = shard.shard %sharded_5 to %ssharded_6 annotate_for_users : tensor<8x16xf32>240 // CHECK-NEXT: return %[[TMP3]] : tensor<5x16xf32>241 return %sharded_6 : tensor<8x16xf32>242}243 244// CHECK-LABEL: func @test_shard_update_halo245// CHECK-SAME: %[[IN1:[A-Za-z0-9_]+]]: tensor<300x1200xi64>246func.func @test_shard_update_halo(%arg0: tensor<1200x1200xi64>) -> tensor<1200x1200xi64> {247 %sharding = shard.sharding @grid_1d_4 split_axes = [[0]] : !shard.sharding248 // CHECK: %[[T:.*]] = tensor.empty() : tensor<304x1200xi64>249 // CHECK: %[[inserted_slice:.*]] = tensor.insert_slice %[[IN1]] into %[[T]][2, 0] [300, 1200] [1, 1] : tensor<300x1200xi64> into tensor<304x1200xi64>250 // CHECK: %[[UH:.*]] = shard.update_halo %[[inserted_slice]] on @grid_1d_4 split_axes = {{\[\[0]]}} halo_sizes = [2, 2] : tensor<304x1200xi64>251 %sharded = shard.shard %arg0 to %sharding : tensor<1200x1200xi64>252 %sharding_0 = shard.sharding @grid_1d_4 split_axes = [[0]] halo_sizes = [2, 2] : !shard.sharding253 %sharded_1 = shard.shard %sharded to %sharding_0 : tensor<1200x1200xi64>254 %sharded_3 = shard.shard %sharded_1 to %sharding_0 annotate_for_users : tensor<1200x1200xi64>255 // CHECK: return %[[UH]] : tensor<304x1200xi64>256 return %sharded_3 : tensor<1200x1200xi64>257}258 259shard.grid @grid4x4(shape = 4x4)260// CHECK-LABEL: func @test_shard_update_halo2d261// CHECK-SAME: %[[IN1:[A-Za-z0-9_]+]]: tensor<300x300xi64>262func.func @test_shard_update_halo2d(%arg0: tensor<1200x1200xi64>) -> tensor<1200x1200xi64> {263 %sharding = shard.sharding @grid4x4 split_axes = [[0], [1]] : !shard.sharding264 // CHECK: %[[T:.*]] = tensor.empty() : tensor<303x307xi64>265 // CHECK: %[[inserted_slice:.*]] = tensor.insert_slice %[[IN1]] into %[[T]][1, 3] [300, 300] [1, 1] : tensor<300x300xi64> into tensor<303x307xi64>266 // CHECK: %[[UH:.*]] = shard.update_halo %[[inserted_slice]] on @grid4x4 split_axes = {{\[\[}}0], [1]] halo_sizes = [1, 2, 3, 4] : tensor<303x307xi64>267 %sharded = shard.shard %arg0 to %sharding : tensor<1200x1200xi64>268 %sharding_0 = shard.sharding @grid4x4 split_axes = [[0], [1]] halo_sizes = [1, 2, 3, 4] : !shard.sharding269 %sharded_1 = shard.shard %sharded to %sharding_0 : tensor<1200x1200xi64>270 %sharded_3 = shard.shard %sharded_1 to %sharding_0 annotate_for_users : tensor<1200x1200xi64>271 // CHECK: return %[[UH]] : tensor<303x307xi64>272 return %sharded_3 : tensor<1200x1200xi64>273}274 275shard.grid @grid(shape = 2)276// CHECK-LABEL: func.func @test_reduce_0d(277// CHECK-SAME: %[[ARG0:[A-Za-z0-9_]+]]: tensor<3x6xi32>278func.func @test_reduce_0d(%arg0: tensor<6x6xi32>) -> (tensor<i32>) {279 %sharding = shard.sharding @grid split_axes = [[0]] : !shard.sharding280 %sharded = shard.shard %arg0 to %sharding annotate_for_users : tensor<6x6xi32>281 %4 = tensor.empty() : tensor<i32>282 %sharding_out = shard.sharding @grid split_axes = [[]] : !shard.sharding283 %sharded_out = shard.shard %4 to %sharding_out : tensor<i32>284 %sharded_in = shard.shard %sharded to %sharding annotate_for_users : tensor<6x6xi32>285 // CHECK: %[[reduced:.*]] = linalg.reduce ins(%arg0 : tensor<3x6xi32>)286 %reduced = linalg.reduce ins(%sharded_in : tensor<6x6xi32>) outs(%sharded_out : tensor<i32>) dimensions = [0, 1] 287 (%in: i32, %init: i32) {288 %6 = arith.addi %in, %init : i32289 linalg.yield %6 : i32290 }291 // CHECK: %[[all_reduce:.*]] = shard.all_reduce %[[reduced]] on @grid grid_axes = [0] : tensor<i32> -> tensor<i32>292 %sharded_red = shard.shard %reduced to %sharding_out : tensor<i32>293 %sharded_ret = shard.shard %sharded_red to %sharding_out annotate_for_users : tensor<i32>294 // CHECK: return %[[all_reduce]] : tensor<i32>295 return %sharded_ret : tensor<i32>296}297 298// CHECK-LABEL: func.func @test_reduce_1d(299// CHECK-SAME: %[[ARG0:[A-Za-z0-9_]+]]: tensor<3x6xi32>300func.func @test_reduce_1d(%arg0: tensor<6x6xi32>) -> (tensor<6xi32>) {301 %sharding = shard.sharding @grid split_axes = [[0]] : !shard.sharding302 %sharded = shard.shard %arg0 to %sharding annotate_for_users : tensor<6x6xi32>303 %4 = tensor.empty() : tensor<6xi32>304 %sharded_out = shard.shard %4 to %sharding : tensor<6xi32>305 %sharded_in = shard.shard %sharded to %sharding annotate_for_users : tensor<6x6xi32>306 // CHECK: %[[reduced:.*]] = linalg.reduce ins(%arg0 : tensor<3x6xi32>)307 %reduced = linalg.reduce ins(%sharded_in : tensor<6x6xi32>) outs(%sharded_out : tensor<6xi32>) dimensions = [1] 308 (%in: i32, %init: i32) {309 %6 = arith.addi %in, %init : i32310 linalg.yield %6 : i32311 }312 // CHECK-NOT: shard.all_reduce313 %sharded_red = shard.shard %reduced to %sharding : tensor<6xi32>314 %sharded_ret = shard.shard %sharded_red to %sharding annotate_for_users : tensor<6xi32>315 // CHECK: return %[[reduced]] : tensor<3xi32>316 return %sharded_ret : tensor<6xi32>317}318