diff --git a/tests/sweep_framework/sweep_utils/sharding_utils.py b/tests/sweep_framework/sweep_utils/sharding_utils.py new file mode 100644 index 000000000000..402cf44ca935 --- /dev/null +++ b/tests/sweep_framework/sweep_utils/sharding_utils.py @@ -0,0 +1,137 @@ +import ttnn +import itertools +import random +from tests.tt_eager.python_api_testing.sweep_tests.generation_funcs import _gen_reshape_args_from_volume + + +def gen_sharded_spec_unary(num_shapes, y, x, max_tensor_size=4 * 1024 * 1024): + # ["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], ["TILE_LAYOUT", "ROW_MAJOR_LAYOUT"] + ): + if sharding_strategy == "tensor_wh": + tensor_hw_as_shard_shape = True + sharding_strategy = "BLOCK" + else: + tensor_hw_as_shard_shape = False + + for _ in range(num_shapes): + if tensor_hw_as_shard_shape: + # 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 + mul_x = random.randint(1, 10) + mul_y = random.randint(1, 64 // mul_x) + + input_shape = random.choice( + _gen_reshape_args_from_volume(mul_y * min_shard_size_y, step=1, out_dims=rank - 1) + ) + input_shape = list(input_shape["reshape_dims"]) + input_shape.append(mul_x * min_shard_size_x) + + 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, + "sharding_strategy": sharding_strategy, + "shard_orientation": shard_orientation, + "tensor_hw_as_shard_shape": tensor_hw_as_shard_shape, + "input_layout": layout, + } + ) + + return spec_list + + +def parse_sharding_spec(input_spec): + input_shape = input_spec["input_shape"] + sharding_strategy = input_spec["sharding_strategy"] + shard_orientation = input_spec["shard_orientation"] + tensor_hw_as_shard_shape = input_spec["tensor_hw_as_shard_shape"] + input_layout = input_spec["input_layout"] + + if sharding_strategy == "HEIGHT": + sharding_strategy = ttnn.ShardStrategy.HEIGHT + elif sharding_strategy == "WIDTH": + sharding_strategy = ttnn.ShardStrategy.WIDTH + else: # sharding_strategy == "BLOCK": + sharding_strategy = ttnn.ShardStrategy.BLOCK + + if shard_orientation == "COL_MAJOR": + shard_orientation = ttnn.ShardOrientation.COL_MAJOR + else: + shard_orientation = ttnn.ShardOrientation.ROW_MAJOR + + if input_layout == "TILE_LAYOUT": + input_layout = ttnn.TILE_LAYOUT + else: + input_layout = ttnn.ROW_MAJOR_LAYOUT + + return input_shape, sharding_strategy, shard_orientation, tensor_hw_as_shard_shape, input_layout diff --git a/tests/sweep_framework/sweeps/eltwise/unary/isfinite/isfinite_sharded.py b/tests/sweep_framework/sweeps/eltwise/unary/isfinite/isfinite_sharded.py index 9eb598f989de..fec93f8f95dc 100644 --- a/tests/sweep_framework/sweeps/eltwise/unary/isfinite/isfinite_sharded.py +++ b/tests/sweep_framework/sweeps/eltwise/unary/isfinite/isfinite_sharded.py @@ -4,14 +4,15 @@ from typing import Optional, Tuple from functools import partial -import itertools +import json import torch import random import ttnn import math from tests.sweep_framework.sweep_utils.utils import gen_shapes, sanitize_shape_rm, get_device_grid_size -from tests.tt_eager.python_api_testing.sweep_tests.generation_funcs import gen_rand_inf, _gen_reshape_args_from_volume +from tests.sweep_framework.sweep_utils.sharding_utils import gen_sharded_spec_unary, parse_sharding_spec +from tests.tt_eager.python_api_testing.sweep_tests.generation_funcs import gen_rand_inf from tests.ttnn.utils_for_testing import check_with_pcc, start_measuring_time, stop_measuring_time from models.utility_functions import torch_random @@ -23,120 +24,13 @@ random.seed(0) -def gen_sharded_spec(num_shapes, y, x, max_tensor_size=8 * 1024 * 1024): - # [ttnn.ShardStrategy.BLOCK, ttnn.ShardStrategy.WIDTH, ttnn.ShardStrategy.HEIGHT, "tensor_wh"] - sharding_strategy_list = ["tensor_wh"] - shard_orientation_list = [ttnn.ShardOrientation.COL_MAJOR, ttnn.ShardOrientation.ROW_MAJOR] - spec_list = [] - - for sharding_strategy, shard_orientation, rank, layout in itertools.product( - sharding_strategy_list, shard_orientation_list, [4, 3, 2], [ttnn.TILE_LAYOUT, ttnn.ROW_MAJOR_LAYOUT] - ): - if sharding_strategy == "tensor_wh": - tensor_hw_as_shard_shape = True - sharding_strategy = ttnn.ShardStrategy.BLOCK - else: - tensor_hw_as_shard_shape = False - - for _ in range(num_shapes): - if tensor_hw_as_shard_shape: - # 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 == ttnn.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 == ttnn.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 == ttnn.ShardOrientation.COL_MAJOR: - tmp = input_shape[-2] - input_shape[-2] = input_shape[-1] - input_shape[-1] = tmp - - elif sharding_strategy == ttnn.ShardStrategy.BLOCK: - min_shard_size_y = 32 * y - min_shard_size_x = 32 * x - mul_x = random.randint(1, 10) - mul_y = random.randint(1, 64 // mul_x) - - input_shape = random.choice( - _gen_reshape_args_from_volume(mul_y * min_shard_size_y, step=1, out_dims=rank - 1) - ) - input_shape = list(input_shape["reshape_dims"]) - input_shape.append(mul_x * min_shard_size_x) - - elif sharding_strategy == ttnn.ShardStrategy.WIDTH or sharding_strategy == ttnn.ShardStrategy.HEIGHT: - # if shard_width % total_cores != 0: raise RuntimeError("Invalid sharding core_grid") - # Shard Size must be multiple of input_tile_size - - if layout == ttnn.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 == ttnn.ROW_MAJOR_LAYOUT: - min_shard_size_x = 1 - min_shard_size_y = x * y - - # Shard Size must be multiple of input_tile_size - mul_32_x = random.choice([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]) - mul_32_y = 1024 // mul_32_x - - min_shard_size_x *= mul_32_x - min_shard_size_y *= mul_32_y - - if sharding_strategy == ttnn.ShardStrategy.HEIGHT: - # Shard width should be multiple of 16 to satisfy L1 alignment - while min_shard_size_x % 16 != 0: - min_shard_size_x *= 2 - - 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 == ttnn.ShardStrategy.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, - "sharding_strategy": sharding_strategy, - "shard_orientation": shard_orientation, - "tensor_hw_as_shard_shape": tensor_hw_as_shard_shape, - "input_layout": layout, - } - ) - - return spec_list - - # Parameters provided to the test vector generator are defined here. # They are defined as dict-type suites that contain the arguments to the run function as keys, and lists of possible inputs as values. # Each suite has a key name (in this case "suite_1" and "suite_2") which will associate the test vectors to this specific suite of inputs. # Developers can create their own generator functions and pass them to the parameters as inputs. parameters = { "nightly": { - "input_spec": gen_sharded_spec(16, Y, X), + "input_spec": gen_sharded_spec_unary(16, Y, X), "input_a_dtype": [ttnn.bfloat16, ttnn.bfloat8_b], }, } @@ -168,12 +62,15 @@ def run( ) -> list: data_seed = random.randint(0, 20000000) torch.manual_seed(data_seed) - input_shape, sharding_strategy, shard_orientation, tensor_hw_as_shard_shape, input_layout = input_spec.values() - print( - f"X {X} Y {Y} input_shape {input_shape} {input_a_dtype} {input_layout} {sharding_strategy} {shard_orientation} tensor_hw_as_shard_shape {tensor_hw_as_shard_shape}" + input_shape, sharding_strategy, shard_orientation, tensor_hw_as_shard_shape, input_layout = parse_sharding_spec( + input_spec ) + # print( + # f"X {X} Y {Y} input_shape {input_shape} {input_a_dtype} {input_layout} {sharding_strategy} {shard_orientation} tensor_hw_as_shard_shape {tensor_hw_as_shard_shape}" + # ) + if input_layout == ttnn.ROW_MAJOR_LAYOUT: input_shape = sanitize_shape_rm(input_shape) @@ -202,34 +99,34 @@ def run( output_tensor = ttnn.to_torch(output_tensor) pcc = check_with_pcc(torch_output_tensor, output_tensor, 0.999) - print(pcc) + # print(pcc) return [check_with_pcc(torch_output_tensor, output_tensor, 0.999), e2e_perf] -# Run sweeps locally -from tests.sweep_framework.framework.permutations import * - -start_time = start_measuring_time() -for suite in parameters.keys(): - device_id = 0 - device = ttnn.open_device(device_id=device_id) - suite_vectors = list(permutations(parameters[suite])) - print(len(suite_vectors)) - for vector in suite_vectors: - invalidate_res = invalidate_vector(vector) - if invalidate_res[0]: - print(f"Invalidated: {invalidate_res[1]}") - continue - try: - passed, _ = run(**vector, device=device) - # if passed[0] != True: - # print(passed) - except Exception as e: - print(e) - - # break - - ttnn.close_device(device) - -e2e_perf = stop_measuring_time(start_time) -print(f"time {e2e_perf / 1000000000}s") +# # Run sweeps locally +# from tests.sweep_framework.framework.permutations import * + +# start_time = start_measuring_time() +# for suite in parameters.keys(): +# device_id = 0 +# device = ttnn.open_device(device_id=device_id) +# suite_vectors = list(permutations(parameters[suite])) +# print(len(suite_vectors)) +# for vector in suite_vectors: +# invalidate_res = invalidate_vector(vector) +# if invalidate_res[0]: +# print(f"Invalidated: {invalidate_res[1]}") +# continue +# try: +# passed, _ = run(**vector, device=device) +# # if passed[0] != True: +# # print(passed) +# except Exception as e: +# print(e) + +# # break + +# ttnn.close_device(device) + +# e2e_perf = stop_measuring_time(start_time) +# print(f"time {e2e_perf / 1000000000}s")