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
53 lines
2.3 KiB
MLIR
53 lines
2.3 KiB
MLIR
// RUN: mlir-opt \
|
|
// RUN: --pass-pipeline="builtin.module(func.func(shard-partition,test-single-fold))" \
|
|
// RUN: %s | FileCheck %s
|
|
|
|
shard.grid @grid_1d_4(shape = 4)
|
|
|
|
// CHECK-LABEL: func @tensor_empty_static_sharded_dims_offsets
|
|
func.func @tensor_empty_static_sharded_dims_offsets() -> () {
|
|
%b = tensor.empty() : tensor<8x16xf32>
|
|
%sharding = shard.sharding @grid_1d_4 split_axes = [[0]] sharded_dims_offsets = [0, 1, 4, 7, 8] : !shard.sharding
|
|
%sharded= shard.shard %b to %sharding : tensor<8x16xf32>
|
|
// CHECK: %[[sharding:.*]] = shard.sharding @grid_1d_4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 1, 4, 7, 8] : !shard.sharding
|
|
// CHECK: %[[proc_multi_idx:.*]] = shard.process_multi_index on @grid_1d_4 : index
|
|
// CHECK: %[[V0:.*]]:2 = shard.shard_shape dims = [8, 16] sharding = %[[sharding]] device = [%[[proc_multi_idx]]
|
|
// CHECK-SAME: ] : index, index
|
|
// CHECK: tensor.empty(%[[V0]]#0) : tensor<?x16xf32>
|
|
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: func @tensor_empty_dynamic_sharded_dims_offsets
|
|
// CHECK-SAME: %[[A0:.*]]: index
|
|
func.func @tensor_empty_dynamic_sharded_dims_offsets(%arg0 : index) -> () {
|
|
%b = tensor.empty(%arg0) : tensor<8x?xf32>
|
|
%sharding = shard.sharding @grid_1d_4 split_axes = [[0]] sharded_dims_offsets = [0, 1, 4, 7, 8] : !shard.sharding
|
|
%sharded= shard.shard %b to %sharding : tensor<8x?xf32>
|
|
// CHECK: %[[sharding:.*]] = shard.sharding @grid_1d_4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 1, 4, 7, 8] : !shard.sharding
|
|
// CHECK: %[[proc_multi_idx:.*]] = shard.process_multi_index on @grid_1d_4 : index
|
|
// CHECK: %[[V0:.*]]:2 = shard.shard_shape dims = [8, %[[A0]]
|
|
// CHECK-SAME: ] sharding = %[[sharding]] device = [%[[proc_multi_idx]]
|
|
// CHECK-SAME: ] : index, index
|
|
// CHECK: tensor.empty(%[[V0]]#0, %[[A0]]) : tensor<?x?xf32>
|
|
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: func @tensor_empty_same_static_dims_sizes
|
|
func.func @tensor_empty_same_static_dims_sizes() -> () {
|
|
%b = tensor.empty() : tensor<16x16xf32>
|
|
%sharding = shard.sharding @grid_1d_4 split_axes = [[0]] sharded_dims_offsets = [0, 4, 8, 12, 16] : !shard.sharding
|
|
%sharded= shard.shard %b to %sharding : tensor<16x16xf32>
|
|
// CHECK-NEXT: tensor.empty() : tensor<4x16xf32>
|
|
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: func @tensor_empty_0d
|
|
func.func @tensor_empty_0d() -> () {
|
|
tensor.empty() : tensor<f32>
|
|
// CHECK-NEXT: tensor.empty() : tensor<f32>
|
|
return
|
|
}
|