Update Python Matmul examples (#434)

* feat: Add utility function to compile GLSL source to SPIR-V

Signed-off-by: n-jay <nuwanjayawardene@gmail.com>

* refactor: Replace kp.Shader.compile_source with utility function compile_source

Signed-off-by: n-jay <nuwanjayawardene@gmail.com>

---------

Signed-off-by: n-jay <nuwanjayawardene@gmail.com>

Author: n-jay

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=}')

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

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(

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()