Fix second implementation, add benchmark script
* Third implementation is broken (WIP)
This commit is contained in:
Corentin
parent
6f04eb9db2
commit
7f4ec27235
6 changed files with 347 additions and 40 deletions
49
examples/python_naive_matmul/benchmark.py
Normal file
49
examples/python_naive_matmul/benchmark.py
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
import time
|
||||
|
||||
import kp
|
||||
import numpy as np
|
||||
from imp1_naive import MatMulOp as MatMulOp1
|
||||
from imp2_tiled import MatMulOp as MatMulOp2
|
||||
from imp3_better_tiling import MatMulOp as MatMulOp3
|
||||
|
||||
|
||||
def main():
|
||||
experiment_count = 1000
|
||||
tensor_size = 512
|
||||
tensor_shape = [tensor_size, tensor_size]
|
||||
mat_1 = np.triu(np.ones(tensor_shape))
|
||||
mat_2 = np.triu(np.ones(tensor_shape))
|
||||
mat_result = mat_1 @ mat_2
|
||||
tensor_shape = [tensor_size, tensor_size]
|
||||
|
||||
print(f'{tensor_shape} input tensors:\n'
|
||||
f'{mat_1}\n'
|
||||
f'{mat_2}\n')
|
||||
print(f'Output :\n{mat_result}')
|
||||
|
||||
mgr = kp.Manager()
|
||||
tensor_in_1 = mgr.tensor(mat_1)
|
||||
tensor_in_2 = mgr.tensor(mat_2)
|
||||
tensor_out = mgr.tensor(np.zeros(tensor_shape))
|
||||
for MatMulOp in [MatMulOp1, MatMulOp2, MatMulOp3]:
|
||||
matmul_op = MatMulOp(mgr)
|
||||
matmul_op(tensor_shape, tensor_in_1, tensor_in_2, tensor_out)
|
||||
|
||||
start_time = time.time()
|
||||
for _ in range(experiment_count):
|
||||
matmul_op(tensor_shape, tensor_in_1, tensor_in_2, tensor_out)
|
||||
end_time = time.time()
|
||||
experiment_time = end_time - start_time
|
||||
op_count = tensor_shape[0] * tensor_shape[1] * (tensor_shape[1] - 1)
|
||||
|
||||
if (tensor_out.data().reshape(tensor_shape) == mat_result).all():
|
||||
print(f'From {MatMulOp.__module__} : {experiment_count} matmul time : '
|
||||
f'{experiment_time * 1000:0.2f}ms => '
|
||||
f'{experiment_count / experiment_time:0.2f}op/s or '
|
||||
f'{experiment_count * op_count / (1e9 * experiment_time):0.2f}GFLOPS')
|
||||
else:
|
||||
print(f'Test failed => output tensor is wrong :\n{tensor_out.data().reshape(tensor_shape)}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
|
@ -41,7 +41,7 @@ class MatMulOp:
|
|||
self.local_size_x = local_size_x
|
||||
self.local_size_y = local_size_y
|
||||
|
||||
self.shader = kp.Shader.compile_source(f'''
|
||||
self.shader = f'''
|
||||
#version 450
|
||||
|
||||
layout (local_size_x = {local_size_x}, local_size_y = {local_size_y}) in;
|
||||
|
|
@ -62,7 +62,8 @@ void main()
|
|||
for(uint k = 0u; k < tensor_size; k++)
|
||||
acc += in_tensor_1[(k * tensor_size) + globalRow] * in_tensor_2[(globalCol * tensor_size) + k];
|
||||
out_tensor[(globalCol * tensor_size) + globalRow] = acc;
|
||||
}}''')
|
||||
}}'''
|
||||
self.compiled_shader = kp.Shader.compile_source(self.shader)
|
||||
self.tensor_shape: tuple[int, int] = (0, 0)
|
||||
self.params: list[kp.Tensor] = []
|
||||
self.algo = None
|
||||
|
|
@ -74,17 +75,18 @@ void main()
|
|||
if self.algo is None or self.tensor_shape != tensor_shape or self.params != params:
|
||||
self.tensor_shape = tensor_shape
|
||||
self.params = params
|
||||
workgroup = (tensor_shape[0] // self.local_size_x, tensor_shape[1] // self.local_size_y, 1)
|
||||
self.algo = self.mgr.algorithm(
|
||||
params, # params
|
||||
self.shader, # spirv
|
||||
(tensor_shape[0] // self.local_size_x, tensor_shape[1] // self.local_size_y, 1), # workgroup
|
||||
self.compiled_shader, # spirv
|
||||
workgroup, # workgroup
|
||||
[float(tensor_shape[0])], # spec_consts
|
||||
[]) # push_consts
|
||||
|
||||
(self.mgr.sequence()
|
||||
.record(kp.OpTensorSyncDevice(self.params))
|
||||
.record(kp.OpAlgoDispatch(self.algo))
|
||||
.record(kp.OpTensorSyncLocal(self.params))
|
||||
.record(kp.OpTensorSyncLocal([tensor_out]))
|
||||
.eval())
|
||||
|
||||
|
||||
|
|
@ -121,11 +123,5 @@ def main():
|
|||
f'{experiment_count * op_count / (1e9 * experiment_time):0.2f}GFLOPS')
|
||||
|
||||
|
||||
def test():
|
||||
main()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
else:
|
||||
test()
|
||||
|
|
@ -24,7 +24,7 @@ class MatMulOp:
|
|||
assert tile_size <= max_workgroup_size[1]
|
||||
self.tile_size = tile_size
|
||||
|
||||
self.shader = kp.Shader.compile_source(f'''
|
||||
self.shader = f'''
|
||||
#version 450
|
||||
|
||||
layout (local_size_x = {tile_size}, local_size_y = {tile_size}) in;
|
||||
|
|
@ -40,20 +40,21 @@ shared float sub_tensor_2[{tile_size}][{tile_size}];
|
|||
|
||||
void main()
|
||||
{{
|
||||
uint row = gl_GlobalInvocationID.x;
|
||||
uint col = gl_GlobalInvocationID.y;
|
||||
uint globalRow = {tile_size} * gl_WorkGroupID.x + row;
|
||||
uint globalCol = {tile_size} * gl_WorkGroupID.y + row;
|
||||
uint row = gl_LocalInvocationID.x; // 0 .. tile_size
|
||||
uint col = gl_LocalInvocationID.y; // 0 .. tile_size
|
||||
// gl_WorkGroupID : 0 .. tensor_size / tile_size
|
||||
uint globalRow = ({tile_size} * gl_WorkGroupID.x) + row;
|
||||
uint globalCol = ({tile_size} * gl_WorkGroupID.y) + col;
|
||||
|
||||
uint tensor_size = uint(tensor_size_f);
|
||||
float acc = 0.0;
|
||||
uint numTiles = tensor_size / {tile_size};
|
||||
for(uint t = 0u; t < numTiles; t++)
|
||||
{{
|
||||
uint tiledRow = {tile_size} * t + row;
|
||||
uint tiledCol = {tile_size} * t + col;
|
||||
sub_tensor_1[col][row] = in_tensor_1[tiledCol * tensor_size + globalRow];
|
||||
sub_tensor_2[col][row] = in_tensor_2[globalCol * tensor_size + tiledRow];
|
||||
uint tiledRow = ({tile_size} * t) + row;
|
||||
uint tiledCol = ({tile_size} * t) + col;
|
||||
sub_tensor_1[col][row] = in_tensor_1[(tiledCol * tensor_size) + globalRow];
|
||||
sub_tensor_2[col][row] = in_tensor_2[(globalCol * tensor_size) + tiledRow];
|
||||
|
||||
memoryBarrierShared();
|
||||
barrier();
|
||||
|
|
@ -63,8 +64,10 @@ void main()
|
|||
|
||||
barrier();
|
||||
}}
|
||||
out_tensor[(globalCol * tensor_size) + globalRow] = acc;
|
||||
}}''')
|
||||
uint globalIndex = (tensor_size * globalCol) + globalRow;
|
||||
out_tensor[globalIndex] = acc;
|
||||
}}'''
|
||||
self.compiled_shader = kp.Shader.compile_source(self.shader)
|
||||
self.tensor_shape: tuple[int, int] = (0, 0)
|
||||
self.params: list[kp.Tensor] = []
|
||||
self.algo = None
|
||||
|
|
@ -76,17 +79,18 @@ void main()
|
|||
if self.algo is None or self.tensor_shape != tensor_shape or self.params != params:
|
||||
self.tensor_shape = tensor_shape
|
||||
self.params = params
|
||||
workgroup = (tensor_shape[0] // self.tile_size, tensor_shape[1] // self.tile_size, 1)
|
||||
self.algo = self.mgr.algorithm(
|
||||
params, # params
|
||||
self.shader, # spirv
|
||||
(tensor_shape[0] // self.tile_size, tensor_shape[1] // self.tile_size, 1), # workgroup
|
||||
self.compiled_shader, # spirv
|
||||
workgroup, # workgroup
|
||||
[float(tensor_shape[0])], # spec_consts
|
||||
[]) # push_consts
|
||||
|
||||
(self.mgr.sequence()
|
||||
.record(kp.OpTensorSyncDevice(self.params))
|
||||
.record(kp.OpAlgoDispatch(self.algo))
|
||||
.record(kp.OpTensorSyncLocal(self.params))
|
||||
.record(kp.OpTensorSyncLocal([tensor_out]))
|
||||
.eval())
|
||||
|
||||
|
||||
156
examples/python_naive_matmul/imp2_tiled_debug.py
Normal file
156
examples/python_naive_matmul/imp2_tiled_debug.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
import time
|
||||
|
||||
import kp
|
||||
import numpy as np
|
||||
|
||||
|
||||
class MatMulOp:
|
||||
def __init__(self, manager: kp.Manager, tile_size: int = -1):
|
||||
self.mgr = manager
|
||||
|
||||
props = self.mgr.get_device_properties()
|
||||
max_workgroup_invocation = props['max_work_group_invocations']
|
||||
max_workgroup_size = props['max_work_group_size']
|
||||
if tile_size < 0:
|
||||
tile_size = 1
|
||||
while (4 * tile_size * tile_size <= max_workgroup_invocation
|
||||
and 2 * tile_size <= max_workgroup_size[0]
|
||||
and 2 * tile_size <= max_workgroup_size[1]):
|
||||
tile_size *= 2
|
||||
|
||||
assert tile_size > 0
|
||||
assert tile_size * tile_size <= max_workgroup_invocation
|
||||
assert tile_size <= max_workgroup_size[0]
|
||||
assert tile_size <= max_workgroup_size[1]
|
||||
self.tile_size = tile_size
|
||||
|
||||
print(f'{tile_size=}')
|
||||
self.shader = f'''
|
||||
#version 450
|
||||
|
||||
layout (local_size_x = {tile_size}, local_size_y = {tile_size}) in;
|
||||
|
||||
layout (set = 0, binding = 0) readonly buffer buf_in_tensor_1 {{ float in_tensor_1[]; }};
|
||||
layout (set = 0, binding = 1) readonly buffer buf_in_tensor_2 {{ float in_tensor_2[]; }};
|
||||
layout (set = 0, binding = 2) writeonly buffer buf_out_tensor {{ float out_tensor[]; }};
|
||||
layout (set = 0, binding = 3) writeonly buffer buf_test1_tensor {{ float test1_tensor[]; }};
|
||||
layout (set = 0, binding = 4) writeonly buffer buf_test2_tensor {{ float test2_tensor[]; }};
|
||||
layout (set = 0, binding = 5) writeonly buffer buf_test3_tensor {{ float test3_tensor[]; }};
|
||||
layout (set = 0, binding = 6) writeonly buffer buf_test4_tensor {{ float test4_tensor[]; }};
|
||||
|
||||
layout (constant_id = 0) const float tensor_size_f = 0;
|
||||
|
||||
shared float sub_tensor_1[{tile_size}][{tile_size}];
|
||||
shared float sub_tensor_2[{tile_size}][{tile_size}];
|
||||
|
||||
void main()
|
||||
{{
|
||||
uint row = gl_LocalInvocationID.x; // 0 .. tile_size
|
||||
uint col = gl_LocalInvocationID.y; // 0 .. tile_size
|
||||
// gl_WorkGroupID : 0 .. tensor_size / tile_size
|
||||
uint globalRow = ({tile_size} * gl_WorkGroupID.x) + row;
|
||||
uint globalCol = ({tile_size} * gl_WorkGroupID.y) + col;
|
||||
|
||||
uint tensor_size = uint(tensor_size_f);
|
||||
float acc = 0.0;
|
||||
uint numTiles = tensor_size / {tile_size};
|
||||
for(uint t = 0u; t < numTiles; t++)
|
||||
{{
|
||||
uint tiledRow = ({tile_size} * t) + row;
|
||||
uint tiledCol = ({tile_size} * t) + col;
|
||||
sub_tensor_1[col][row] = in_tensor_1[(tiledCol * tensor_size) + globalRow];
|
||||
sub_tensor_2[col][row] = in_tensor_2[(globalCol * tensor_size) + tiledRow];
|
||||
|
||||
memoryBarrierShared();
|
||||
barrier();
|
||||
|
||||
for(uint k = 0u; k < {tile_size}; k++)
|
||||
acc += sub_tensor_1[k][row] * sub_tensor_2[col][k];
|
||||
|
||||
barrier();
|
||||
}}
|
||||
uint globalIndex = (tensor_size * globalCol) + globalRow;
|
||||
out_tensor[globalIndex] = acc;
|
||||
test1_tensor[globalIndex] = row;
|
||||
test2_tensor[globalIndex] = col;
|
||||
test3_tensor[globalIndex] = gl_WorkGroupID.x;
|
||||
test4_tensor[globalIndex] = gl_WorkGroupID.y;
|
||||
}}'''
|
||||
print(self.shader)
|
||||
self.compiled_shader = kp.Shader.compile_source(self.shader)
|
||||
self.tensor_shape: tuple[int, int] = (0, 0)
|
||||
self.params: list[kp.Tensor] = []
|
||||
self.algo = None
|
||||
|
||||
def __call__(self, tensor_shape: tuple[int, int], tensor_in_1: kp.Tensor, tensor_in_2: kp.Tensor,
|
||||
tensor_out: kp.Tensor, tensor_test_1: kp.Tensor, tensor_test_2: kp.Tensor,
|
||||
tensor_test_3: kp.Tensor, tensor_test_4: kp.Tensor):
|
||||
# params = [tensor_in_1, tensor_in_2, tensor_out]
|
||||
params = [tensor_in_1, tensor_in_2, tensor_out, tensor_test_1, tensor_test_2, tensor_test_3, tensor_test_4]
|
||||
|
||||
if self.algo is None or self.tensor_shape != tensor_shape or self.params != params:
|
||||
self.tensor_shape = tensor_shape
|
||||
self.params = params
|
||||
workgroup = (tensor_shape[0] // self.tile_size, tensor_shape[1] // self.tile_size, 1)
|
||||
# workgroup = (2, 2, 1)
|
||||
print(f'{float(tensor_shape[0])=} {self.tile_size=} {workgroup=}')
|
||||
self.algo = self.mgr.algorithm(
|
||||
params, # params
|
||||
self.compiled_shader, # spirv
|
||||
workgroup, # workgroup
|
||||
[float(tensor_shape[0])], # spec_consts
|
||||
[]) # push_consts
|
||||
|
||||
(self.mgr.sequence()
|
||||
.record(kp.OpTensorSyncDevice(self.params))
|
||||
.record(kp.OpAlgoDispatch(self.algo))
|
||||
.record(kp.OpTensorSyncLocal(self.params[2:]))
|
||||
# .record(kp.OpTensorSyncLocal([tensor_out]))
|
||||
.eval())
|
||||
|
||||
|
||||
def main():
|
||||
mgr = kp.Manager()
|
||||
|
||||
matmul_op = MatMulOp(mgr, 4)
|
||||
|
||||
tensor_size = 8
|
||||
tensor_shape = [tensor_size, tensor_size]
|
||||
tensor_in_1 = mgr.tensor(np.triu(np.ones(tensor_shape)))
|
||||
tensor_in_2 = mgr.tensor(np.triu(np.ones(tensor_shape)))
|
||||
tensor_out = mgr.tensor(np.zeros(tensor_shape))
|
||||
tensor_test_1 = mgr.tensor(np.zeros(tensor_shape))
|
||||
tensor_test_2 = mgr.tensor(np.zeros(tensor_shape))
|
||||
tensor_test_3 = mgr.tensor(np.zeros(tensor_shape))
|
||||
tensor_test_4 = mgr.tensor(np.zeros(tensor_shape))
|
||||
|
||||
print(f'{tensor_shape} input tensors:\n'
|
||||
f'{tensor_in_1.data().reshape(tensor_shape)}\n'
|
||||
f'{tensor_in_2.data().reshape(tensor_shape)}\n')
|
||||
|
||||
# matmul_op(tensor_shape, tensor_in_1, tensor_in_2, tensor_out)
|
||||
matmul_op(tensor_shape, tensor_in_1, tensor_in_2,
|
||||
tensor_out, tensor_test_1, tensor_test_2, tensor_test_3, tensor_test_4)
|
||||
|
||||
# experiment_count = 10
|
||||
# start_time = time.time()
|
||||
# for _ in range(experiment_count):
|
||||
# matmul_op(tensor_shape, tensor_in_1, tensor_in_2, tensor_out)
|
||||
# end_time = time.time()
|
||||
# experiment_time = end_time - start_time
|
||||
# op_count = tensor_shape[0] * tensor_shape[1] * (tensor_shape[1] - 1)
|
||||
|
||||
print(f'Output :\n{tensor_out.data().reshape(tensor_shape)}')
|
||||
print(f'test_1 :\n{tensor_test_1.data().reshape(tensor_shape)}')
|
||||
print(f'test_2 :\n{tensor_test_2.data().reshape(tensor_shape)}')
|
||||
print(f'test_3 :\n{tensor_test_3.data().reshape(tensor_shape)}')
|
||||
print(f'test_4 :\n{tensor_test_4.data().reshape(tensor_shape)}')
|
||||
|
||||
# print(f'{experiment_count} matmul time : '
|
||||
# f'{experiment_time * 1000:0.2f}ms => '
|
||||
# f'{experiment_count / experiment_time:0.2f}op/s or '
|
||||
# f'{experiment_count * op_count / (1e9 * experiment_time):0.2f}GFLOPS')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
|
@ -30,11 +30,12 @@ class MatMulOp:
|
|||
self.tile_size = tile_size
|
||||
self.thread_work_ratio = thread_work_ratio
|
||||
|
||||
local_size_y = tile_size // thread_work_ratio
|
||||
self.shader = kp.Shader.compile_source(f'''
|
||||
self.local_size_x = tile_size
|
||||
self.local_size_y = tile_size // thread_work_ratio
|
||||
self.shader = f'''
|
||||
#version 450
|
||||
|
||||
layout (local_size_x = {tile_size}, local_size_y = {local_size_y}) in;
|
||||
layout (local_size_x = {tile_size}, local_size_y = {self.local_size_y}) in;
|
||||
|
||||
layout (set = 0, binding = 0) readonly buffer buf_in_tensor_1 {{ float in_tensor_1[]; }};
|
||||
layout (set = 0, binding = 1) readonly buffer buf_in_tensor_2 {{ float in_tensor_2[]; }};
|
||||
|
|
@ -47,8 +48,8 @@ shared float sub_tensor_2[{tile_size}][{tile_size}];
|
|||
|
||||
void main()
|
||||
{{
|
||||
uint row = gl_GlobalInvocationID.x;
|
||||
uint col = gl_GlobalInvocationID.y;
|
||||
uint row = gl_LocalInvocationID.x;
|
||||
uint col = gl_LocalInvocationID.y;
|
||||
uint globalRow = {tile_size} * gl_WorkGroupID.x + row;
|
||||
uint globalCol = {tile_size} * gl_WorkGroupID.y + row;
|
||||
|
||||
|
|
@ -62,23 +63,24 @@ void main()
|
|||
{{
|
||||
uint tiledRow = {tile_size} * t + row;
|
||||
uint tiledCol = {tile_size} * t + col;
|
||||
sub_tensor_1[col + t * {local_size_y}][row] = in_tensor_1[
|
||||
(tiledCol + t * {local_size_y}) * tensor_size + globalRow];
|
||||
sub_tensor_2[col + t * {local_size_y}][row] = in_tensor_2[
|
||||
(globalCol + t * {local_size_y})* tensor_size + tiledRow];
|
||||
sub_tensor_1[col + t * {self.local_size_y}][row] = in_tensor_1[
|
||||
(tiledCol + t * {self.local_size_y}) * tensor_size + globalRow];
|
||||
sub_tensor_2[col + t * {self.local_size_y}][row] = in_tensor_2[
|
||||
(globalCol + t * {self.local_size_y})* tensor_size + tiledRow];
|
||||
|
||||
memoryBarrierShared();
|
||||
barrier();
|
||||
|
||||
for(uint k = 0u; k < {tile_size}; k++)
|
||||
for(uint l = 0u; l < {thread_work_ratio}; l++)
|
||||
acc[l] += sub_tensor_1[k][row] * sub_tensor_2[col + l * {local_size_y}][k];
|
||||
acc[l] += sub_tensor_1[k][row] * sub_tensor_2[col + l * {self.local_size_y}][k];
|
||||
|
||||
barrier();
|
||||
}}
|
||||
for(uint l = 0u; l < {thread_work_ratio}; l++)
|
||||
out_tensor[(globalCol + l * {local_size_y}) * tensor_size + globalRow] = acc[l];
|
||||
}}''')
|
||||
out_tensor[(globalCol + l * {self.local_size_y}) * tensor_size + globalRow] = acc[l];
|
||||
}}'''
|
||||
self.compiled_shader = kp.Shader.compile_source(self.shader)
|
||||
self.tensor_shape: tuple[int, int] = (0, 0)
|
||||
self.params: list[kp.Tensor] = []
|
||||
self.algo = None
|
||||
|
|
@ -90,17 +92,20 @@ void main()
|
|||
if self.algo is None or self.tensor_shape != tensor_shape or self.params != params:
|
||||
self.tensor_shape = tensor_shape
|
||||
self.params = params
|
||||
print(
|
||||
tensor_shape, self.local_size_x, self.local_size_y,
|
||||
(tensor_shape[0] // self.local_size_x, tensor_shape[1] // self.local_size_y, 1))
|
||||
self.algo = self.mgr.algorithm(
|
||||
params, # params
|
||||
self.shader, # spirv
|
||||
(tensor_shape[0] // self.tile_size, tensor_shape[1] // self.tile_size, 1), # workgroup
|
||||
self.compiled_shader, # spirv
|
||||
(tensor_shape[0] // self.local_size_x, tensor_shape[1] // self.local_size_y, 1), # workgroup
|
||||
[float(tensor_shape[0])], # spec_consts
|
||||
[]) # push_consts
|
||||
|
||||
(self.mgr.sequence()
|
||||
.record(kp.OpTensorSyncDevice(self.params))
|
||||
.record(kp.OpAlgoDispatch(self.algo))
|
||||
.record(kp.OpTensorSyncLocal(self.params))
|
||||
.record(kp.OpTensorSyncLocal([tensor_out]))
|
||||
.eval())
|
||||
|
||||
|
||||
97
examples/python_naive_matmul/matmul_plot.py
Normal file
97
examples/python_naive_matmul/matmul_plot.py
Normal file
|
|
@ -0,0 +1,97 @@
|
|||
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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue