50 lines · plain
1// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation{traversal=forward}))" %s | FileCheck %s2 3#map = affine_map<(d0, d1) -> (d0, d1)>4module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "mpich", "MPI:comm_world_rank" = 0 : i32>} {5 shard.grid @grid(shape = 1) {sym_visibility = "private"}6 func.func @test_forward() -> (tensor<6x6xi32>, tensor<6x6xi32>, tensor<i32>) attributes {llvm.emit_c_interface} {7 %c1_i32 = arith.constant 1 : i328 // CHECK: [[v3:%.*]] = tensor.empty() : tensor<6x6xi32>9 %0 = tensor.empty() : tensor<6x6xi32>10 // CHECK: [[v1:%.*]] = linalg.fill ins11 // CHECK: [[vsharding_0:%.*]] = shard.sharding @grid split_axes = {{\[\[}}0]] : !shard.sharding12 // CHECK: [[vsharded_1:%.*]] = shard.shard [[v1]] to [[vsharding_0]] : tensor<6x6xi32>13 %1 = linalg.fill ins(%c1_i32 : i32) outs(%0 : tensor<6x6xi32>) -> tensor<6x6xi32>14 %sharding = shard.sharding @grid split_axes = [[0]] : !shard.sharding15 %sharded = shard.shard %1 to %sharding : tensor<6x6xi32>16 // CHECK: [[v2:%.*]] = tensor.empty() : tensor<6x6xi32>17 // CHECK: [[vsharding_2:%.*]] = shard.sharding @grid split_axes = {{\[\[}}0]] : !shard.sharding18 // CHECK: [[vsharded_3:%.*]] = shard.shard [[vsharded_1]] to [[vsharding_2]] annotate_for_users : tensor<6x6xi32>19 %3 = tensor.empty() : tensor<6x6xi32>20 // CHECK: [[vsharding_4:%.*]] = shard.sharding @grid split_axes = {{\[\[}}0]] : !shard.sharding21 // CHECK: [[vsharded_5:%.*]] = shard.shard [[v2]] to [[vsharding_4]] annotate_for_users : tensor<6x6xi32>22 // CHECK: [[v3:%.*]] = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]}23 // CHECK-SAME: ins([[vsharded_3]], [[vsharded_3]] : tensor<6x6xi32>, tensor<6x6xi32>) outs([[vsharded_5]] : tensor<6x6xi32>) {24 // CHECK: [[vsharding_6:%.*]] = shard.sharding @grid split_axes = {{\[\[}}0]] : !shard.sharding25 // CHECK: [[vsharded_7:%.*]] = shard.shard [[v3]] to [[vsharding_6]] : tensor<6x6xi32>26 %4 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%sharded, %sharded27 : tensor<6x6xi32>, tensor<6x6xi32>) outs(%3 : tensor<6x6xi32>) {28 ^bb0(%in: i32, %in_2: i32, %out: i32):29 %9 = arith.addi %in, %in_2 : i3230 linalg.yield %9 : i3231 } -> tensor<6x6xi32>32 %c0_i32 = arith.constant 0 : i3233 %6 = tensor.empty() : tensor<i32>34 %7 = linalg.fill ins(%c0_i32 : i32) outs(%6 : tensor<i32>) -> tensor<i32>35 // CHECK: [[vreduced:%.*]] = linalg.reduce ins36 // CHECK: [[vsharding_12:%.*]] = shard.sharding @grid split_axes = [] : !shard.sharding37 // CHECK: [[vsharded_13:%.*]] = shard.shard [[vreduced]] to [[vsharding_12]] : tensor<i32>38 %reduced = linalg.reduce ins(%4 : tensor<6x6xi32>) outs(%7 : tensor<i32>) dimensions = [0, 1] 39 (%in: i32, %init: i32) {40 %9 = arith.addi %in, %init : i3241 linalg.yield %9 : i3242 }43 // CHECK: [[vsharding_14:%.*]] = shard.sharding @grid split_axes = {{\[\[}}]] : !shard.sharding44 %sharding_0 = shard.sharding @grid split_axes = [[]] : !shard.sharding45 // CHECK: [[vsharded_15:%.*]] = shard.shard [[vsharded_13]] to [[vsharding_14]] annotate_for_users : tensor<i32>46 %sharded_1 = shard.shard %reduced to %sharding_0 annotate_for_users : tensor<i32>47 return %sharded, %4, %sharded_1 : tensor<6x6xi32>, tensor<6x6xi32>, tensor<i32>48 }49}50