[Requirement|需求建议]: CATLASS DSL 支持以类实例组织 Kernel-catlass-AtomGit

已关闭

[Requirement|需求建议]: CATLASS DSL 支持以类实例组织 Kernel #559

rxtfeng创建于 9月12日关闭于 5 天前

9月12日创建

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

Backgroud / 背景信息 (Mandatory / 必填)

背景

当前 CATLASS DSL 主要以模块级 @tla.kernel 函数作为编译入口。对于包含较多编译期配置
的复杂算子，kernel、配置校验、fake tensor 构造和编译缓存往往分散在不同的模块级函数或
全局变量中，难以表达“一个配置实例对应一个 kernel specialization”的关系。

CuteDSL 支持将编译期配置保存在实例中，由 @jit __call__ 选择
实例 @kernel，再通过 @cache @staticmethod compile 创建并缓存编译结果。这种方式让
配置、kernel 和编译入口形成清晰的单一抽象。

当前如下按照cutedsl的组织方式会导致编译报错

class VectorAdd:
    def __init__(self, hidden_size: int):
        self.hidden_size = hidden_size

    @tla.jit
    def __call__(self, lhs, rhs, output):
        self.kernel(lhs, rhs, output)

    @tla.kernel
    def kernel(self, lhs, rhs, output):
        seqlen = lhs.origin_shape[0]
        for token_index in tla.range(0, seqlen, 1):
            for hidden_index in tla.range(0, self.hidden_size, 1):
                output[token_index, hidden_index] = (
                    lhs[token_index, hidden_index]
                    + rhs[token_index, hidden_index]
                )

    @cache
    @staticmethod
    def compile(seqlen: int, hidden_size: int):
        fake = make_fake_tensor((seqlen, hidden_size))
        kernel = VectorAdd(hidden_size)
        return tla.compile(kernel, fake, fake, fake, options="--npu-arch 3510")


compiled = VectorAdd.compile(seqlen=4, hidden_size=16)
compiled(lhs, rhs, output, block_num=1)