From c2dfe37bd1a960d211349fe51c5891ac38fe15d9 Mon Sep 17 00:00:00 2001 From: amalbasaTT Date: Tue, 3 Dec 2024 13:46:51 +0000 Subject: [PATCH] #11512: Include both gen_sharded_spec_unary functions --- .../sweep_utils/sharding_utils.py | 115 +++++++++++++++++- .../eltwise/unary/isinf/isinf_sharded.py | 4 +- .../eltwise/unary/isnan/isnan_sharded.py | 4 +- .../unary/isneginf/isneginf_sharded.py | 4 +- .../unary/isposinf/isposinf_sharded.py | 4 +- .../eltwise/unary/lgamma/lgamma_sharded.py | 4 +- .../sweeps/eltwise/unary/mish/mish_sharded.py | 4 +- 7 files changed, 123 insertions(+), 16 deletions(-) diff --git a/tests/sweep_framework/sweep_utils/sharding_utils.py b/tests/sweep_framework/sweep_utils/sharding_utils.py index 499a33d150c..aca108bf2bc 100644 --- a/tests/sweep_framework/sweep_utils/sharding_utils.py +++ b/tests/sweep_framework/sweep_utils/sharding_utils.py @@ -29,10 +29,117 @@ def roundup(a, b): return result +def gen_sharded_spec_unary(num_shapes, max_tensor_size=4 * 1024 * 1024, layouts=["TILE_LAYOUT", "ROW_MAJOR_LAYOUT"]): + # device.compute_with_storage_grid_size() + y = 8 + x = 8 + + # ["BLOCK", "WIDTH", "HEIGHT", "tensor_wh"] + sharding_strategy_list = ["BLOCK", "WIDTH", "HEIGHT", "TENSOR_WH"] + shard_orientation_list = ["COL_MAJOR", "ROW_MAJOR"] + spec_list = [] + + for sharding_strategy, shard_orientation, rank, layout in itertools.product( + sharding_strategy_list, shard_orientation_list, [4, 3, 2], layouts + ): + for _ in range(num_shapes): + if sharding_strategy == "TENSOR_WH": + # Gets stuck: + # X 8 Y 8 input_shape [1, 17792, 8] DataType.BFLOAT8_B Layout.TILE ShardStrategy.BLOCK ShardOrientation.COL_MAJOR tensor_hw_as_shard_shape True + + if layout == "TILE_LAYOUT": + # In shard mode ShardMode::PHYSICAL, physical shard shape {12, 13312} is not compatible with alignment Alignment([32, 32])! + min_shard_size_x = 32 + min_shard_size_y = 32 + else: # if layout == "ROW_MAJOR_LAYOUT": + # Shard Size must be multiple of input_tile_size (width * height is multiple of 1024) + min_shard_size_x = random.choice([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]) + min_shard_size_y = 1024 // min_shard_size_x + + rest_volume = random.randint(1, max_tensor_size // (min_shard_size_x * min_shard_size_y * x * y)) + input_shape = random.choice(_gen_reshape_args_from_volume(rest_volume, step=1, out_dims=rank)) + input_shape = list(input_shape["reshape_dims"]) + input_shape[-2] = input_shape[-2] * min_shard_size_x + input_shape[-1] = input_shape[-1] * min_shard_size_y + + # Shard width should be multiple of 16 to satisfy L1 alignment (width = multiple 8 for bfloat16) + while input_shape[-1] % 16 != 0: + input_shape[-1] *= 2 + input_shape[-2] //= 2 + + if shard_orientation == "COL_MAJOR": + tmp = input_shape[-2] + input_shape[-2] = input_shape[-1] + input_shape[-1] = tmp + + elif sharding_strategy == "BLOCK": + min_shard_size_y = 32 * y + min_shard_size_x = 32 * x + + rest_volume = random.randint(1, max_tensor_size // (min_shard_size_x * min_shard_size_y)) + physical_shape = random.choice(_gen_reshape_args_from_volume(rest_volume, step=1, out_dims=2)) + physical_shape = list(physical_shape["reshape_dims"]) + physical_shape[1] *= min_shard_size_y + physical_shape[0] *= min_shard_size_x + + input_shape = random.choice(_gen_reshape_args_from_volume(physical_shape[0], step=1, out_dims=rank - 1)) + input_shape = list(input_shape["reshape_dims"]) + input_shape.append(physical_shape[1]) + + elif sharding_strategy == "WIDTH" or sharding_strategy == "HEIGHT": + # if shard_width % total_cores != 0: raise RuntimeError("Invalid sharding core_grid") + # Shard Size must be multiple of input_tile_size + + if layout == "TILE_LAYOUT": + # In shard mode ShardMode::PHYSICAL, physical shard shape {12, 13312} is not compatible with alignment Alignment([32, 32])! + min_shard_size_x = 32 + min_shard_size_y = 32 * x * y + else: # if layout == "ROW_MAJOR_LAYOUT": + # Shard Size must be multiple of input_tile_size + # Shard width should be multiple of 16 to satisfy L1 alignment + mul_32_y = random.choice([16, 32, 64, 128, 256, 512, 1024]) + mul_32_x = 1024 // mul_32_y + + if sharding_strategy == "HEIGHT": + # Shard width should be multiple of 16 to satisfy L1 alignment + while mul_32_x % 16 != 0: + mul_32_x *= 2 + mul_32_y //= 2 + + min_shard_size_x = mul_32_x + min_shard_size_y = mul_32_y * x * y + + rest_volume = random.randint(1, max_tensor_size // (min_shard_size_x * min_shard_size_y)) + input_shape = random.choice(_gen_reshape_args_from_volume(rest_volume, step=1, out_dims=rank)) + input_shape = list(input_shape["reshape_dims"]) + input_shape[-2] = input_shape[-2] * min_shard_size_x + input_shape[-1] = input_shape[-1] * min_shard_size_y + + if sharding_strategy == "HEIGHT": + tmp = input_shape[-2] + input_shape[-2] = input_shape[-1] + input_shape[-1] = tmp + + # print(input_shape) + + spec_list.append( + { + "input_shape": input_shape, + "core_grid_size": (y, x), + "sharding_strategy": sharding_strategy, + "shard_orientation": shard_orientation, + "shard_height_mul_of_32": False, + "input_layout": layout, + } + ) + + return spec_list + + def gen_sharded_spec_unary_2( num_shapes, - layouts=["ROW_MAJOR_LAYOUT", "TILED_LAYOUT"], max_tensor_size_per_core=256 * 256, + layouts=["ROW_MAJOR_LAYOUT", "TILE_LAYOUT"], ): sharding_strategy_list = ["HEIGHT", "WIDTH", "BLOCK", "TENSOR_HW"] shard_orientation_list = ["COL_MAJOR", "ROW_MAJOR"] @@ -47,7 +154,7 @@ def gen_sharded_spec_unary_2( x = random.randint(1, X) max_tensor_size = y * x * max_tensor_size_per_core if sharding_strategy == "TENSOR_HW": - if input_layout == "TILED_LAYOUT": + if input_layout == "TILE_LAYOUT": min_tensor_height = 32 min_tensor_width = 32 max_tensor_height = int(math.sqrt(max_tensor_size_per_core)) @@ -197,7 +304,7 @@ def parse_sharding_spec(input_spec): input_layout = input_spec["input_layout"] assert sharding_strategy in ["HEIGHT", "WIDTH", "BLOCK", "TENSOR_HW"] - assert input_layout in ["TILED_LAYOUT", "ROW_MAJOR_LAYOUT"] + assert input_layout in ["TILE_LAYOUT", "ROW_MAJOR_LAYOUT"] tensor_hw_as_shard_shape = False @@ -216,7 +323,7 @@ def parse_sharding_spec(input_spec): else: shard_orientation = ttnn.ShardOrientation.ROW_MAJOR - if input_layout == "TILED_LAYOUT": + if input_layout == "TILE_LAYOUT": input_layout = ttnn.TILE_LAYOUT else: input_layout = ttnn.ROW_MAJOR_LAYOUT diff --git a/tests/sweep_framework/sweeps/eltwise/unary/isinf/isinf_sharded.py b/tests/sweep_framework/sweeps/eltwise/unary/isinf/isinf_sharded.py index 86e15b22a5a..ffa6d31499c 100644 --- a/tests/sweep_framework/sweeps/eltwise/unary/isinf/isinf_sharded.py +++ b/tests/sweep_framework/sweeps/eltwise/unary/isinf/isinf_sharded.py @@ -12,7 +12,7 @@ import math from tests.sweep_framework.sweep_utils.utils import gen_shapes, sanitize_shape_rm from tests.sweep_framework.sweep_utils.sharding_utils import ( - gen_sharded_spec_unary_2, + gen_sharded_spec_unary, parse_sharding_spec, invalidate_vector_sharding, roundup, @@ -35,7 +35,7 @@ # Developers can create their own generator functions and pass them to the parameters as inputs. parameters = { "nightly": { - "input_spec": gen_sharded_spec_unary_2(16), + "input_spec": gen_sharded_spec_unary(16), "input_a_dtype": [ttnn.bfloat16, ttnn.bfloat8_b], }, } diff --git a/tests/sweep_framework/sweeps/eltwise/unary/isnan/isnan_sharded.py b/tests/sweep_framework/sweeps/eltwise/unary/isnan/isnan_sharded.py index 7b157785e07..f8a95860a08 100644 --- a/tests/sweep_framework/sweeps/eltwise/unary/isnan/isnan_sharded.py +++ b/tests/sweep_framework/sweeps/eltwise/unary/isnan/isnan_sharded.py @@ -12,7 +12,7 @@ import math from tests.sweep_framework.sweep_utils.utils import gen_shapes, sanitize_shape_rm from tests.sweep_framework.sweep_utils.sharding_utils import ( - gen_sharded_spec_unary_2, + gen_sharded_spec_unary, parse_sharding_spec, invalidate_vector_sharding, roundup, @@ -35,7 +35,7 @@ # Developers can create their own generator functions and pass them to the parameters as inputs. parameters = { "nightly": { - "input_spec": gen_sharded_spec_unary_2(16), + "input_spec": gen_sharded_spec_unary(16), "input_a_dtype": [ttnn.bfloat16, ttnn.bfloat8_b], }, } diff --git a/tests/sweep_framework/sweeps/eltwise/unary/isneginf/isneginf_sharded.py b/tests/sweep_framework/sweeps/eltwise/unary/isneginf/isneginf_sharded.py index 54149dd0089..111dbded56d 100644 --- a/tests/sweep_framework/sweeps/eltwise/unary/isneginf/isneginf_sharded.py +++ b/tests/sweep_framework/sweeps/eltwise/unary/isneginf/isneginf_sharded.py @@ -12,7 +12,7 @@ import math from tests.sweep_framework.sweep_utils.utils import gen_shapes, sanitize_shape_rm from tests.sweep_framework.sweep_utils.sharding_utils import ( - gen_sharded_spec_unary_2, + gen_sharded_spec_unary, parse_sharding_spec, invalidate_vector_sharding, roundup, @@ -35,7 +35,7 @@ # Developers can create their own generator functions and pass them to the parameters as inputs. parameters = { "nightly": { - "input_spec": gen_sharded_spec_unary_2(16), + "input_spec": gen_sharded_spec_unary(16), "input_a_dtype": [ttnn.bfloat16, ttnn.bfloat8_b], }, } diff --git a/tests/sweep_framework/sweeps/eltwise/unary/isposinf/isposinf_sharded.py b/tests/sweep_framework/sweeps/eltwise/unary/isposinf/isposinf_sharded.py index 83d4ac2e1bc..bc91d23bbe5 100644 --- a/tests/sweep_framework/sweeps/eltwise/unary/isposinf/isposinf_sharded.py +++ b/tests/sweep_framework/sweeps/eltwise/unary/isposinf/isposinf_sharded.py @@ -12,7 +12,7 @@ import math from tests.sweep_framework.sweep_utils.utils import gen_shapes, sanitize_shape_rm from tests.sweep_framework.sweep_utils.sharding_utils import ( - gen_sharded_spec_unary_2, + gen_sharded_spec_unary, parse_sharding_spec, invalidate_vector_sharding, roundup, @@ -35,7 +35,7 @@ # Developers can create their own generator functions and pass them to the parameters as inputs. parameters = { "nightly": { - "input_spec": gen_sharded_spec_unary_2(16), + "input_spec": gen_sharded_spec_unary(16), "input_a_dtype": [ttnn.bfloat16, ttnn.bfloat8_b], }, } diff --git a/tests/sweep_framework/sweeps/eltwise/unary/lgamma/lgamma_sharded.py b/tests/sweep_framework/sweeps/eltwise/unary/lgamma/lgamma_sharded.py index 7be73037f16..8f1ff28addc 100644 --- a/tests/sweep_framework/sweeps/eltwise/unary/lgamma/lgamma_sharded.py +++ b/tests/sweep_framework/sweeps/eltwise/unary/lgamma/lgamma_sharded.py @@ -12,7 +12,7 @@ import math from tests.sweep_framework.sweep_utils.utils import gen_shapes, sanitize_shape_rm from tests.sweep_framework.sweep_utils.sharding_utils import ( - gen_sharded_spec_unary_2, + gen_sharded_spec_unary, parse_sharding_spec, invalidate_vector_sharding, roundup, @@ -35,7 +35,7 @@ # Developers can create their own generator functions and pass them to the parameters as inputs. parameters = { "nightly": { - "input_spec": gen_sharded_spec_unary_2(16), + "input_spec": gen_sharded_spec_unary(16), "input_a_dtype": [ttnn.bfloat16, ttnn.bfloat8_b], }, } diff --git a/tests/sweep_framework/sweeps/eltwise/unary/mish/mish_sharded.py b/tests/sweep_framework/sweeps/eltwise/unary/mish/mish_sharded.py index 0f9dd45d1d5..2f3989c982e 100644 --- a/tests/sweep_framework/sweeps/eltwise/unary/mish/mish_sharded.py +++ b/tests/sweep_framework/sweeps/eltwise/unary/mish/mish_sharded.py @@ -12,7 +12,7 @@ import math from tests.sweep_framework.sweep_utils.utils import gen_shapes, sanitize_shape_rm from tests.sweep_framework.sweep_utils.sharding_utils import ( - gen_sharded_spec_unary_2, + gen_sharded_spec_unary, parse_sharding_spec, invalidate_vector_sharding, roundup, @@ -35,7 +35,7 @@ # Developers can create their own generator functions and pass them to the parameters as inputs. parameters = { "nightly": { - "input_spec": gen_sharded_spec_unary_2(16), + "input_spec": gen_sharded_spec_unary(16), "input_a_dtype": [ttnn.bfloat16, ttnn.bfloat8_b], }, }