evofabric.app.kernel_evolve#

class evofabric.app.kernel_evolve.BaseEvaluator#

内核演化过程的抽象评估器

evaluate(initial_code: str, evolve_code: str) → Metrics#

执行生成代码的评估

参数:

initial_code (str) -- 内核的初始代码
evolve_code (str) -- 内核的演化代码

返回:

生成内核代码的评估指标

返回类型:

Metrics

Raises:

NotImplementedError

class evofabric.app.kernel_evolve.GPUEvaluator#

基于GPU设备实现的内核演化评估器

evaluate(initial_code: str, evolve_code: str) → Metrics#

执行生成代码的评估

参数:

initial_code (str) -- 内核的初始代码
evolve_code (str) -- 内核的演化代码

返回:

生成内核代码的评估指标

返回类型:

Metrics

class evofabric.app.kernel_evolve.Metrics(speedup: float, original_time: float, optimized_time: float, error: str | None, traceback: str | None)#

内核演化的指标类

参数:

speedup (float) -- 内核重写后的加速比
original_time (float) -- 初始代码的执行时间
optimized_time (float) -- 优化后生成代码的执行时间
error (str) -- 评估过程中产生的错误
traceback (str) -- 评估过程的错误追踪信息

class evofabric.app.kernel_evolve.LLMConfig(model_class: str, model_name: str, api_key: str, base_url: str, **kwargs)#

内核演化的LLM配置类

参数:

model_class (str) -- 模型类名
model_name (str) -- 模型名称
api_key (str) -- 模型的API密钥
base_url (str) -- 模型的基础URL
**kwargs -- 任意关键字参数

示例用法：

from evofabric.app.kernel_evolve import LLMConfig

LLMConfig(
    model_class="OpenAIChatClient",
    model_name='your-model-name',
    api_key="xxxx",
    base_url="xxxx",
)

class evofabric.app.kernel_evolve.KernelEvolve(initial_code: str, llm_config: LLMConfig, evaluator: BaseEvaluator)#

内核演化控制器，用于启动内核演化

参数:

initial_code (str) -- 待演化的操作代码
llm_config (LLMConfig) -- 内核演化的LLM配置类
evaluator (BaseEvaluator) -- 用于生成内核代码的评估器实现

evolve()#

执行自我演化并返回生成的内核代码

返回:

flag (boolean): 执行成功的标志
result (str): 生成的内核代码或错误信息

返回类型:

tuple

示例用法：

from evofabric.app.kernel_evolve import BaseEvaluator, LLMConfig
from evofabric.app.kernel_evolve.core.controller import KernelEvolve
original_code = '''
        import torch
        import torch.nn as nn

        class Model(nn.Module):
            """
            calculate C = diag(A) * B + D
            A: (N,)
            B: (N, M)
            D: (N, M)
            C: (N, M)
            """

            def __init__(self, BLOCK_M=128):
                super(Model, self).__init__()
                self.BLOCK_M = BLOCK_M

            def forward(self, A, B, D):
                return torch.diag(A) @ B + D


        def get_inputs():
            N, M = 4096, 4096
            A = torch.randn(N, dtype=torch.float32)
            B = torch.randn(N, M, dtype=torch.float32)
            D = torch.randn(N, M, dtype=torch.float32)
            return [A, B, D]


        def get_init_inputs():
            return []
        '''
config = LLMConfig(
    model_class="PanguClient",
    model_name='Pangu_38b',
    api_key="xxxx",
    base_url="xxxx",
    default_headers={"csb-token": "xxxx"}
)
evaluator = GPUKernelEvaluator()
kernel_evolve = KernelEvolve(
    initial_code=original_code,
    llm_config=config,
    evaluator=evaluator)
success_flag, result = kernel_evolve.evolve()

evofabric.app.kernel_evolve

目录

evofabric.app.kernel_evolve#