evofabric.app.kernel_evolve

evofabric.app.kernel_evolve#

class evofabric.app.kernel_evolve.BaseEvaluator#

Abstract Evaluator of the Kernel Evolution Process

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

Execute the evaluation of generated code

Parameters:

  • initial_code (str) – Initial kernel code

  • evolve_code (str) – Kernel evolution code

Returns:

Evaluation metrics for generating kernel code

Return type:

Metrics

Raises:

NotImplementedError

class evofabric.app.kernel_evolve.GPUEvaluator#

GPU-based Kernel Evolution Evaluator

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

Execute the evaluation of generated code

Parameters:

  • initial_code (str) – Initial kernel code

  • evolve_code (str) – Kernel evolution code

Returns:

Evaluation metrics for generating kernel code

Return type:

Metrics

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

Kernel Evolution Indicator Class

Parameters:

  • speedup (float) – Speedup Ratio After Kernel Rewrite

  • original_time (float) – Initial code execution time

  • optimized_time (float) – Execution time of the optimized generated code

  • error (str) – Errors generated during the evaluation process

  • traceback (str) – Evaluation process error tracking information

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

Kernel Evolution LLM Configuration Class

Parameters:

  • model_class (str) – Model class name

  • model_name (str) – Model Name

  • api_key (str) – Model API Key

  • base_url (str) – Model’s base URL

  • **kwargs – Arbitrary keyword arguments

Example Usage:

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)#

Kernel Evolution Controller, for initiating kernel evolution

Parameters:

  • initial_code (str) – Operation Code Pending Evolution

  • llm_config (LLMConfig) – Kernel Evolution LLM Configuration Class

  • evaluator (BaseEvaluator) – Evaluator implementation for generating kernel code

evolve()#

Execute self-evolution and return the generated kernel code

Returns:

  • flag (boolean): execution success flag

  • result (str): generated kernel code or error message

Return type:

tuple

Example Usage:

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()
Contents