sketch-design

UnifiedSketch 设计

目标与原则

目标

用最小DSL表达算子设计意图，便于LLM理解和Coder实现。

原则

• 极简原语：只有少数核心操作（alloc/load/store/compute/...）
• 统一语法：所有操作都是函数调用风格，无语法差异
• 标准控制流：使用Python for/range语法，不发明新语法
• hint分离：复杂优化用hint表达，不影响主逻辑清晰性

核心语法元素

结构声明

sketch <op_name> {
  symbols: M, N, K;                    # 符号变量声明
  tensors: A[M, K]: f16; B[K, N]: f16; C[M, N]: f32;  # 张量声明
  constexpr: m0, k0, n0
}

@llm_hint 装饰器详解

基本语法

@llm_hint 用于给 LLM 提供优化提示，帮助 coder 选择最优的实现策略。

@llm_hint("optimization_type")              # 单一提示
@llm_hint("optimization_type", "context")   # 带上下文的提示
@llm_hint("opt1", "opt2", "opt3")          # 多重提示

优化类型

• "parallel" - 并行化此循环
• "pipeline" - 流水线优化
• "vectorize" - 向量化
• "unroll" - 循环展开

硬件上下文提示

• "grididx" - GPU grid 级别并行（对应 blockIdx）
• "threadidx" - GPU thread 级别并行（对应 threadIdx）
• "coreidx" - NPU core 级别并行
• "warp" - GPU warp 级别优化
• "simd" - CPU/NPU SIMD 向量化

for循环表达

# GPU风格：grid + thread 两级并行
@llm_hint("parallel", "grididx.x")
for i in range(0, M, 128):                  # block级别
    @llm_hint("parallel", "threadidx.x") 
    for j in range(0, N, 32):               # thread级别
        @llm_hint("pipeline")
        for k in range(0, K, k_tile):
            # 计算逻辑

# NPU风格：core级别并行
@llm_hint("parallel", "coreidx")
for core_idx in range(num_cores):
    @llm_hint("pipeline")
    for k in range(0, K, k_tile):
        # 每个core的计算

# CPU风格：SIMD向量化
@llm_hint("parallel")                       # OpenMP并行
for i in range(0, M, tile_size):
    @llm_hint("vectorize", "simd")          # SIMD向量化
    for j in range(tile_size):
        # 向量化计算

核心操作

1. alloc - 内存分配
2. load - 数据加载
3. store - 数据存储
4. compute函数 - 计算操作

语法概览

sketch matmul {
  symbols: M, N, K;
  tensors: A[M, K]: f16; B[K, N]: f16; C[M, N]: f32;
  
  m0, k0, n0 = 128, 256, 256

  @llm_hint("parallel")
  for i_outer in range(0, ceil(M, m0)):
    @llm_hint("parallel")
    for j_outer in range(0, ceil(N, n0)):

      # 内存分配
      c_tile = alloc([m0, n0], llm_hint=["accumulator", "init_zero"])
      a_tile = alloc([m0, k0], llm_hint=["fast", "input_cache"])
      b_tile = alloc([k0, n0], llm_hint=["fast", "input_cache"])

      @llm_hint("pipeline")
      for k_outer in range(0, ceil(K, k0)):
        # 数据搬移
        load(A[i_outer:i_outer+m0, k_outer:k_outer+k0] -> a_tile)
        load(B[k_outer:k_outer+k0, j_outer:j_outer+n0] -> b_tile)
        
        # 计算操作
        gemm(a_tile, b_tile, dst=c_tile)
        
      # 数据写回
      store(c_tile -> C[i_outer:i_outer+m0, j_outer:j_outer+n0])
}