fix

Author: yaoyaoding

examples/quantization/matmul_a16wx.py

@@ -81,8 +81,9 @@ def __call__(self, k_size: int, n_size: int, src_ptr: void_p, dst_ptr: void_p):
 
         g_src = self.global_view(src_ptr, dtype=self.b_dtype, shape=[k_size, n_size])
         r_src = self.load_global(
-            g_src, offsets=[offset_k, offset_n], layout=self.tile_layout
+            g_src, offsets=[offset_k, offset_n], shape=self.tile_layout.shape
         )
+        self.annotate_layout(r_src, self.tile_layout)
         r_dst = self.view(r_src, layout=self.flatten_tile_layout, dtype=uint8)
         g_dst = self.global_view(
             dst_ptr,
@@ -127,8 +128,9 @@ def __call__(self, k_size: int, n_size: int, src_ptr: void_p, dst_ptr: void_p):
         r_src = self.load_global(
             g_src,
             offsets=[self.blockIdx.x, self.blockIdx.y, 0],
-            layout=self.flatten_tile_layout,
+            shape=self.flatten_tile_layout.shape,
         )
+        self.annotate_layout(r_src, self.flatten_tile_layout)
         r_dst = self.view(r_src, layout=self.tile_layout, dtype=self.b_dtype)
         g_dst = self.global_view(dst_ptr, dtype=self.b_dtype, shape=[k_size, n_size])
         self.store_global(g_dst, r_dst, offsets=[offset_k, offset_n], dims=[2])

python/tilus/lang/transpiler.py

@@ -689,7 +689,10 @@ def visit_Call(self, expr: ast.Call) -> Any:
                         raise TilusProgramError(self, expr, str(e))
                 else:
                     # case 4
-                    ret = func(*args, **kwargs)
+                    try:
+                        ret = func(*args, **kwargs)
+                    except TypeError as e:
+                        raise TilusProgramError(self, expr, str(e)) from e
             elif isinstance(func, types.FunctionType):
                 # case 4
                 ret = func(*args, **kwargs)

tests/instructions/test_where.py

@@ -41,13 +41,15 @@ def __call__(self, m: int32, n: int32, cond_ptr: ~boolean, x_ptr: ~int32, y_ptr:
         gy = self.global_view(y_ptr, dtype=int32, shape=(m, n))
         go = self.global_view(out_ptr, dtype=int32, shape=(m, n))
 
-        rc = self.load_global(gc, offsets=[m_offset, n_offset], layout=self.layout)
+        rc = self.load_global(gc, offsets=[m_offset, n_offset], shape=self.layout.shape)
         rx = self.load_global(gx, offsets=[m_offset, n_offset], shape=[self.block_m, self.block_n])
         ry = self.load_global(gy, offsets=[m_offset, n_offset], shape=[self.block_m, self.block_n])
         ro = self.where(rc, rx, ry)
 
         self.store_global(go, ro, offsets=[m_offset, n_offset])
 
+        self.annotate_layout(rc, self.layout)
+
 
 @pytest.mark.parametrize("m, n, layout", [[16, 16, spatial(4, 8)], [128, 128, local(2, 2).spatial(4, 8).local(2, 2)]])
 def test_where(
@@ -65,3 +67,7 @@ def test_where(
     kernel(m, n, cond, x, y, actual)
 
     assert torch.allclose(actual, expected), f"Failed for layout {layout} with m={m}, n={n}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])