28 lines · plain
1// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation{traversal=forward-backward}))" %s | FileCheck %s2 3#map = affine_map<(d0, d1) -> (d0, d1)>4module {5 shard.grid @grid(shape = 1) {sym_visibility = "private"}6 func.func @test_forward() -> tensor<6x6xi32> {7 %c1_i32 = arith.constant 1 : i328 // CHECK: tensor.empty()9 %0 = tensor.empty() : tensor<6x6xi32>10 // CHECK-COUNT-3: shard.sharding @grid split_axes = {{\[\[0}}]]11 %sharding_row = shard.sharding @grid split_axes = [[0]] : !shard.sharding12 %annotated_row = shard.shard %0 to %sharding_row : tensor<6x6xi32>13 %1 = linalg.fill ins(%c1_i32 : i32) outs(%annotated_row : tensor<6x6xi32>) -> tensor<6x6xi32>14 %2 = tensor.empty() : tensor<6x6xi32>15 // CHECK-COUNT-4: shard.sharding @grid split_axes = {{\[\[1}}]]16 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%2, %117 : tensor<6x6xi32>, tensor<6x6xi32>) outs(%2 : tensor<6x6xi32>) {18 ^bb0(%in: i32, %in_2: i32, %out: i32):19 %9 = arith.addi %in, %in_2 : i3220 linalg.yield %9 : i3221 } -> tensor<6x6xi32>22 %sharding_col = shard.sharding @grid split_axes = [[1]] : !shard.sharding23 %annotated_col = shard.shard %3 to %sharding_col : tensor<6x6xi32>24 // CHECK: return25 return %annotated_col : tensor<6x6xi32>26 }27}28