27 lines · plain
1// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation{traversal=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 %sharding = shard.sharding @grid split_axes = [[0]] : !shard.sharding11 // CHECK-COUNT-2: shard.shard12 %sharded = shard.shard %0 to %sharding : tensor<6x6xi32>13 %1 = linalg.fill ins(%c1_i32 : i32) outs(%sharded : tensor<6x6xi32>) -> tensor<6x6xi32>14 // CHECK: tensor.empty()15 // CHECK-NOT: shard.shard @16 %2 = tensor.empty() : tensor<6x6xi32>17 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%1, %118 : tensor<6x6xi32>, tensor<6x6xi32>) outs(%2 : tensor<6x6xi32>) {19 ^bb0(%in: i32, %in_2: i32, %out: i32):20 %9 = arith.addi %in, %in_2 : i3221 linalg.yield %9 : i3222 } -> tensor<6x6xi32>23 // CHECK: return24 return %3 : tensor<6x6xi32>25 }26}27