diff --git a/examples/python_naive_matmul/imp1_naive.py b/examples/python_naive_matmul/imp1_naive.py index c65124ca..e5340b84 100644 --- a/examples/python_naive_matmul/imp1_naive.py +++ b/examples/python_naive_matmul/imp1_naive.py @@ -2,6 +2,7 @@ import kp import numpy as np +from utils import compile_source class MatMulOp: @@ -63,7 +64,7 @@ def __init__(self, manager: kp.Manager, local_size_x: int = -1, local_size_y: in 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.format( + self.compiled_shader = compile_source(self.shader.format( local_size_x=self.local_size_x, local_size_y=self.local_size_y)) self.tensor_shape: tuple[int, int] = (0, 0) self.params: list[kp.Tensor] = [] @@ -78,7 +79,7 @@ def __call__(self, tensor_shape: tuple[int, int], tensor_in_1: kp.Tensor, tensor self.params = params local_size_x = min(self.local_size_x, tensor_shape[0]) local_size_y = min(self.local_size_y, tensor_shape[1]) - self.compiled_shader = kp.Shader.compile_source(self.shader.format( + self.compiled_shader = compile_source(self.shader.format( local_size_x=local_size_x, local_size_y=local_size_y)) workgroup = (tensor_shape[0] // local_size_x, tensor_shape[1] // local_size_y, 1) print(f'{workgroup=} {self.local_size_x=} {self.local_size_y=}') diff --git a/examples/python_naive_matmul/imp2_tiled.py b/examples/python_naive_matmul/imp2_tiled.py index 4c97c91c..ebac8f0d 100644 --- a/examples/python_naive_matmul/imp2_tiled.py +++ b/examples/python_naive_matmul/imp2_tiled.py @@ -2,7 +2,7 @@ import kp import numpy as np - +from utils import compile_source class MatMulOp: def __init__(self, manager: kp.Manager, tile_size: int = -1): @@ -66,7 +66,7 @@ def __init__(self, manager: kp.Manager, tile_size: int = -1): }} out_tensor[tensor_size * globalCol + globalRow] = acc; }}''' - self.compiled_shader = kp.Shader.compile_source(self.shader.format(tile_size=tile_size)) + self.compiled_shader = compile_source(self.shader.format(tile_size=tile_size)) self.tensor_shape: tuple[int, int] = (0, 0) self.params: list[kp.Tensor] = [] self.algo = None @@ -79,7 +79,7 @@ def __call__(self, tensor_shape: tuple[int, int], tensor_in_1: kp.Tensor, tensor self.tensor_shape = tensor_shape self.params = params tile_size = min(tensor_shape[0], tensor_shape[1], self.tile_size) - self.compiled_shader = kp.Shader.compile_source(self.shader.format(tile_size=tile_size)) + self.compiled_shader = compile_source(self.shader.format(tile_size=tile_size)) workgroup = [tensor_shape[0] // tile_size, tensor_shape[1] // tile_size, 1] self.algo = self.mgr.algorithm( params, # params diff --git a/examples/python_naive_matmul/imp3_better_tiling.py b/examples/python_naive_matmul/imp3_better_tiling.py index 2f2be60e..da0633b3 100644 --- a/examples/python_naive_matmul/imp3_better_tiling.py +++ b/examples/python_naive_matmul/imp3_better_tiling.py @@ -2,7 +2,7 @@ import kp import numpy as np - +from utils import compile_source class MatMulOp: def __init__(self, manager: kp.Manager, tile_size: int = -1, thread_work_ratio: int = 16): @@ -83,7 +83,7 @@ def __init__(self, manager: kp.Manager, tile_size: int = -1, thread_work_ratio: for(uint w = 0u; w < {thread_work_ratio}; w++) out_tensor[(globalCol + w * {local_size_y}) * tensor_size + globalRow] = acc[w]; }}''' - self.compiled_shader = kp.Shader.compile_source(self.shader.format( + self.compiled_shader = compile_source(self.shader.format( tile_size=tile_size, thread_work_ratio=thread_work_ratio, local_size_y=local_size_y)) self.tensor_shape: tuple[int, int] = (0, 0) self.params: list[kp.Tensor] = [] @@ -99,7 +99,7 @@ def __call__(self, tensor_shape: tuple[int, int], tensor_in_1: kp.Tensor, tensor tile_size = min(self.tensor_shape[0], self.tile_size) thread_work_ratio = min(self.tensor_shape[1] // self.tile_size, self.thread_work_ratio) local_size_y = tile_size // thread_work_ratio - self.compiled_shader = kp.Shader.compile_source(self.shader.format( + self.compiled_shader = compile_source(self.shader.format( tile_size=tile_size, thread_work_ratio=thread_work_ratio, local_size_y=local_size_y)) workgroup = (tensor_shape[0] // self.local_size_x, tensor_shape[1] // self.local_size_y, 1) self.algo = self.mgr.algorithm( diff --git a/examples/python_naive_matmul/utils.py b/examples/python_naive_matmul/utils.py new file mode 100644 index 00000000..d02f6add --- /dev/null +++ b/examples/python_naive_matmul/utils.py @@ -0,0 +1,7 @@ +import os + + +def compile_source(source): + open("tmp_kp_shader.comp", "w").write(source) + os.system("glslangValidator -V tmp_kp_shader.comp -o tmp_kp_shader.comp.spv") + return open("tmp_kp_shader.comp.spv", "rb").read()