Skip to content

[Feature] Support L1 calling convention in tvm-ffi #54

@yaoyaoding

Description

@yaoyaoding

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

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions