367 lines · plain
1// RUN: mlir-opt %s -convert-shard-to-mpi -canonicalize -split-input-file | FileCheck %s2 3// -----4// CHECK: shard.grid @grid05shard.grid @grid0(shape = 3x4x5)6func.func @process_multi_index() -> (index, index, index) {7 // CHECK: mpi.comm_rank8 // CHECK-DAG: %[[v4:.*]] = arith.remsi9 // CHECK-DAG: %[[v0:.*]] = arith.remsi10 // CHECK-DAG: %[[v1:.*]] = arith.remsi11 %0:3 = shard.process_multi_index on @grid0 axes = [] : index, index, index12 // CHECK: return %[[v1]], %[[v0]], %[[v4]] : index, index, index13 return %0#0, %0#1, %0#2 : index, index, index14}15 16// CHECK-LABEL: func @process_linear_index17func.func @process_linear_index() -> index {18 // CHECK: %[[RES:.*]], %[[rank:.*]] = mpi.comm_rank19 // CHECK: %[[cast:.*]] = arith.index_cast %[[rank]] : i32 to index20 %0 = shard.process_linear_index on @grid0 : index21 // CHECK: return %[[cast]] : index22 return %0 : index23}24 25// CHECK-LABEL: func @neighbors_dim026func.func @neighbors_dim0(%arg0 : tensor<120x120x120xi8>) -> (index, index) {27 %c0 = arith.constant 0 : index28 %c1 = arith.constant 1 : index29 %c4 = arith.constant 4 : index30 // CHECK-DAG: [[up:%.*]] = arith.constant 44 : index31 // CHECK-DAG: [[down:%.*]] = arith.constant 4 : index32 %idx:2 = shard.neighbors_linear_indices on @grid0[%c1, %c0, %c4] split_axes = [0] : index, index33 // CHECK: return [[down]], [[up]] : index, index34 return %idx#0, %idx#1 : index, index35}36 37// CHECK-LABEL: func @neighbors_dim138func.func @neighbors_dim1(%arg0 : tensor<120x120x120xi8>) -> (index, index) {39 %c0 = arith.constant 0 : index40 %c1 = arith.constant 1 : index41 %c4 = arith.constant 4 : index42 // CHECK-DAG: [[up:%.*]] = arith.constant 29 : index43 // CHECK-DAG: [[down:%.*]] = arith.constant -1 : index44 %idx:2 = shard.neighbors_linear_indices on @grid0[%c1, %c0, %c4] split_axes = [1] : index, index45 // CHECK: return [[down]], [[up]] : index, index46 return %idx#0, %idx#1 : index, index47}48 49// CHECK-LABEL: func @neighbors_dim250func.func @neighbors_dim2(%arg0 : tensor<120x120x120xi8>) -> (index, index) {51 %c0 = arith.constant 0 : index52 %c1 = arith.constant 1 : index53 %c4 = arith.constant 4 : index54 // CHECK-DAG: [[up:%.*]] = arith.constant -1 : index55 // CHECK-DAG: [[down:%.*]] = arith.constant 23 : index56 %idx:2 = shard.neighbors_linear_indices on @grid0[%c1, %c0, %c4] split_axes = [2] : index, index57 // CHECK: return [[down]], [[up]] : index, index58 return %idx#0, %idx#1 : index, index59}60 61// -----62// CHECK: shard.grid @grid063module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 24> } {64 shard.grid @grid0(shape = 3x4x5)65 func.func @process_multi_index() -> (index, index, index) {66 // CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index67 // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index68 // CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index69 %0:3 = shard.process_multi_index on @grid0 axes = [] : index, index, index70 // CHECK: return %[[c1]], %[[c0]], %[[c4]] : index, index, index71 return %0#0, %0#1, %0#2 : index, index, index72 }73 74 // CHECK-LABEL: func @process_linear_index75 func.func @process_linear_index() -> index {76 // CHECK: %[[c24:.*]] = arith.constant 24 : index77 %0 = shard.process_linear_index on @grid0 : index78 // CHECK: return %[[c24]] : index79 return %0 : index80 }81}82 83// -----84module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 7> } {85 shard.grid @grid0(shape = 3x4x5)86 // CHECK-LABEL: func.func @allreduce_tensor(87 func.func @allreduce_tensor(88 // CHECK-SAME: [[varg0:%.*]]: tensor<3x4xf32>89 %arg0 : tensor<3x4xf32>) -> tensor<3x4xf32> {90 // CHECK-DAG: [[vc4_i32:%.*]] = arith.constant 4 : i3291 // CHECK-DAG: [[vc2_i32:%.*]] = arith.constant 2 : i3292 // CHECK: [[v0:%.*]] = bufferization.to_buffer [[varg0]] : tensor<3x4xf32> to memref<3x4xf32>93 // CHECK: [[valloc:%.*]] = memref.alloc() : memref<3x4xf32>94 // CHECK: linalg.copy ins([[v0]] : memref<3x4xf32>) outs([[valloc]] : memref<3x4xf32>)95 // CHECK: [[v1:%.*]] = mpi.comm_world : !mpi.comm96 // CHECK: [[vnewcomm:%.*]] = mpi.comm_split([[v1]], [[vc2_i32]], [[vc4_i32]]) : !mpi.comm97 // CHECK: [[vcollapse_shape:%.*]] = memref.collapse_shape [[valloc]] {{\[\[}}0, 1]] : memref<3x4xf32> into memref<12xf32>98 // CHECK: mpi.allreduce([[vcollapse_shape]], [[vcollapse_shape]], MPI_MAX, [[vnewcomm]]) : memref<12xf32>, memref<12xf32>99 // CHECK: [[v2:%.*]] = bufferization.to_tensor [[valloc]] restrict : memref<3x4xf32> to tensor<3x4xf32>100 %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0, 1] reduction = max : tensor<3x4xf32> -> tensor<3x4xf32>101 // CHECK: return [[v2]] : tensor<3x4xf32>102 return %0 : tensor<3x4xf32>103 }104 105 // CHECK-LABEL: func.func @allreduce_memref(106 func.func @allreduce_memref(107 // CHECK-SAME: [[varg0:%.*]]: memref<3x4xf32>108 %arg0 : memref<3x4xf32>) -> memref<3x4xf32> {109 // CHECK: [[vc4_i32:%.*]] = arith.constant 4 : i32110 // CHECK: [[vc2_i32:%.*]] = arith.constant 2 : i32111 // CHECK: [[valloc:%.*]] = memref.alloc() : memref<3x4xf32>112 // CHECK: linalg.copy ins([[varg0]] : memref<3x4xf32>) outs([[valloc]] : memref<3x4xf32>)113 // CHECK: [[v0:%.*]] = mpi.comm_world : !mpi.comm114 // CHECK: [[vnewcomm:%.*]] = mpi.comm_split([[v0]], [[vc2_i32]], [[vc4_i32]]) : !mpi.comm115 // CHECK: [[vcollapse_shape:%.*]] = memref.collapse_shape [[valloc]] {{\[\[}}0, 1]] : memref<3x4xf32> into memref<12xf32>116 // CHECK: mpi.allreduce([[vcollapse_shape]], [[vcollapse_shape]], MPI_MAX, [[vnewcomm]]) : memref<12xf32>, memref<12xf32>117 %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0, 1] reduction = max : memref<3x4xf32> -> memref<3x4xf32>118 // CHECK: return [[valloc]] : memref<3x4xf32>119 return %0 : memref<3x4xf32>120 }121 122 // CHECK-LABEL: func.func @allreduce_new_type(123 func.func @allreduce_new_type(124 // CHECK-SAME: [[varg0:%.*]]: memref<3x4xf32>125 %arg0 : memref<3x4xf32>) -> memref<3x4xf64> {126 // CHECK: [[vc4_i32:%.*]] = arith.constant 4 : i32127 // CHECK: [[vc2_i32:%.*]] = arith.constant 2 : i32128 // CHECK: [[valloc:%.*]] = memref.alloc() : memref<3x4xf64>129 // CHECK: linalg.copy ins([[varg0]] : memref<3x4xf32>) outs([[valloc]] : memref<3x4xf64>)130 // CHECK: [[v0:%.*]] = mpi.comm_world : !mpi.comm131 // CHECK: [[vnewcomm:%.*]] = mpi.comm_split([[v0]], [[vc2_i32]], [[vc4_i32]]) : !mpi.comm132 // CHECK: [[vcollapse_shape:%.*]] = memref.collapse_shape [[valloc]] {{\[\[}}0, 1]] : memref<3x4xf64> into memref<12xf64>133 // CHECK: mpi.allreduce([[vcollapse_shape]], [[vcollapse_shape]], MPI_MAX, [[vnewcomm]]) : memref<12xf64>, memref<12xf64>134 %0 = shard.all_reduce %arg0 on @grid0 grid_axes = [0, 1] reduction = max : memref<3x4xf32> -> memref<3x4xf64>135 // CHECK: return [[valloc]] : memref<3x4xf64>136 return %0 : memref<3x4xf64>137 }138}139 140// -----141shard.grid @grid0(shape = 3x4x5)142// CHECK-LABEL: func @update_halo_1d_first143func.func @update_halo_1d_first(144 // CHECK-SAME: [[arg0:%.*]]: memref<120x120x120xi8>145 %arg0 : memref<120x120x120xi8>) -> memref<120x120x120xi8> {146 // CHECK: memref.subview [[arg0]][115, 0, 0] [2, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<2x120x120xi8147 // CHECK: mpi.send(148 // CHECK-SAME: : memref<2x120x120xi8>, i32, i32149 // CHECK: mpi.recv(150 // CHECK-SAME: : memref<2x120x120xi8>, i32, i32151 // CHECK: memref.subview [[arg0]][0, 0, 0] [2, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<2x120x120xi8152 // CHECK: memref.subview [[arg0]][2, 0, 0] [3, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<3x120x120xi8153 // CHECK: mpi.send(154 // CHECK-SAME: : memref<3x120x120xi8>, i32, i32155 // CHECK: mpi.recv(156 // CHECK-SAME: : memref<3x120x120xi8>, i32, i32157 // CHECK: memref.subview [[arg0]][117, 0, 0] [3, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<3x120x120xi8158 %res = shard.update_halo %arg0 on @grid0 split_axes = [[0]] halo_sizes = [2, 3] : memref<120x120x120xi8>159 // CHECK: return [[res:%.*]] : memref<120x120x120xi8>160 return %res : memref<120x120x120xi8>161}162 163// -----164module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 1> } {165 shard.grid @grid0(shape = 4)166 // CHECK-LABEL: func @update_halo_1d_with_zero167 func.func @update_halo_1d_with_zero (168 // CHECK-SAME: [[varg0:%.*]]: memref<120x120x120xi8>169 %arg0 : memref<120x120x120xi8>) -> memref<120x120x120xi8> {170 // CHECK-DAG: [[vc91_i32:%.*]] = arith.constant 91 : i32171 // CHECK-DAG: [[vc0_i32:%.*]] = arith.constant 0 : i32172 // CHECK-DAG: [[vc2_i32:%.*]] = arith.constant 2 : i32173 // CHECK: [[v0:%.*]] = mpi.comm_world : !mpi.comm174 // CHECK: [[valloc:%.*]] = memref.alloc() : memref<2x120x120xi8>175 // CHECK: [[vsubview:%.*]] = memref.subview [[varg0]][118, 0, 0] [2, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<2x120x120xi8, strided<[14400, 120, 1], offset: 1699200>>176 // CHECK: memref.copy [[vsubview]], [[valloc]] : memref<2x120x120xi8, strided<[14400, 120, 1], offset: 1699200>> to memref<2x120x120xi8>177 // CHECK: mpi.send([[valloc]], [[vc91_i32]], [[vc2_i32]], [[v0]]) : memref<2x120x120xi8>, i32, i32178 // CHECK: mpi.recv([[valloc]], [[vc91_i32]], [[vc0_i32]], [[v0]]) : memref<2x120x120xi8>, i32, i32179 // CHECK: [[vsubview_0:%.*]] = memref.subview [[varg0]][0, 0, 0] [2, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<2x120x120xi8, strided<[14400, 120, 1]>>180 // CHECK: memref.copy [[valloc]], [[vsubview_0]] : memref<2x120x120xi8> to memref<2x120x120xi8, strided<[14400, 120, 1]>>181 // CHECK: memref.dealloc [[valloc]] : memref<2x120x120xi8>182 %res = shard.update_halo %arg0 on @grid0 split_axes = [[0]] halo_sizes = [2, 0] : memref<120x120x120xi8>183 // CHECK: return [[varg0]] : memref<120x120x120xi8>184 return %res : memref<120x120x120xi8>185 }186}187 188// -----189module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 24> } {190 shard.grid @grid0(shape = 3x4x5)191 // CHECK-LABEL: func @update_halo_3d192 func.func @update_halo_3d(193 // CHECK-SAME: [[varg0:%.*]]: memref<120x120x120xi8>194 %arg0 : memref<120x120x120xi8>) -> memref<120x120x120xi8> {195 // CHECK-DAG: [[vc23_i32:%.*]] = arith.constant 23 : i32196 // CHECK-DAG: [[vc29_i32:%.*]] = arith.constant 29 : i32197 // CHECK-DAG: [[vc91_i32:%.*]] = arith.constant 91 : i32198 // CHECK-DAG: [[vc4_i32:%.*]] = arith.constant 4 : i32199 // CHECK-DAG: [[vc44_i32:%.*]] = arith.constant 44 : i32200 // CHECK: [[v0:%.*]] = mpi.comm_world : !mpi.comm201 // CHECK: [[valloc:%.*]] = memref.alloc() : memref<117x113x5xi8>202 // CHECK: [[vsubview:%.*]] = memref.subview [[varg0]][1, 3, 109] [117, 113, 5] [1, 1, 1] : memref<120x120x120xi8> to memref<117x113x5xi8, strided<[14400, 120, 1], offset: 14869>>203 // CHECK: memref.copy [[vsubview]], [[valloc]] : memref<117x113x5xi8, strided<[14400, 120, 1], offset: 14869>> to memref<117x113x5xi8>204 // CHECK: mpi.send([[valloc]], [[vc91_i32]], [[vc44_i32]], [[v0]]) : memref<117x113x5xi8>, i32, i32205 // CHECK: mpi.recv([[valloc]], [[vc91_i32]], [[vc4_i32]], [[v0]]) : memref<117x113x5xi8>, i32, i32206 // CHECK: [[vsubview_0:%.*]] = memref.subview [[varg0]][1, 3, 0] [117, 113, 5] [1, 1, 1] : memref<120x120x120xi8> to memref<117x113x5xi8, strided<[14400, 120, 1], offset: 14760>>207 // CHECK: memref.copy [[valloc]], [[vsubview_0]] : memref<117x113x5xi8> to memref<117x113x5xi8, strided<[14400, 120, 1], offset: 14760>>208 // CHECK: memref.dealloc [[valloc]] : memref<117x113x5xi8>209 // CHECK: [[valloc_1:%.*]] = memref.alloc() : memref<117x113x6xi8>210 // CHECK: [[vsubview_2:%.*]] = memref.subview [[varg0]][1, 3, 5] [117, 113, 6] [1, 1, 1] : memref<120x120x120xi8> to memref<117x113x6xi8, strided<[14400, 120, 1], offset: 14765>>211 // CHECK: memref.copy [[vsubview_2]], [[valloc_1]] : memref<117x113x6xi8, strided<[14400, 120, 1], offset: 14765>> to memref<117x113x6xi8>212 // CHECK: mpi.send([[valloc_1]], [[vc91_i32]], [[vc4_i32]], [[v0]]) : memref<117x113x6xi8>, i32, i32213 // CHECK: mpi.recv([[valloc_1]], [[vc91_i32]], [[vc44_i32]], [[v0]]) : memref<117x113x6xi8>, i32, i32214 // CHECK: [[vsubview_3:%.*]] = memref.subview [[varg0]][1, 3, 114] [117, 113, 6] [1, 1, 1] : memref<120x120x120xi8> to memref<117x113x6xi8, strided<[14400, 120, 1], offset: 14874>>215 // CHECK: memref.copy [[valloc_1]], [[vsubview_3]] : memref<117x113x6xi8> to memref<117x113x6xi8, strided<[14400, 120, 1], offset: 14874>>216 // CHECK: memref.dealloc [[valloc_1]] : memref<117x113x6xi8>217 // CHECK: [[v1:%.*]] = mpi.comm_world : !mpi.comm218 // CHECK: [[valloc_4:%.*]] = memref.alloc() : memref<117x3x120xi8>219 // CHECK: [[vsubview_5:%.*]] = memref.subview [[varg0]][1, 113, 0] [117, 3, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<117x3x120xi8, strided<[14400, 120, 1], offset: 27960>>220 // CHECK: memref.copy [[vsubview_5]], [[valloc_4]] : memref<117x3x120xi8, strided<[14400, 120, 1], offset: 27960>> to memref<117x3x120xi8>221 // CHECK: mpi.send([[valloc_4]], [[vc91_i32]], [[vc29_i32]], [[v1]]) : memref<117x3x120xi8>, i32, i32222 // CHECK: memref.dealloc [[valloc_4]] : memref<117x3x120xi8>223 // CHECK: [[valloc_6:%.*]] = memref.alloc() : memref<117x4x120xi8>224 // CHECK: mpi.recv([[valloc_6]], [[vc91_i32]], [[vc29_i32]], [[v1]]) : memref<117x4x120xi8>, i32, i32225 // CHECK: [[vsubview_7:%.*]] = memref.subview [[varg0]][1, 116, 0] [117, 4, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<117x4x120xi8, strided<[14400, 120, 1], offset: 28320>>226 // CHECK: memref.copy [[valloc_6]], [[vsubview_7]] : memref<117x4x120xi8> to memref<117x4x120xi8, strided<[14400, 120, 1], offset: 28320>>227 // CHECK: memref.dealloc [[valloc_6]] : memref<117x4x120xi8>228 // CHECK: [[v2:%.*]] = mpi.comm_world : !mpi.comm229 // CHECK: [[valloc_8:%.*]] = memref.alloc() : memref<1x120x120xi8>230 // CHECK: mpi.recv([[valloc_8]], [[vc91_i32]], [[vc23_i32]], [[v2]]) : memref<1x120x120xi8>, i32, i32231 // CHECK: [[vsubview_9:%.*]] = memref.subview [[varg0]][0, 0, 0] [1, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<1x120x120xi8, strided<[14400, 120, 1]>>232 // CHECK: memref.copy [[valloc_8]], [[vsubview_9]] : memref<1x120x120xi8> to memref<1x120x120xi8, strided<[14400, 120, 1]>>233 // CHECK: memref.dealloc [[valloc_8]] : memref<1x120x120xi8>234 // CHECK: [[valloc_10:%.*]] = memref.alloc() : memref<2x120x120xi8>235 // CHECK: [[vsubview_11:%.*]] = memref.subview [[varg0]][1, 0, 0] [2, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<2x120x120xi8, strided<[14400, 120, 1], offset: 14400>>236 // CHECK: memref.copy [[vsubview_11]], [[valloc_10]] : memref<2x120x120xi8, strided<[14400, 120, 1], offset: 14400>> to memref<2x120x120xi8>237 // CHECK: mpi.send([[valloc_10]], [[vc91_i32]], [[vc23_i32]], [[v2]]) : memref<2x120x120xi8>, i32, i32238 // CHECK: memref.dealloc [[valloc_10]] : memref<2x120x120xi8>239 %res = shard.update_halo %arg0 on @grid0 split_axes = [[2], [1], [0]] halo_sizes = [1, 2, 3, 4, 5, 6] : memref<120x120x120xi8>240 // CHECK: return [[varg0]] : memref<120x120x120xi8>241 return %res : memref<120x120x120xi8>242 }243 244 // CHECK-LABEL: func @update_halo_3d_tensor245 func.func @update_halo_3d_tensor(246 // CHECK-SAME: [[varg0:%.*]]: tensor<120x120x120xi8>247 %arg0 : tensor<120x120x120xi8>) -> tensor<120x120x120xi8> {248 // CHECK-DAG: [[vc23_i32:%.*]] = arith.constant 23 : i32249 // CHECK-DAG: [[vc29_i32:%.*]] = arith.constant 29 : i32250 // CHECK-DAG: [[vc44_i32:%.*]] = arith.constant 44 : i32251 // CHECK-DAG: [[vc4_i32:%.*]] = arith.constant 4 : i32252 // CHECK-DAG: [[vc91_i32:%.*]] = arith.constant 91 : i32253 // CHECK: [[v0:%.*]] = bufferization.to_buffer [[varg0]] : tensor<120x120x120xi8> to memref<120x120x120xi8>254 // CHECK: [[v1:%.*]] = mpi.comm_world : !mpi.comm255 // CHECK: [[valloc:%.*]] = memref.alloc() : memref<117x113x5xi8>256 // CHECK: [[vsubview:%.*]] = memref.subview [[v0]][1, 3, 109] [117, 113, 5] [1, 1, 1] : memref<120x120x120xi8> to memref<117x113x5xi8, strided<[14400, 120, 1], offset: 14869>>257 // CHECK: memref.copy [[vsubview]], [[valloc]] : memref<117x113x5xi8, strided<[14400, 120, 1], offset: 14869>> to memref<117x113x5xi8>258 // CHECK: mpi.send([[valloc]], [[vc91_i32]], [[vc44_i32]], [[v1]]) : memref<117x113x5xi8>, i32, i32259 // CHECK: mpi.recv([[valloc]], [[vc91_i32]], [[vc4_i32]], [[v1]]) : memref<117x113x5xi8>, i32, i32260 // CHECK: [[vsubview_0:%.*]] = memref.subview [[v0]][1, 3, 0] [117, 113, 5] [1, 1, 1] : memref<120x120x120xi8> to memref<117x113x5xi8, strided<[14400, 120, 1], offset: 14760>>261 // CHECK: memref.copy [[valloc]], [[vsubview_0]] : memref<117x113x5xi8> to memref<117x113x5xi8, strided<[14400, 120, 1], offset: 14760>>262 // CHECK: memref.dealloc [[valloc]] : memref<117x113x5xi8>263 // CHECK: [[valloc_1:%.*]] = memref.alloc() : memref<117x113x6xi8>264 // CHECK: [[vsubview_2:%.*]] = memref.subview [[v0]][1, 3, 5] [117, 113, 6] [1, 1, 1] : memref<120x120x120xi8> to memref<117x113x6xi8, strided<[14400, 120, 1], offset: 14765>>265 // CHECK: memref.copy [[vsubview_2]], [[valloc_1]] : memref<117x113x6xi8, strided<[14400, 120, 1], offset: 14765>> to memref<117x113x6xi8>266 // CHECK: mpi.send([[valloc_1]], [[vc91_i32]], [[vc4_i32]], [[v1]]) : memref<117x113x6xi8>, i32, i32267 // CHECK: mpi.recv([[valloc_1]], [[vc91_i32]], [[vc44_i32]], [[v1]]) : memref<117x113x6xi8>, i32, i32268 // CHECK: [[vsubview_3:%.*]] = memref.subview [[v0]][1, 3, 114] [117, 113, 6] [1, 1, 1] : memref<120x120x120xi8> to memref<117x113x6xi8, strided<[14400, 120, 1], offset: 14874>>269 // CHECK: memref.copy [[valloc_1]], [[vsubview_3]] : memref<117x113x6xi8> to memref<117x113x6xi8, strided<[14400, 120, 1], offset: 14874>>270 // CHECK: memref.dealloc [[valloc_1]] : memref<117x113x6xi8>271 // CHECK: [[v2:%.*]] = mpi.comm_world : !mpi.comm272 // CHECK: [[valloc_4:%.*]] = memref.alloc() : memref<117x3x120xi8>273 // CHECK: [[vsubview_5:%.*]] = memref.subview [[v0]][1, 113, 0] [117, 3, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<117x3x120xi8, strided<[14400, 120, 1], offset: 27960>>274 // CHECK: memref.copy [[vsubview_5]], [[valloc_4]] : memref<117x3x120xi8, strided<[14400, 120, 1], offset: 27960>> to memref<117x3x120xi8>275 // CHECK: mpi.send([[valloc_4]], [[vc91_i32]], [[vc29_i32]], [[v2]]) : memref<117x3x120xi8>, i32, i32276 // CHECK: memref.dealloc [[valloc_4]] : memref<117x3x120xi8>277 // CHECK: [[valloc_6:%.*]] = memref.alloc() : memref<117x4x120xi8>278 // CHECK: mpi.recv([[valloc_6]], [[vc91_i32]], [[vc29_i32]], [[v2]]) : memref<117x4x120xi8>, i32, i32279 // CHECK: [[vsubview_7:%.*]] = memref.subview [[v0]][1, 116, 0] [117, 4, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<117x4x120xi8, strided<[14400, 120, 1], offset: 28320>>280 // CHECK: memref.copy [[valloc_6]], [[vsubview_7]] : memref<117x4x120xi8> to memref<117x4x120xi8, strided<[14400, 120, 1], offset: 28320>>281 // CHECK: memref.dealloc [[valloc_6]] : memref<117x4x120xi8>282 // CHECK: [[v3:%.*]] = mpi.comm_world : !mpi.comm283 // CHECK: [[valloc_8:%.*]] = memref.alloc() : memref<1x120x120xi8>284 // CHECK: mpi.recv([[valloc_8]], [[vc91_i32]], [[vc23_i32]], [[v3]]) : memref<1x120x120xi8>, i32, i32285 // CHECK: [[vsubview_9:%.*]] = memref.subview [[v0]][0, 0, 0] [1, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<1x120x120xi8, strided<[14400, 120, 1]>>286 // CHECK: memref.copy [[valloc_8]], [[vsubview_9]] : memref<1x120x120xi8> to memref<1x120x120xi8, strided<[14400, 120, 1]>>287 // CHECK: memref.dealloc [[valloc_8]] : memref<1x120x120xi8>288 // CHECK: [[valloc_10:%.*]] = memref.alloc() : memref<2x120x120xi8>289 // CHECK: [[vsubview_11:%.*]] = memref.subview [[v0]][1, 0, 0] [2, 120, 120] [1, 1, 1] : memref<120x120x120xi8> to memref<2x120x120xi8, strided<[14400, 120, 1], offset: 14400>>290 // CHECK: memref.copy [[vsubview_11]], [[valloc_10]] : memref<2x120x120xi8, strided<[14400, 120, 1], offset: 14400>> to memref<2x120x120xi8>291 // CHECK: mpi.send([[valloc_10]], [[vc91_i32]], [[vc23_i32]], [[v3]]) : memref<2x120x120xi8>, i32, i32292 // CHECK: memref.dealloc [[valloc_10]] : memref<2x120x120xi8>293 // CHECK: [[v4:%.*]] = bufferization.to_tensor [[v0]] restrict writable : memref<120x120x120xi8> to tensor<120x120x120xi8>294 %res = shard.update_halo %arg0 on @grid0 split_axes = [[2], [1], [0]] halo_sizes = [1, 2, 3, 4, 5, 6] : tensor<120x120x120xi8>295 // CHECK: return [[v4]] : tensor<120x120x120xi8>296 return %res : tensor<120x120x120xi8>297 }298}299 300// -----301shard.grid @grid0(shape = 2x2x4)302// CHECK-LABEL: func.func @return_sharding(303// CHECK-SAME: [[varg0:%.*]]: tensor<2x4xf32>) -> (tensor<2x4xf32>, tensor<?x?xi16>, tensor<?x?xi64>, tensor<?x?xi64>) {304func.func @return_sharding(%arg0: tensor<2x4xf32>) -> (tensor<2x4xf32>, !shard.sharding) {305 %sharding = shard.sharding @grid0 split_axes = [[0, 1], [2]] : !shard.sharding306 // CHECK: [[vcst:%.*]] = arith.constant dense<2> : tensor<1xi16>307 // CHECK: [[vcst_0:%.*]] = arith.constant dense<[0, 1]> : tensor<2xi16>308 // CHECK: [[vcm1_i16:%.*]] = arith.constant -1 : i16309 // CHECK: [[v0:%.*]] = tensor.empty() : tensor<2x2xi16>310 // CHECK: [[v1:%.*]] = linalg.fill ins([[vcm1_i16]] : i16) outs([[v0]] : tensor<2x2xi16>) -> tensor<2x2xi16>311 // CHECK: [[vinserted_slice:%.*]] = tensor.insert_slice [[vcst_0]] into [[v1]][0, 0] [1, 2] [1, 1] : tensor<2xi16> into tensor<2x2xi16>312 // CHECK: [[vinserted_slice_1:%.*]] = tensor.insert_slice [[vcst]] into [[vinserted_slice]][1, 0] [1, 1] [1, 1] : tensor<1xi16> into tensor<2x2xi16>313 // CHECK: [[v2:%.*]] = tensor.empty() : tensor<0x0xi64>314 // CHECK: [[v3:%.*]] = tensor.empty() : tensor<0x0xi64>315 // CHECK: [[vcast:%.*]] = tensor.cast [[vinserted_slice_1]] : tensor<2x2xi16> to tensor<?x?xi16>316 // CHECK: [[vcast_2:%.*]] = tensor.cast [[v2]] : tensor<0x0xi64> to tensor<?x?xi64>317 // CHECK: [[vcast_3:%.*]] = tensor.cast [[v3]] : tensor<0x0xi64> to tensor<?x?xi64>318 // CHECK: return [[varg0]], [[vcast]], [[vcast_2]], [[vcast_3]] : tensor<2x4xf32>, tensor<?x?xi16>, tensor<?x?xi64>, tensor<?x?xi64>319 return %arg0, %sharding : tensor<2x4xf32>, !shard.sharding320}321 322// CHECK-LABEL: func.func @return_sharding_halos(323// CHECK-SAME: [[varg0:%.*]]: tensor<6x8xf32>) -> (tensor<6x8xf32>, tensor<?x?xi16>, tensor<?x?xi64>, tensor<?x?xi64>) {324func.func @return_sharding_halos(%arg0: tensor<6x8xf32>) -> (tensor<6x8xf32>, !shard.sharding) {325 %sharding = shard.sharding @grid0 split_axes = [[0, 1], [2]] halo_sizes = [0, 4, 3, 1] : !shard.sharding326 // CHECK: [[vcst:%.*]] = arith.constant dense<{{\[\[}}0, 4], [3, 1]]> : tensor<2x2xi64>327 // CHECK: [[vcst_0:%.*]] = arith.constant dense<2> : tensor<1xi16>328 // CHECK: [[vcst_1:%.*]] = arith.constant dense<[0, 1]> : tensor<2xi16>329 // CHECK: [[vcm1_i16:%.*]] = arith.constant -1 : i16330 // CHECK: [[v0:%.*]] = tensor.empty() : tensor<2x2xi16>331 // CHECK: [[v1:%.*]] = linalg.fill ins([[vcm1_i16]] : i16) outs([[v0]] : tensor<2x2xi16>) -> tensor<2x2xi16>332 // CHECK: [[vinserted_slice:%.*]] = tensor.insert_slice [[vcst_1]] into [[v1]][0, 0] [1, 2] [1, 1] : tensor<2xi16> into tensor<2x2xi16>333 // CHECK: [[vinserted_slice_2:%.*]] = tensor.insert_slice [[vcst_0]] into [[vinserted_slice]][1, 0] [1, 1] [1, 1] : tensor<1xi16> into tensor<2x2xi16>334 // CHECK: [[v2:%.*]] = tensor.empty() : tensor<0x0xi64>335 // CHECK: [[vcast:%.*]] = tensor.cast [[vinserted_slice_2]] : tensor<2x2xi16> to tensor<?x?xi16>336 // CHECK: [[vcast_3:%.*]] = tensor.cast [[vcst]] : tensor<2x2xi64> to tensor<?x?xi64>337 // CHECK: [[vcast_4:%.*]] = tensor.cast [[v2]] : tensor<0x0xi64> to tensor<?x?xi64>338 // CHECK: return [[varg0]], [[vcast]], [[vcast_3]], [[vcast_4]] : tensor<6x8xf32>, tensor<?x?xi16>, tensor<?x?xi64>, tensor<?x?xi64>339 return %arg0, %sharding : tensor<6x8xf32>, !shard.sharding340}341 342// CHECK-LABEL: func.func @return_sharding_offs(343// CHECK-SAME: [[varg0:%.*]]: tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xi16>, tensor<?x?xi64>, tensor<?x?xi64>) {344func.func @return_sharding_offs(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>, !shard.sharding) {345 %sharding = shard.sharding @grid0 split_axes = [[0, 1], [2]] sharded_dims_offsets = [0, 3, 5, 7, 8, 0, 0, 5, 10, 16] : !shard.sharding346 // CHECK: [[vcst:%.*]] = arith.constant dense<[0, 0, 5, 10, 16]> : tensor<5xi64>347 // CHECK: [[vcst_0:%.*]] = arith.constant dense<[0, 3, 5, 7, 8]> : tensor<5xi64>348 // CHECK: [[vcm9223372036854775808_i64:%.*]] = arith.constant -9223372036854775808 : i64349 // CHECK: [[vcst_1:%.*]] = arith.constant dense<2> : tensor<1xi16>350 // CHECK: [[vcst_2:%.*]] = arith.constant dense<[0, 1]> : tensor<2xi16>351 // CHECK: [[vcm1_i16:%.*]] = arith.constant -1 : i16352 // CHECK: [[v0:%.*]] = tensor.empty() : tensor<2x2xi16>353 // CHECK: [[v1:%.*]] = linalg.fill ins([[vcm1_i16]] : i16) outs([[v0]] : tensor<2x2xi16>) -> tensor<2x2xi16>354 // CHECK: [[vinserted_slice:%.*]] = tensor.insert_slice [[vcst_2]] into [[v1]][0, 0] [1, 2] [1, 1] : tensor<2xi16> into tensor<2x2xi16>355 // CHECK: [[vinserted_slice_3:%.*]] = tensor.insert_slice [[vcst_1]] into [[vinserted_slice]][1, 0] [1, 1] [1, 1] : tensor<1xi16> into tensor<2x2xi16>356 // CHECK: [[v2:%.*]] = tensor.empty() : tensor<0x0xi64>357 // CHECK: [[v3:%.*]] = tensor.empty() : tensor<2x5xi64>358 // CHECK: [[v4:%.*]] = linalg.fill ins([[vcm9223372036854775808_i64]] : i64) outs([[v3]] : tensor<2x5xi64>) -> tensor<2x5xi64>359 // CHECK: [[vinserted_slice_4:%.*]] = tensor.insert_slice [[vcst_0]] into [[v4]][0, 0] [1, 5] [1, 1] : tensor<5xi64> into tensor<2x5xi64>360 // CHECK: [[vinserted_slice_5:%.*]] = tensor.insert_slice [[vcst]] into [[vinserted_slice_4]][1, 0] [1, 5] [1, 1] : tensor<5xi64> into tensor<2x5xi64>361 // CHECK: [[vcast:%.*]] = tensor.cast [[vinserted_slice_3]] : tensor<2x2xi16> to tensor<?x?xi16>362 // CHECK: [[vcast_6:%.*]] = tensor.cast [[v2]] : tensor<0x0xi64> to tensor<?x?xi64>363 // CHECK: [[vcast_7:%.*]] = tensor.cast [[vinserted_slice_5]] : tensor<2x5xi64> to tensor<?x?xi64>364 // CHECK: return [[varg0]], [[vcast]], [[vcast_6]], [[vcast_7]] : tensor<?x?xf32>, tensor<?x?xi16>, tensor<?x?xi64>, tensor<?x?xi64>365 return %arg0, %sharding : tensor<?x?xf32>, !shard.sharding366}367