-
Notifications
You must be signed in to change notification settings - Fork 15
Open
Description
The current calling convention in Tilus:
import tilus
from tilus import float16, int32, GlobalTensor
class Example(tilus.Script):
def __call__(
self, m_size: int32, n_size: int, k_size: int, a_ptr: ~flaot16, b_ptr: ~flaot16, c_ptr: ~flaot16
): # m_size is dynamic, n_size and k_size are JIT constants
"""
Generate:
void example(int m, int n, int k, float16 *a, float16 *b, float16 *c) {...}
"""
...
g_a: GlobalTensor = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size])
...The L1 calling convention:
class ExampleL1(tilus.Script):
def __call__(
self,
g_a: GlobalTensor[float16, ['m', 'K']], # m is dynamic,
g_b: GlobalTensor[float16, ['K', 'N']], # K and N are JIT constants
g_c: GlobalTensor[float16, ['m', 'N']],
):
"""
Generate:
void exampleL1(tvm::ffi::TensorView g_a, tvm::ffi::TensorView g_b, tvm::ffi::TensorView g_c) {...}
"""
...We can also use
g_a: 'float16[m, K]',
g_b: 'float16[K, N]',
g_c: 'float16[m, N]'as type annotation.
Metadata
Metadata
Assignees
Labels
No labels