Skip to content

Commit

Permalink
#11512: Improve isfinite sharded sweeps
Browse files Browse the repository at this point in the history
  • Loading branch information
@nemanjagrujic
nemanjagrujic committed Nov 20, 2024
1 parent 2b9b22b commit b759cd1
Show file tree
Hide file tree
Showing 2 changed files with 175 additions and 141 deletions.
137 changes: 137 additions & 0 deletions tests/sweep_framework/sweep_utils/sharding_utils.py
View file
Original file line number Diff line number Diff line change
@@ -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
179 changes: 38 additions & 141 deletions tests/sweep_framework/sweeps/eltwise/unary/isfinite/isfinite_sharded.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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],
},
}
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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")

0 comments on commit b759cd1

Please sign in to comment.