mirror of
https://github.com/nomic-ai/kompute.git
synced 2026-06-30 03:17:12 +00:00
98 lines
4.1 KiB
Python
98 lines
4.1 KiB
Python
from argparse import ArgumentParser
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def plot_tensor(window_name: str, tensor: np.ndarray, coord_highlight: tuple[int, int] = None):
|
|
font_size = 48
|
|
image = np.zeros((tensor.shape[1] * font_size, tensor.shape[0] * font_size, 3), dtype=np.uint8)
|
|
|
|
for y in range(tensor.shape[1]):
|
|
for x in range(tensor.shape[0]):
|
|
if coord_highlight and x == coord_highlight[1] and y == coord_highlight[0]:
|
|
cv2.putText(
|
|
image, str(int(tensor[y, x])), (x * font_size, int((y + 0.8) * font_size)),
|
|
cv2.FONT_HERSHEY_TRIPLEX, 1., (127, 127, 255))
|
|
else:
|
|
cv2.putText(
|
|
image, str(int(tensor[y, x])), (x * font_size, int((y + 0.8) * font_size)),
|
|
cv2.FONT_HERSHEY_TRIPLEX, 1., (255, 255, 255))
|
|
|
|
cv2.imshow(window_name, image)
|
|
|
|
|
|
def main():
|
|
parser = ArgumentParser()
|
|
parser.add_argument('tensor_size', type=int, help='Size of the square tensors')
|
|
parser.add_argument('tile_size', type=int)
|
|
parser.add_argument('local_size', type=int, nargs=2)
|
|
parser.add_argument('workgroup', type=int, nargs=2)
|
|
arguments = parser.parse_args()
|
|
|
|
tensor_size: int = arguments.tensor_size
|
|
tile_size: int = arguments.tile_size
|
|
local_size: tuple[int, int, int] = tuple(arguments.local_size)
|
|
workgroup: tuple[int, int, int] = tuple(arguments.workgroup)
|
|
|
|
tensor_shape = (tensor_size, tensor_size)
|
|
tensor_1 = np.triu(np.ones(tensor_shape))
|
|
tensor_2 = np.triu(np.ones(tensor_shape))
|
|
tensor_out = np.zeros(tensor_shape)
|
|
tensor_test_1 = np.zeros(tensor_shape)
|
|
tensor_test_2 = np.zeros(tensor_shape)
|
|
tensor_test_3 = np.zeros(tensor_shape)
|
|
tensor_test_4 = np.zeros(tensor_shape)
|
|
tensor_test_5 = np.zeros(tensor_shape)
|
|
|
|
plot_tensor('tensor_1', tensor_1)
|
|
plot_tensor('tensor_2', tensor_2)
|
|
plot_tensor('tensor_out', tensor_out)
|
|
plot_tensor('tensor_test_1', tensor_test_1)
|
|
plot_tensor('tensor_test_2', tensor_test_2)
|
|
plot_tensor('tensor_test_3', tensor_test_3)
|
|
plot_tensor('tensor_test_4', tensor_test_4)
|
|
plot_tensor('tensor_test_5', tensor_test_5)
|
|
cv2.waitKey(-1)
|
|
|
|
print(f'{workgroup=} {local_size=}')
|
|
for workgroup_x in range(workgroup[0]):
|
|
for workgroup_y in range(workgroup[1]):
|
|
for invocation_x in range(workgroup_x * local_size[0], (workgroup_x + 1) * local_size[0]):
|
|
for invocation_y in range(workgroup_y * local_size[1], (workgroup_y + 1) * local_size[1]):
|
|
row = invocation_x
|
|
col = invocation_y
|
|
globalRow = (tile_size * workgroup_x) + row
|
|
globalCol = (tile_size * workgroup_y) + col
|
|
try:
|
|
tensor_out[row, col] = row
|
|
tensor_test_1[row, col] = col
|
|
tensor_test_2[row, col] = workgroup_x
|
|
tensor_test_3[row, col] = workgroup_y
|
|
tensor_test_4[row, col] = globalRow
|
|
tensor_test_5[row, col] = globalCol
|
|
plot_tensor('tensor_out', tensor_out, (row, col))
|
|
plot_tensor('tensor_test_1', tensor_test_1, (row, col))
|
|
plot_tensor('tensor_test_2', tensor_test_2, (row, col))
|
|
plot_tensor('tensor_test_3', tensor_test_3, (row, col))
|
|
plot_tensor('tensor_test_4', tensor_test_4, (row, col))
|
|
plot_tensor('tensor_test_5', tensor_test_5, (row, col))
|
|
cv2.waitKey(-1)
|
|
except IndexError as error:
|
|
print(f'{workgroup_x=} {workgroup_y=} {row=} {col=}')
|
|
raise error
|
|
|
|
plot_tensor('tensor_1', tensor_1)
|
|
plot_tensor('tensor_2', tensor_2)
|
|
plot_tensor('tensor_out', tensor_out)
|
|
plot_tensor('tensor_test_1', tensor_test_1)
|
|
plot_tensor('tensor_test_2', tensor_test_2)
|
|
plot_tensor('tensor_test_3', tensor_test_3)
|
|
plot_tensor('tensor_test_4', tensor_test_4)
|
|
plot_tensor('tensor_test_5', tensor_test_5)
|
|
cv2.waitKey(-1)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|