Add Torch and Triton kernels for Torch2FlyDSL and Triton2FlyDSL translation

.gitignore

@@ -1,6 +1,7 @@
 .DS_Store
 logs
 __pycache__
+build/
 .vscode
 .idea
 .claude

src/prompt_builder.py

@@ -252,6 +252,10 @@ def prompt_builder(task_config_dir: str, workspace_directory: Path, eval_config:
         task_type_prompt = task_type.instruction2triton_task_type()
     elif task_type_name == 'flydsl2flydsl':
         task_type_prompt = task_type.flydsl2flydsl_task_type()
+    elif task_type_name == 'torch2flydsl':
+        task_type_prompt = task_type.torch2flydsl_task_type()
+    elif task_type_name == 'triton2flydsl':
+        task_type_prompt = task_type.triton2flydsl_task_type()
     elif task_type_name == 'repository':
         task_type_prompt = task_type.repository_task_type()
     else:

src/prompts/task_type.py

@@ -19,5 +19,13 @@ def flydsl2flydsl_task_type() -> str:
     return '''You are a Kernel Optimization Specialist with expertise in FlyDSL (FlyDSL Python DSL) programming for AMD GPUs. Your core mission is to systematically optimize existing FlyDSL kernels for maximum performance while ensuring strict numerical correctness and functional equivalence to the original code. You understand FlyDSL's @flyc.kernel decorator, fx.Tensor buffer APIs, shared-memory reduction patterns, vectorized buffer_load/store copy atoms, and how to leverage ROCm architecture features for optimal throughput on AMD Instinct accelerators.'''
 
 
+def torch2flydsl_task_type() -> str:
+    return '''You are a GPU Kernel Development Specialist with deep expertise in both PyTorch and FlyDSL (FlyDSL Python DSL) programming for AMD GPUs. Your core mission is to translate PyTorch operations and models into highly optimized FlyDSL kernels for AMD Instinct accelerators, while ensuring numerical correctness and functional equivalence to the original PyTorch implementation. You are given a PyTorch reference in KernelBench format (a `class Model(nn.Module)` with `get_inputs()` and `get_init_inputs()`); your job is to implement an equivalent FlyDSL kernel for `Model.forward`. You understand FlyDSL's @flyc.kernel decorator, fx.Tensor buffer APIs, tiling and shared-memory patterns, and how to map high-level PyTorch tensor ops onto ROCm primitives. The target MUST be FlyDSL — do NOT rewrite it in HIP, CUDA, or Triton.'''
+
+
+def triton2flydsl_task_type() -> str:
+    return '''You are a GPU Kernel Translation Specialist with deep expertise in both Triton and FlyDSL (FlyDSL Python DSL) programming for AMD GPUs. Your core mission is to translate an existing Triton kernel into an equivalent, highly optimized FlyDSL kernel for AMD Instinct accelerators, while ensuring numerical correctness and functional equivalence to the original Triton implementation. You are given a STANDALONE Triton source (depends only on `triton`/`torch`) exposing a public entry function plus one or more `@triton.jit` kernels; your job is to implement an equivalent FlyDSL kernel that preserves the public entry function's signature and produces matching outputs. You understand Triton's block-based programming model (program ids, `tl.load`/`tl.store` with masks, `tl.dot`, online-softmax / flash-attention patterns) AND FlyDSL's @flyc.kernel decorator, fx.Tensor buffer APIs, tiling and shared-memory patterns, and how to map Triton block semantics onto ROCm primitives. The target MUST be FlyDSL — do NOT merely re-optimize the Triton kernel and do NOT rewrite it in HIP, CUDA, or plain Triton. Preserve the function signature and the numerical correctness gate of the original.'''
+
+
 def repository_task_type() -> str:
     return '''You are a GPU performance engineer working on Level-3 (repository-scope) tasks. You are given a full checkout of an upstream project—not an isolated snippet. Your job is to explore the real directory layout, build system, tests, and dependencies, then improve the target kernels or hot paths the task describes while preserving correct behavior. The task config selects the language stack (HIP or Triton) for the knowledge section via `repository_language`; follow that stack and the project’s own conventions. The task’s compile, correctness, and performance commands are the source of truth. Prioritize measurable speedups on the target AMD GPU without breaking the project’s validation story.'''

tasks/torch2flydsl/batched_gemm_a8w8_kernel/config.yaml

@@ -0,0 +1,12 @@
+source_file_path:
+- kernel.py
+target_kernel_functions:
+- flydsl_batched_gemm_a8w8
+compile_command:
+- python3 -c "import model, kernel; print('compile ok')"
+correctness_command:
+- python3 test_kernel_harness.py --correctness
+performance_command:
+- python3 test_kernel_harness.py --full-benchmark
+task_type: torch2flydsl
+task_result_template: null

tasks/torch2flydsl/batched_gemm_a8w8_kernel/kernel.py

@@ -0,0 +1,11 @@
+# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved.
+"""FlyDSL starter for this torch2flydsl task.
+
+Implement the target entry points below using FlyDSL. ``model.py`` contains the
+PyTorch reference/specification and ``test_kernel_harness.py`` contains the
+correctness and performance checks. These stubs intentionally do not call the
+reference implementation, so an unimplemented task cannot pass validation.
+"""
+
+def flydsl_batched_gemm_a8w8(*args, **kwargs):
+    raise NotImplementedError("Implement flydsl_batched_gemm_a8w8 using FlyDSL for the batched_gemm_a8w8_kernel task.")

tasks/torch2flydsl/batched_gemm_a8w8_kernel/model.py

@@ -0,0 +1,83 @@
+# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved.
+"""Pure-PyTorch reference for the batched INT8 GEMM ``batched_gemm_a8w8``.
+
+Computes a per-batch ``out[b] = (x[b] @ w[b].T) * (x_scale[b] @ w_scale[b])``
+where the activation ``x`` (``[B, M, K]``) is quantized to INT8 with a per-token
+scale and the weight ``w`` (``[B, N, K]``) is quantized to INT8 with a
+per-output-channel scale, matching the AMD runtime batched a8w8 GEMM:
+
+- activation scale ``x_scale`` is ``[B, M, 1]`` (one per row of each batch);
+- weight scale ``w_scale`` is ``[B, 1, N]`` (one per output channel of each
+  batch);
+- the GEMM forms the INT8 product, accumulates in fp32, multiplies by the
+  outer-product of the per-token and per-channel scales, and truncates to bf16.
+
+``quantize_batched_a8w8`` is the single source of the INT8 operands and scales
+(the harness reuses it to drive the real AMD runtime op), so the reference and
+the hardware kernel operate on byte-identical inputs and differ only by the
+accumulation precision (fp32 reference vs int32 hardware accumulate).
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# INT8 saturation magnitude used as the per-row scale divisor.
+_INT8_MAX = 127.0
+
+
+def _quantize_rowwise_i8(t):
+    """Per-row INT8 quantization of a ``[B, R, K]`` tensor.
+
+    Returns the INT8 codes and the fp32 row scales (``[B, R, 1]``) such that
+    ``t ~= codes.float() * scale``."""
+    tf = t.float()
+    amax = tf.abs().amax(dim=-1, keepdim=True)
+    scale = (amax / _INT8_MAX).clamp_min(1e-12)
+    q = (tf / scale).round().clamp_(-_INT8_MAX, _INT8_MAX).to(torch.int8)
+    return q, scale
+
+
+def quantize_batched_a8w8(x, w):
+    """Quantize a high-precision batched activation/weight pair to the INT8
+    operands the deployed batched GEMM consumes.
+
+    Returns ``(x_i8, x_scale, w_i8, w_scale)`` where ``x_scale`` is ``[B, M, 1]``
+    and ``w_scale`` is ``[B, 1, N]`` (the layouts the AMD runtime op expects)."""
+    x_i8, x_scale = _quantize_rowwise_i8(x)
+    w_i8, w_scale_rows = _quantize_rowwise_i8(w)
+    w_scale = w_scale_rows.transpose(1, 2)
+    return x_i8, x_scale, w_i8, w_scale
+
+
+def _dequant_batched_matmul(x_i8, x_scale, w_i8, w_scale):
+    """INT8 batched GEMM with fp32 accumulation and per-token/per-channel scale."""
+    b, m, _ = x_i8.shape
+    n = w_i8.shape[1]
+    out = torch.empty(b, m, n, dtype=torch.float32, device=x_i8.device)
+    for i in range(b):
+        prod = F.linear(x_i8[i].float(), w_i8[i].float())
+        out[i] = prod * torch.matmul(x_scale[i], w_scale[i])
+    return out
+
+
+class Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, w):
+        x_i8, x_scale, w_i8, w_scale = quantize_batched_a8w8(x, w)
+        out = _dequant_batched_matmul(x_i8, x_scale, w_i8, w_scale)
+        return out.to(torch.bfloat16)
+
+
+def get_inputs():
+    # Representative batched a8w8 shape (B, M, N, K) = (16, 128, 1280, 8192); the
+    # harness sweeps more real shapes from configs/a8w8_untuned_batched_gemm.csv.
+    b, m, n, k = 16, 128, 1280, 8192
+    x = torch.randn(b, m, k, dtype=torch.bfloat16)
+    w = torch.randn(b, n, k, dtype=torch.bfloat16)
+    return [x, w]
+
+
+def get_init_inputs():
+    return []