derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
249 lines
9.1 KiB
MLIR
249 lines
9.1 KiB
MLIR
// RUN: mlir-opt --canonicalize %s | FileCheck %s
|
|
|
|
shard.grid @grid0(shape = 2x4)
|
|
|
|
// CHECK-LABEL: func @all_reduce_empty_grid_axes
|
|
func.func @all_reduce_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.all_reduce
|
|
%0 = shard.all_reduce %arg0 on @grid0
|
|
grid_axes = []
|
|
: tensor<4xf32> -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @all_reduce_empty_grid_axes_different_return_type
|
|
func.func @all_reduce_empty_grid_axes_different_return_type(
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf64> {
|
|
// CHECK: shard.all_reduce
|
|
%0 = shard.all_reduce %arg0 on @grid0
|
|
// CHECK-NOT: grid_axes
|
|
grid_axes = []
|
|
: tensor<4xf32> -> tensor<4xf64>
|
|
return %0 : tensor<4xf64>
|
|
}
|
|
|
|
// CHECK-LABEL: func @all_reduce_default_reduction
|
|
func.func @all_reduce_default_reduction(
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf64> {
|
|
%0 = shard.all_reduce %arg0 on @grid0
|
|
grid_axes = [0]
|
|
// CHECK-NOT: reduction
|
|
reduction = sum
|
|
: tensor<4xf32> -> tensor<4xf64>
|
|
return %0 : tensor<4xf64>
|
|
}
|
|
|
|
// CHECK-LABEL: func @all_to_all_empty_grid_axes
|
|
func.func @all_to_all_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<8xf32>
|
|
%arg0 : tensor<8xf32>) -> tensor<8xf32> {
|
|
// CHECK-NOT: shard.all_to_all
|
|
%0 = shard.all_to_all %arg0 on @grid0
|
|
grid_axes = []
|
|
split_axis = 0
|
|
concat_axis = 0
|
|
: tensor<8xf32> -> tensor<8xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<8xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @all_gather_empty_grid_axes
|
|
func.func @all_gather_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.all_gather
|
|
%0 = shard.all_gather %arg0 on @grid0
|
|
grid_axes = []
|
|
gather_axis = 0
|
|
: tensor<4xf32> -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @all_slice_empty_grid_axes
|
|
func.func @all_slice_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.scatter
|
|
%0 = shard.all_slice %arg0 on @grid0
|
|
grid_axes = []
|
|
slice_axis = 0
|
|
: tensor<4xf32> -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @broadcast_empty_grid_axes
|
|
func.func @broadcast_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.broadcast
|
|
%0 = shard.broadcast %arg0 on @grid0
|
|
grid_axes = []
|
|
root = []
|
|
: (tensor<4xf32>) -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @gather_empty_grid_axes
|
|
func.func @gather_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.gather
|
|
%0 = shard.gather %arg0 on @grid0
|
|
grid_axes = []
|
|
gather_axis = 0
|
|
root = []
|
|
: (tensor<4xf32>) -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @receive_empty_grid_axes
|
|
func.func @receive_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.recv
|
|
%0 = shard.recv %arg0 on @grid0
|
|
grid_axes = []
|
|
: (tensor<4xf32>) -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @reduce_empty_grid_axes
|
|
func.func @reduce_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.reduce
|
|
%0 = shard.reduce %arg0 on @grid0
|
|
grid_axes = []
|
|
root = []
|
|
: (tensor<4xf32>) -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @reduce_scatter_empty_grid_axes
|
|
func.func @reduce_scatter_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.reduce_scatter
|
|
%0 = shard.reduce_scatter %arg0 on @grid0
|
|
grid_axes = []
|
|
scatter_dim = 0
|
|
: tensor<4xf32> -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @reduce_scatter_empty_grid_axes_different_return_type
|
|
func.func @reduce_scatter_empty_grid_axes_different_return_type(
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf64> {
|
|
// CHECK: shard.reduce_scatter
|
|
%0 = shard.reduce_scatter %arg0 on @grid0
|
|
// CHECK-NOT: grid_axes
|
|
grid_axes = []
|
|
scatter_dim = 0
|
|
: tensor<4xf32> -> tensor<4xf64>
|
|
return %0 : tensor<4xf64>
|
|
}
|
|
|
|
// CHECK-LABEL: func @reduce_scatter_default_reduction
|
|
func.func @reduce_scatter_default_reduction(
|
|
%arg0 : tensor<4xf32>) -> tensor<2xf64> {
|
|
%0 = shard.reduce_scatter %arg0 on @grid0
|
|
grid_axes = [0]
|
|
// CHECK-NOT: reduction
|
|
reduction = sum
|
|
scatter_dim = 0
|
|
: tensor<4xf32> -> tensor<2xf64>
|
|
return %0 : tensor<2xf64>
|
|
}
|
|
|
|
// CHECK-LABEL: func @scatter_empty_grid_axes
|
|
func.func @scatter_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.scatter
|
|
%0 = shard.scatter %arg0 on @grid0
|
|
grid_axes = []
|
|
scatter_dim = 0
|
|
root = []
|
|
: (tensor<4xf32>) -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @send_empty_grid_axes
|
|
func.func @send_empty_grid_axes(
|
|
// CHECK-SAME: %[[ARG:.*]]: tensor<4xf32>
|
|
%arg0 : tensor<4xf32>) -> tensor<4xf32> {
|
|
// CHECK-NOT: shard.send
|
|
%0 = shard.send %arg0 on @grid0
|
|
grid_axes = []
|
|
destination = []
|
|
: (tensor<4xf32>) -> tensor<4xf32>
|
|
// CHECK: return %[[ARG]]
|
|
return %0 : tensor<4xf32>
|
|
}
|
|
|
|
shard.grid @grid4x4(shape = 4x4)
|
|
// CHECK-LABEL: func @test_halo_sizes
|
|
func.func @test_halo_sizes() -> !shard.sharding {
|
|
%c2_i64 = arith.constant 2 : i64
|
|
// CHECK shard.sharding @grid4x4 split_axes = [[0], [1]] halo_sizes = [1, 2, 2, 22] : !shard.sharding
|
|
%sharding = shard.sharding @grid4x4 split_axes = [[0], [1]] halo_sizes = [1, %c2_i64, %c2_i64, 22] : !shard.sharding
|
|
return %sharding : !shard.sharding
|
|
}
|
|
|
|
// CHECK-LABEL: func @test_shard_offs
|
|
func.func @test_shard_offs() -> !shard.sharding {
|
|
%c2_i64 = arith.constant 2 : i64
|
|
// CHECK shard.sharding @grid4x4 split_axes = [[0], [1]] sharded_dims_offsets = [0, 1, 2, 3, 4, 0, 2, 3, 4, 22] : !shard.sharding
|
|
%sharding = shard.sharding @grid4x4 split_axes = [[0], [1]] sharded_dims_offsets = [0, 1, %c2_i64, 3, 4, 0, %c2_i64, 3, 4, 22] : !shard.sharding
|
|
return %sharding : !shard.sharding
|
|
}
|
|
|
|
// CHECK-LABEL: func @test_duplicate_shardops
|
|
func.func @test_duplicate_shardops() -> (tensor<1024x1024xf32>, tensor<1024x1024xf32>) attributes {llvm.emit_c_interface} {
|
|
// CHECK-NEXT: [[vcst:%.*]] = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>
|
|
%cst_1 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>
|
|
// CHECK-NEXT: [[vsharding:%.*]] = shard.sharding @grid4x4 split_axes = {{\[\[}}0, 1]] : !shard.sharding
|
|
%sharding_1 = shard.sharding @grid4x4 split_axes = [[0, 1]] : !shard.sharding
|
|
%cst_2 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>
|
|
%sharding_2 = shard.sharding @grid4x4 split_axes = [[0, 1]] : !shard.sharding
|
|
%sharded_2 = shard.shard %cst_2 to %sharding_2 : tensor<1024x1024xf32>
|
|
%cst_3 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>
|
|
%sharding_3 = shard.sharding @grid4x4 split_axes = [[0, 1]] : !shard.sharding
|
|
%sharded_3 = shard.shard %cst_3 to %sharding_3 : tensor<1024x1024xf32>
|
|
// CHECK-NEXT: [[vsharded:%.*]] = shard.shard [[vcst]] to [[vsharding]] : tensor<1024x1024xf32>
|
|
%sharded_1 = shard.shard %cst_1 to %sharding_1 : tensor<1024x1024xf32>
|
|
// CHECK-NEXT: return [[vsharded]], [[vsharded]] : tensor<1024x1024xf32>, tensor<1024x1024xf32>
|
|
return %sharded_1, %sharded_2 : tensor<1024x1024xf32>, tensor<1024x1024xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: func @test_duplicate_shardops_diff
|
|
func.func @test_duplicate_shardops_diff() -> (tensor<1024x1024xf32>, tensor<1024x1024xf32>) attributes {llvm.emit_c_interface} {
|
|
// CHECK-NEXT: [[vcst:%.*]] = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>
|
|
%cst_1 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>
|
|
// CHECK-NEXT: [[vsharding:%.*]] = shard.sharding @grid4x4 split_axes = {{\[\[}}0]] : !shard.sharding
|
|
%sharding_1 = shard.sharding @grid4x4 split_axes = [[0]] : !shard.sharding
|
|
%cst_2 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>
|
|
// CHECK-NEXT: [[vsharding_0:%.*]] = shard.sharding @grid4x4 split_axes = {{\[\[}}0, 1]] : !shard.sharding
|
|
%sharding_2 = shard.sharding @grid4x4 split_axes = [[0, 1]] : !shard.sharding
|
|
// CHECK-NEXT: [[vsharded:%.*]] = shard.shard [[vcst]] to [[vsharding_0]] : tensor<1024x1024xf32>
|
|
%sharded_2 = shard.shard %cst_2 to %sharding_2 : tensor<1024x1024xf32>
|
|
%cst_3 = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>
|
|
%sharding_3 = shard.sharding @grid4x4 split_axes = [[0]] : !shard.sharding
|
|
%sharded_3 = shard.shard %cst_3 to %sharding_3 : tensor<1024x1024xf32>
|
|
// CHECK-NEXT: [[vsharded_1:%.*]] = shard.shard [[vsharded]] to [[vsharding]] : tensor<1024x1024xf32>
|
|
%sharded_1 = shard.shard %cst_1 to %sharding_1 : tensor<1024x1024xf32>
|
|
// CHECK-NEXT: return [[vsharded_1]], [[vsharded]] : tensor<1024x1024xf32>, tensor<1024x1024xf32>
|
|
return %sharded_1, %sharded_2 : tensor<1024x1024xf32>, tensor<1024x1024xf32>
|
|
}
|