53 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_4(shape = 4)6 7// CHECK-LABEL: func @tensor_empty_static_sharded_dims_offsets8func.func @tensor_empty_static_sharded_dims_offsets() -> () {9 %b = tensor.empty() : tensor<8x16xf32>10 %sharding = shard.sharding @grid_1d_4 split_axes = [[0]] sharded_dims_offsets = [0, 1, 4, 7, 8] : !shard.sharding11 %sharded= shard.shard %b to %sharding : tensor<8x16xf32>12 // CHECK: %[[sharding:.*]] = shard.sharding @grid_1d_4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 1, 4, 7, 8] : !shard.sharding13 // CHECK: %[[proc_multi_idx:.*]] = shard.process_multi_index on @grid_1d_4 : index14 // CHECK: %[[V0:.*]]:2 = shard.shard_shape dims = [8, 16] sharding = %[[sharding]] device = [%[[proc_multi_idx]]15 // CHECK-SAME: ] : index, index16 // CHECK: tensor.empty(%[[V0]]#0) : tensor<?x16xf32>17 18 return19}20 21// CHECK-LABEL: func @tensor_empty_dynamic_sharded_dims_offsets22// CHECK-SAME: %[[A0:.*]]: index23func.func @tensor_empty_dynamic_sharded_dims_offsets(%arg0 : index) -> () {24 %b = tensor.empty(%arg0) : tensor<8x?xf32>25 %sharding = shard.sharding @grid_1d_4 split_axes = [[0]] sharded_dims_offsets = [0, 1, 4, 7, 8] : !shard.sharding26 %sharded= shard.shard %b to %sharding : tensor<8x?xf32>27 // CHECK: %[[sharding:.*]] = shard.sharding @grid_1d_4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 1, 4, 7, 8] : !shard.sharding28 // CHECK: %[[proc_multi_idx:.*]] = shard.process_multi_index on @grid_1d_4 : index29 // CHECK: %[[V0:.*]]:2 = shard.shard_shape dims = [8, %[[A0]]30 // CHECK-SAME: ] sharding = %[[sharding]] device = [%[[proc_multi_idx]]31 // CHECK-SAME: ] : index, index32 // CHECK: tensor.empty(%[[V0]]#0, %[[A0]]) : tensor<?x?xf32>33 34 return35}36 37// CHECK-LABEL: func @tensor_empty_same_static_dims_sizes38func.func @tensor_empty_same_static_dims_sizes() -> () {39 %b = tensor.empty() : tensor<16x16xf32>40 %sharding = shard.sharding @grid_1d_4 split_axes = [[0]] sharded_dims_offsets = [0, 4, 8, 12, 16] : !shard.sharding41 %sharded= shard.shard %b to %sharding : tensor<16x16xf32>42 // CHECK-NEXT: tensor.empty() : tensor<4x16xf32>43 44 return45}46 47// CHECK-LABEL: func @tensor_empty_0d48func.func @tensor_empty_0d() -> () {49 tensor.empty() : tensor<f32>50 // CHECK-NEXT: tensor.empty() : tensor<f32>51 return52}53