diff --git a/skills/triton/matmul-related-gen/k_axis_offset/skill.md b/skills/triton/matmul-related-gen/k_axis_offset/skill.md new file mode 100644 index 00000000..2b41d56f --- /dev/null +++ b/skills/triton/matmul-related-gen/k_axis_offset/skill.md @@ -0,0 +1,448 @@ +# K-Axis Offset Optimization Algorithm for GEMM + +## Inputs + +- `input_triton_code` - Original Triton kernel code containing matrix multiplication + - Must contain `tl.dot()` operation (indicates matrix multiplication) + - Should use `tl.make_block_ptr()` for pointer management + - Should have K-axis reduction loop + +## Outputs + +- `triton_code_with_koffset` - Optimized Triton code with K-axis offset technique applied + - All K positions are computed exactly once + - Uses cyclic K-axis traversal + +## Overview + +K-Axis Offset Optimization is a technique for optimizing GEMM (General Matrix Multiply) operations on NPU by distributing K-axis computation across different blocks with different starting positions. This reduces memory access conflicts and improves parallelism. + +## Preconditions + +Before applying this optimization, verify the following conditions: + +1. **Contains Matrix Multiplication**: The code must contain `tl.dot()` operation +2. **K-Axis Reduction Loop**: There should be a loop that iterates over the K dimension +3. **Block Pointer Usage**: Code should use `tl.make_block_ptr()` for pointer management +4. **Sufficient K Dimension**: K should be significantly larger than BLOCK_SIZE_K for meaningful benefit +5. **Multi-Core Target**: Target hardware should have multiple cores for parallelism benefit + +## Code Pattern Recognition + +### Pattern to Match + +Look for the following code patterns that indicate applicability of K-axis offset optimization: + +```python +# Standard K-axis loop pattern (BEFORE optimization) +for k in range(0, K, BLOCK_SIZE_K): + a_val = tl.load(a_ptr, boundary_check=(0, 1)) + b_val = tl.load(b_ptr, boundary_check=(0, 1)) + acc += tl.dot(a_val, b_val, out_dtype=tl.int32) + a_ptr = tl.advance(a_ptr, (0, BLOCK_SIZE_K)) + b_ptr = tl.advance(b_ptr, (BLOCK_SIZE_K, 0)) +``` + +### Key Indicators + +1. **`tl.dot()` presence**: The code contains `tl.dot()` operation +2. **K-axis loop**: A loop iterating over K dimension with step size `BLOCK_SIZE_K` +3. **Pointer advancement**: Uses `tl.advance()` to move pointers along K dimension +4. **Accumulator pattern**: Uses an accumulator variable (e.g., `acc`) that accumulates dot products + +### Detection Logic + +```python +def can_apply_koffset(code): + indicators = [ + "tl.dot(" in code, + "tl.advance(" in code, + "tl.make_block_ptr(" in code, + any(keyword in code for keyword in ["BLOCK_SIZE_K", "for k", "range(0, K"]) + ] + return sum(indicators) >= 3 +``` + +## Algorithm Principle + +### Core Idea + +Different blocks start from different K positions and traverse the entire K axis cyclically: + +1. **Calculate K Offset**: Each block calculates its starting K position using modulo arithmetic +2. **Cyclic Traversal**: Traverse from offset to K end, then wrap around to 0 and continue to offset +3. **Complete Coverage**: Ensure all K positions are computed exactly once + +### Mathematical Foundation + +For a block with index `block_idx`: +```python +num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K) +k_offset = (2 * block_idx) % num_k_blocks + +for k_iter in range(num_k_blocks): + k_idx = (k_iter + k_offset) % num_k_blocks + current_k = k_idx * BLOCK_SIZE_K +``` + +The K-axis traversal follows this pattern: +- **Phase 1**: From `k_offset * BLOCK_SIZE_K` to `K` (exclusive) +- **Phase 2**: From `0` to `k_offset * BLOCK_SIZE_K` (exclusive) + +This ensures complete K-axis coverage with cyclic behavior. + +## Implementation + +### Method 1: Fusion Approach with Per-Iteration Pointer Creation (Recommended) + +This approach is based on the reference implementation and provides an optimized fusion pattern: + +```python +num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K) +k_offset = (2 * block_idx) % num_k_blocks + +acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.int32) + +for k_iter in range(num_k_blocks): + k_idx = (k_iter + k_offset) % num_k_blocks + current_k = k_idx * BLOCK_SIZE_K + + # Create new pointers with current offset each iteration + a_ptr = tl.make_block_ptr( + base=a, + shape=(M, K), + strides=(stride_am, stride_ak), + offsets=(block_m * BLOCK_SIZE_M, current_k), + block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K), + order=(1, 0), + ) + + b_ptr = tl.make_block_ptr( + base=b, + shape=(K, N), + strides=(stride_bk, stride_bn), + offsets=(current_k, block_n * BLOCK_SIZE_N), + block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N), + order=(1, 0), + ) + + a_val = tl.load(a_ptr, boundary_check=(0, 1)) + b_val = tl.load(b_ptr, boundary_check=(0, 1)) + acc += tl.dot(a_val, b_val, out_dtype=tl.int32) +``` + +**Pros**: +- Simplest code structure +- No complex pointer management +- Direct fusion pattern +- No conditional logic + +**Cons**: +- Does not use `tl.advance()` (violates checklist requirement) +- More pointer creation overhead +- Potentially less efficient than hybrid approach + +### Method 2: Two-Phase Approach (Alternative) + +```python +num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K) +k_offset = ((2 * block_idx) % num_k_blocks) * BLOCK_SIZE_K + +# Phase 1: From offset to K +a_ptr = tl.make_block_ptr(..., offsets=(..., k_offset), ...) +b_ptr = tl.make_block_ptr(..., offsets=(k_offset, ...), ...) + +num_k_blocks_phase1 = tl.cdiv(K - k_offset, BLOCK_SIZE_K) +for k_iter in range(num_k_blocks_phase1): + a_val = tl.load(a_ptr, boundary_check=(0, 1)) + b_val = tl.load(b_ptr, boundary_check=(0, 1)) + acc += tl.dot(a_val, b_val, out_dtype=tl.int32) + a_ptr = tl.advance(a_ptr, (0, BLOCK_SIZE_K)) + b_ptr = tl.advance(b_ptr, (BLOCK_SIZE_K, 0)) + +# Phase 2: From 0 to offset +a_ptr = tl.make_block_ptr(..., offsets=(..., 0), ...) +b_ptr = tl.make_block_ptr(..., offsets=(0, ...), ...) + +num_k_blocks_phase2 = tl.cdiv(k_offset, BLOCK_SIZE_K) +for k_iter in range(num_k_blocks_phase2): + a_val = tl.load(a_ptr, boundary_check=(0, 1)) + b_val = tl.load(b_ptr, boundary_check=(0, 1)) + acc += tl.dot(a_val, b_val, out_dtype=tl.int32) + a_ptr = tl.advance(a_ptr, (0, BLOCK_SIZE_K)) + b_ptr = tl.advance(b_ptr, (BLOCK_SIZE_K, 0)) +``` + +**Pros**: +- Clear separation of two phases +- Uses `tl.advance()` efficiently (satisfies checklist requirement) + +**Cons**: +- Two separate loops +- More code duplication +- Manual calculation of phase lengths + +## Key Techniques + +### 1. Modulo Arithmetic for Wrap-Around + +```python +num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K) +k_offset = (2 * block_idx) % num_k_blocks + +for k_iter in range(num_k_blocks): + k_idx = (k_iter + k_offset) % num_k_blocks + current_k = k_idx * BLOCK_SIZE_K +``` + +This formula automatically calculates the current K position and handles wrap-around when exceeding K. + +### 2. Hybrid Pointer Management + +- **Primary**: Use `tl.advance()` for forward movement (efficient) +- **Secondary**: Use `tl.make_block_ptr()` for wrap-around (necessary) + +### 3. Conditional Pointer Recreation + +```python +if k_iter > 0 and current_k == 0: + # Recreate pointers at position 0 +``` + +Only recreate pointers when actually wrapping around to avoid overhead. + +## Workflow + +Follow these steps to apply K-axis offset optimization to your Triton code: + +### Step 1: Analyze Input Code + +- Check if the code contains `tl.dot()` operation +- Identify the K-axis loop structure +- Locate `tl.make_block_ptr()` and `tl.advance()` usage +- Extract matrix dimensions (M, N, K) and BLOCK_SIZE_K + +### Step 2: Extract Parameters + +- Extract K dimension from the code +- Identify the block indexing scheme (e.g., `block_idx`, `pid`, `pgid`) +- Find the current K-loop boundaries + +### Step 3: Apply Optimization + +1. **Calculate K Offset**: Add `num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K)` and `k_offset = (2 * block_idx) % num_k_blocks` +2. **Initialize Accumulator**: Create zero-initialized accumulator +3. **Modify Loop Structure**: Change from simple K-loop to cyclic traversal +4. **Implement Chosen Method**: + - **Method 1 (Fusion Approach)**: Create new pointers per iteration with current_k + - **Method 2 (Two-Phase)**: Use two separate loops with tl.advance() +5. **Handle Wrap-Around**: Ensure complete K-axis coverage + +### Method Selection Guide + +- **Method 1 (Fusion Approach)**: Recommended - optimized for fusion with simplest code structure +- **Method 2 (Two-Phase)**: Alternative - clearer separation when needed + +## Use Cases + +### When to Use K-Axis Offset + +1. **Large K Dimension**: When K is significantly larger than BLOCK_SIZE_K +2. **Memory-Bound Operations**: When memory access is the bottleneck +3. **Multi-Core Systems**: When utilizing multiple NPU cores +4. **Regular GEMM**: Standard matrix multiplication without special patterns + +### When NOT to Use + +1. **Small K Dimension**: When K is close to BLOCK_SIZE_K +2. **Irregular Access Patterns**: When K-axis access is already optimized +3. **Single-Core Execution**: No benefit with single core + +## Implementation Checklist + +- [ ] Calculate `k_offset` using `k_offset = (2 * block_idx) % num_k_blocks` +- [ ] Initialize pointers appropriately based on chosen method +- [ ] Implement cyclic K-axis traversal +- [ ] Handle wrap-around correctly +- [ ] Ensure complete K-axis coverage +- [ ] Verify correctness with allclose test + +### Method-Specific Requirements + +#### Method 1 (Fusion Approach) +- [ ] Create new pointers per iteration +- [ ] Use current_k directly for offsets + +#### Method 2 (Two-Phase) +- [ ] Calculate phase lengths manually +- [ ] Use `tl.advance()` for both phases + +## Common Pitfalls + +### 1. Missing Wrap-Around Logic + +**Wrong**: +```python +for k_iter in range(num_k_blocks): + k_idx = (k_iter + k_offset) % num_k_blocks + current_k = k_idx * BLOCK_SIZE_K + # Always advance, even when wrapping around + a_ptr = tl.advance(a_ptr, (0, BLOCK_SIZE_K)) +``` + +**Correct**: +```python +for k_iter in range(num_k_blocks): + k_idx = (k_iter + k_offset) % num_k_blocks + current_k = k_idx * BLOCK_SIZE_K + + if k_iter > 0 and current_k == 0: + # Recreate pointers at 0 + a_ptr = tl.make_block_ptr(..., offsets=(..., 0), ...) + + # Only advance when not wrapping + if k_iter < num_k_blocks - 1: + a_ptr = tl.advance(a_ptr, (0, BLOCK_SIZE_K)) +``` + +### 2. Incorrect Offset Calculation + +**Wrong**: +```python +k_offset = block_idx * BLOCK_SIZE_K # Can exceed K +``` + +**Correct**: +```python +num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K) +k_offset = (2 * block_idx) % num_k_blocks # Always within [0, num_k_blocks) +``` + +### 3. Not Using tl.advance() + +**Wrong** (violates checklist): +```python +for k_iter in range(num_k_blocks): + k_idx = (k_iter + k_offset) % num_k_blocks + current_k = k_idx * BLOCK_SIZE_K + a_ptr = tl.make_block_ptr(..., offsets=(..., current_k), ...) + # Always recreate pointers +``` + +**Correct** (satisfies checklist): +```python +# Use advance for forward movement +a_ptr = tl.advance(a_ptr, (0, BLOCK_SIZE_K)) +# Only recreate when necessary +if k_iter > 0 and current_k == 0: + a_ptr = tl.make_block_ptr(..., offsets=(..., 0), ...) +``` + +## Integration with Existing Code + +### Step-by-Step Integration + +1. **Identify K-loop**: Find the K-axis loop in your GEMM kernel +2. **Calculate Offset**: Add `num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K)` and `k_offset = (2 * block_idx) % num_k_blocks` +3. **Initialize Accumulator**: Create zero-initialized accumulator +4. **Modify Loop Structure**: Change from simple K-loop to cyclic traversal +5. **Implement Fusion Approach**: Create new pointers per iteration with current_k +6. **Test**: Verify correctness and performance + +### Example Integration + +**Before** (standard GEMM): +```python +for k in range(0, K, BLOCK_SIZE_K): + a_val = tl.load(a_ptr, boundary_check=(0, 1)) + b_val = tl.load(b_ptr, boundary_check=(0, 1)) + acc += tl.dot(a_val, b_val, out_dtype=tl.int32) + a_ptr = tl.advance(a_ptr, (0, BLOCK_SIZE_K)) + b_ptr = tl.advance(b_ptr, (BLOCK_SIZE_K, 0)) +``` + +**After** (K-axis offset - Fusion Approach): +```python +num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K) +k_offset = (2 * block_idx) % num_k_blocks +acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.int32) + +for k_iter in range(num_k_blocks): + k_idx = (k_iter + k_offset) % num_k_blocks + current_k = k_idx * BLOCK_SIZE_K + + # Create new pointers with current offset each iteration + a_ptr = tl.make_block_ptr( + base=a, + shape=(M, K), + strides=(stride_am, stride_ak), + offsets=(block_m * BLOCK_SIZE_M, current_k), + block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K), + order=(1, 0), + ) + + b_ptr = tl.make_block_ptr( + base=b, + shape=(K, N), + strides=(stride_bk, stride_bn), + offsets=(current_k, block_n * BLOCK_SIZE_N), + block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N), + order=(1, 0), + ) + + a_val = tl.load(a_ptr, boundary_check=(0, 1)) + b_val = tl.load(b_ptr, boundary_check=(0, 1)) + acc += tl.dot(a_val, b_val, out_dtype=tl.int32) +``` + +## Advanced Topics + +### Tuning Parameters + +1. **BLOCK_SIZE_K**: Larger values reduce loop iterations but may increase register pressure +2. **CORE_NUM**: Should match hardware core count for optimal parallelism +3. **Offset Granularity**: Can adjust offset calculation for different access patterns + +### Performance Results + +Experimental results on M=128, N=4096, K=7168 matrix multiplication: + +| Configuration | K offset Calculation | Task Duration (us) | Performance Improvement | +|--------------|---------------------|-------------------|------------------------| +| Without K offset | Standard K-loop | 69.98 | Baseline | +| With K offset | `(2 * block_idx) % num_k_blocks` | 62.32 | **11% faster** | + +The modulo operation with multiplier 2 in the offset calculation provides better performance by reducing memory access conflicts. + +### Combining with Other Optimizations + +K-axis offset can be combined with: +- Split-K parallelization +- Double buffering +- Pipeline optimization +- Memory coalescing + +### Debugging Tips + +1. **Print current_k**: Verify wrap-around behavior +2. **Check coverage**: Ensure all K positions are visited +3. **Validate offsets**: Confirm pointer positions are correct +4. **Monitor performance**: Compare with baseline implementation + +## References + +- Triton Programming Guide: https://triton-lang.org/main/index.html +- Ascend NPU Documentation: https://ascend.github.io/triton-ascend/ +- GEMM Optimization Techniques: See `reference/operator_examples/` directory + +## Version History + +- **v1.0** (2026-04-14): Initial version with two-phase approach +- **v1.1** (2026-04-14): Unified loop with hybrid pointer management +- **v1.2** (2026-04-14): Added comprehensive documentation and examples +- **v1.3** (2026-04-14): Updated to use fusion approach as recommended method +- **v1.4** (2026-04-25): Updated to use `(2 * block_idx) % num_k_blocks` formula for offset calculation + +## License + +This optimization technique is provided for use in Triton-Ascend projects. Please refer to the project license for usage terms. \ No newline at end of file diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/README.md b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/README.md new file mode 100644 index 00000000..f8516458 --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/README.md @@ -0,0 +1,64 @@ +# Matmul Code Generator + +特别注意,需要环境中事先安装好 `msprof op`(强调不是 `msprof`)。 + +## 使用方法 + +使用方式有 2 种: + +**1. 单shape测试**(单个shape约需20~30分钟,例如 M=128 N=7168 K=4096 性能约为57us): + +直接在终端执行: + +``` +请按照 SKILL.md 生成 M=128 N=7168 K=4096 的 FP16 高性能算子 +或 请按照 SKILL.md 生成 M=128 N=7168 K=4096 的 FP16 ND格式算子 +或 请按照 SKILL.md 生成 M=128 N=7168 K=4096 的 INT8 高性能算子 +或 请按照 SKILL.md 生成 M=128 N=7168 K=4096 的 INT8 ND格式算子 +``` + +**2. 多shape测试**(使用 `shape_list.md` 中的 shape 列表逐一测试): + +直接在终端执行: + +``` +请按照 SKILL.md 生成 FP16 高性能算子,且请使用 shape_list 中的 shape 进行测试 +或 请按照 SKILL.md 生成 FP16 ND格式算子,且请使用 shape_list 中的 shape 进行测试 +或 请按照 SKILL.md 生成 INT8 高性能算子,且请使用 shape_list 中的 shape 进行测试 +或 请按照 SKILL.md 生成 INT8 ND格式算子,且请使用 shape_list 中的 shape 进行测试 +``` + +代码和测试报告会存储在 `output` 文件夹下。 + +## 服务器配置 + +编辑 `script/server_config.json`: + +```json +{ + "ip": "你的服务器IP", + "username": "root", + "auth_method": "key", + "password": "", + "ssh_key_path": "你的密钥文件路径", + "docker_container": "容器名称", + "host_temp_dir": "/tmp/triton_upload", + "docker_working_dir": "/root/MyAICode", + "npu_device": "NPU设备ID" +} +``` + +| 字段 | 说明 | +|------|------| +| `ip` | 服务器 IP 地址 | +| `auth_method` | `"password"` 密码认证 或 `"key"` 密钥认证 | +| `password` | 密码认证时填写 | +| `ssh_key_path` | 密钥认证时填写 .pem 文件路径 | +| `docker_container` | Docker 容器名称 | +| `npu_device` | NPU 设备 ID(0-7) | + +配置完成后运行检测: + +```bash +python script/check_config.py +``` diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/SKILL.md b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/SKILL.md new file mode 100644 index 00000000..cb02c740 --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/SKILL.md @@ -0,0 +1,202 @@ +--- +name: "matmul-code-generator" +description: "Generate high-performance Triton matmul kernels (FP16/INT8, standard ND or ND2NZ input format). Invoke when user wants to create or modify matmul kernel code with various optimization strategies (tiling, koffset, diagonal, a_fuse/w_fuse/aw_fuse)." +--- + +# Matmul 代码生成器 + +特别注意,需要环境中事先安装好 `msprof op`(强调不是 `msprof`)。 + +生成高性能 Triton 矩阵乘法 Kernel,支持 FP16/INT8 两种精度 × 标准 ND / ND2NZ 两种输入格式,含 tiling、koffset、diagonal、a_fuse/w_fuse/aw_fuse 等优化策略。 + +## 🔄 Workflow + +严格按顺序执行,每步成功后再进入下一步,失败则重试。 + +``` +Step 0: CONFIG_CHECK (自动检测) + 🔍 **自动检测**: 检查 script/server_config.json 是否存在且有效 + + 检测命令: + ```bash + python script/check_config.py + ``` + + - ✅ 配置已存在且有效 → 跳过配置,直接进入 Step 1 + - ❌ 配置不存在或无效 → 按下方模板配置后重新检测 + + 测试 SSH 连接 (可选): + ```bash + python script/check_config.py --test-connection + ``` + + 配置模板 (仅首次或配置变更时需要): + + 方式1: 密码认证 (推荐新手使用) + ```json + { + "ip": "你的服务器IP", + "username": "root", + "auth_method": "password", + "password": "你的密码", + "docker_container": "容器名称", + "npu_device": "NPU设备ID" + } + ``` + + 方式2: 密钥认证 (推荐生产环境) + ```json + { + "ip": "你的服务器IP", + "username": "root", + "auth_method": "key", + "password": "", + "ssh_key_path": "密钥文件路径", + "docker_container": "容器名称", + "npu_device": "NPU设备ID" + } + ``` + +Step 1: CLEAN + python script/server_check.py clean + 验证: 服务器 /root/MyAICode/batch_verification/ 已删除 + +Step 2: UPLOAD + START + # 启动上传(verify_code.py 在启动时初始化 script/.upload_status.json 为 {"status": "false"}, + # 全部上传+启动完成后自动设为 {"status": "true"}) + python script/verify_code.py --code-subdir --server-config script/server_config.json --shape "" --npu-device + 或 + python script/verify_code.py --code-subdir --server-config script/server_config.json --shape-list shape_list.md --npu-device + + ⚠️ **Repeatedly execute** the following command until the `status` field in `script/.upload_status.json` becomes `"true"`: + + ``` + python script/check_file_exist.py --code-subdir + ``` + + Once `status` is confirmed as `"true"`, proceed with subsequent startup steps. + + **⚠️ CRITICAL: Do NOT perform any other operations during this polling/monitoring process. Focus exclusively on running the command above in a loop until verification succeeds. Repeat this polling cycle no more than 100 times; stop if the limit is reached.** + +Step 3: VERIFY + python script/server_check.py status + 验证: 输出包含 "python.*batch_verification" 进程 + +--- 直接退出,等待用户自行操作 --- +--- 用户自行操作 --- + 查看进度: python script/server_check.py progress + 下载结果: python script/analyze_results.py → output/ +``` + +## 📁 项目结构 + +``` +├── SKILL.md # 本文件 +├── code_lib.zip # 所有 kernel 代码 (上传到服务器) +├── shape_list.md # Shape 列表模板 +└── script/ + ├── verify_code.py # 主入口: clean + 上传 + 启动测试 + ├── batch_verification.py # 服务器端批量测试 (自动运行) + ├── server_check.py # 查看状态/进度/结果/clean/kill + ├── analyze_results.py # 下载最优代码 + 生成 Excel + ├── check_config.py # 配置检查 + ├── check_file_exist.py # 上传状态轮询 + ├── reorder_excel.py # Excel 排序 + ├── server_config.json # 服务器连接配置 + └── .upload_status.json # 上传状态标记 +``` + +> **注意**: `code_lib.zip` 解压后包含 `code-fp16/`、`code-fp16-fuse/`、`code-int8/`、`code-int8-fuse/` 四个子目录,对应四种 kernel 变体。运行时需通过 `--code-subdir` 指定要测试的目录。 + +## 🔧 配置 + +`script/server_config.json`: +```json +{ + "ip": "你的服务器IP", + "username": "root", + "auth_method": "key", + "password": "", + "ssh_key_path": "你的密钥文件路径", + "docker_container": "你的容器名", + "npu_device": "5" +} +``` + +**认证方式说明:** +- `auth_method`: `"password"` 使用密码认证,`"key"` 使用密钥认证 +- `username`: SSH 用户名,默认为 `"root"` +- `password`: 当 `auth_method="password"` 时填写密码 +- `ssh_key_path`: 当 `auth_method="key"` 时填写密钥文件路径 + +测试 SSH: +- 密码方式: `ssh root@你的服务器IP` +- 密钥方式: `ssh -i 你的密钥文件路径 root@你的服务器IP` + +## 🎯 四个任务 → 参数映射 + +| 用户需求 | `--code-subdir` | 精度 | 格式 | 文件 | +|----------|-----------------|------|------|------| +| 生成 FP16 高性能 matmul | `code-fp16` | FP16 | 伪NZ | 36 | +| 生成 FP16 ND 格式 matmul | `code-fp16-fuse` | FP16 | ND + 融合 | 72 | +| 生成 INT8 高性能 matmul | `code-int8` | INT8 | 伪NZ | 36 | +| 生成 INT8 ND 格式 matmul | `code-int8-fuse` | INT8 | ND + 融合 | 72 | + +## 📋 脚本 I/O 说明 + +### verify_code.py — 主入口 + +| 参数 | 必需 | 取值 | +|------|------|------| +| `--code-subdir` | ✅ | `code-fp16` / `code-fp16-fuse` / `code-int8` / `code-int8-fuse` | +| `--server-config` | ✅ | `script/server_config.json` | +| `--shape` | 二选一 | `"M N K"` 如 `"128 4096 7168"` | +| `--shape-list` | 二选一 | `shape_list.md` | +| `--npu-device` | ❌ | 0-7,默认取 config | +| `--pattern` | ❌ | 默认 `*.py`,如 `"gemm*.py"` | + +输入: code_lib.zip + batch_verification.py + (shape_list.md) +输出: 服务器后台启动 batch_verification → 生成 results/ 目录 + +### server_check.py — 状态/进度 + +```bash +python script/server_check.py clean # 删除服务器整个 batch_verification 目录 +python script/server_check.py status # 查看测试进程是否在运行 +python script/server_check.py progress # 查看日志末尾 + 结果目录文件数 +python script/server_check.py results # 列出已生成的结果目录 +python script/server_check.py summary # 查看 summary.json +python script/server_check.py durations # 按性能排序显示所有 kernel +python script/server_check.py kill # 安全停止当前 NPU 上的测试进程 +``` + +### analyze_results.py — 下载结果 + +输入: 服务器 `/root/MyAICode/batch_verification/results/summary.json` +输出 (到本地 `output/`): +- `best_kernels_report.xlsx` — 每个 shape 最优 kernel 一行 (Shape / M/N/K / Duration / Kernel File / Kernel Name) +- `*.py` — 所有最优 kernel 源代码 + +### batch_verification.py — 服务器端 (无需手动调用) + +由 verify_code.py 自动启动。扫描 codes//*.py,用 msprof 逐个测试,生成 AAAA.json / summary.json。 + +## 📝 Shape 列表格式 + +`shape_list.md` 每行一个 shape,空格或逗号分隔,支持 `#` 注释: + +``` +128 4096 7168 +256,8192,14336 +# LLM 推理 +1 4096 4096 +``` + +## ❓ FAQ + +| 问题 | 解决 | +|------|------| +| SSH 连接失败 | `chmod 600 密钥文件路径` | +| code_lib.zip 不存在 | 确保项目根目录有 `code_lib.zip` | +| 测试太慢 | 减少 shape_list 中的 shape 数量;用 `--pattern` 过滤 kernel | +| 中断测试 | `python script/server_check.py kill` | diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/code_lib.zip b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/code_lib.zip new file mode 100644 index 00000000..02948284 Binary files /dev/null and b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/code_lib.zip differ diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/.upload_status.json b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/.upload_status.json new file mode 100644 index 00000000..a8a6bbe9 --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/.upload_status.json @@ -0,0 +1 @@ +{"status": "true"} \ No newline at end of file diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/analyze_results.py b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/analyze_results.py new file mode 100644 index 00000000..31bf8dbf --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/analyze_results.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +""" +分析测试结果 — 单shape / 多shape 通用 +产出到 output/: + - best_kernels_report.xlsx 每个shape最优kernel汇总 + - *.py 所有最优kernel代码 +""" +import paramiko +import json +import os +import shutil +from openpyxl import Workbook +from openpyxl.styles import Font, PatternFill, Alignment + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +OUTPUT_DIR = os.path.join(os.path.dirname(SCRIPT_DIR), 'output') + +with open(os.path.join(SCRIPT_DIR, 'server_config.json')) as f: + config = json.load(f) + +CONTAINER = config['docker_container'] +RESULTS = '/root/MyAICode/batch_verification/results' +CODES = '/root/MyAICode/batch_verification/codes' + +def ssh_cmd(ssh, cmd): + stdin, stdout, stderr = ssh.exec_command(f'docker exec {CONTAINER} bash -c "{cmd}"', timeout=30) + return stdout.read().decode('utf-8'), stderr.read().decode('utf-8') + +def get_summary(ssh): + out, _ = ssh_cmd(ssh, f'cat {RESULTS}/summary.json 2>/dev/null') + if not out.strip(): + return None + return json.loads(out) + +def get_shape_summary(ssh, shape): + filename = f'summary_M{shape["M"]}_N{shape["N"]}_K{shape["K"]}.json' + out, _ = ssh_cmd(ssh, f'cat {RESULTS}/{filename} 2>/dev/null') + if not out.strip(): + return None + return json.loads(out) + +def download_kernel(ssh, filename): + cmd = f'find {CODES} -name "{filename}" -exec cat {{}} \\;' + out, _ = ssh_cmd(ssh, cmd) + if not out.strip(): + return None + local = os.path.join(OUTPUT_DIR, filename) + with open(local, 'w', encoding='utf-8') as f: + f.write(out) + return local + +def generate_excel(best_kernels): + path = os.path.join(OUTPUT_DIR, 'best_kernels_report.xlsx') + wb = Workbook() + ws = wb.active + ws.title = "Best Kernels" + + headers = ["Shape", "M", "N", "K", "Task Duration (us)", "Passed", "Kernel File", "Kernel Name"] + ws.append(headers) + header_fill = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid") + header_font = Font(bold=True, color="FFFFFF") + for cell in ws[1]: + cell.fill = header_fill + cell.font = header_font + cell.alignment = Alignment(horizontal="center", vertical="center") + + for item in best_kernels: + s = item['shape'] + k = item['kernel'] + ws.append([ + f"M={s['M']}, N={s['N']}, K={s['K']}" if s else "default", + s['M'] if s else "", s['N'] if s else "", s['K'] if s else "", + k['duration'], k['passed'], k['file'], k['kernel'] + ]) + + for col in ws.columns: + max_len = max((len(str(c.value or '')) for c in col), default=0) + ws.column_dimensions[col[0].column_letter].width = min(max_len + 2, 50) + + wb.save(path) + return path + +ssh = paramiko.SSHClient() +ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) +ssh.connect(hostname=config['ip'], username='root', key_filename=config['ssh_key_path'], timeout=30) + +try: + summary = get_summary(ssh) + if not summary: + print('No summary.json found on server') + exit(1) + + if os.path.exists(OUTPUT_DIR): + shutil.rmtree(OUTPUT_DIR) + os.makedirs(OUTPUT_DIR, exist_ok=True) + + total_shapes = summary.get('total_shapes', 1) + all_results = summary.get('results', []) + + best_kernels = [] + + if total_shapes == 1 and all_results: + shape = all_results[0].get('shape') if all_results else None + passed = [r for r in all_results if r.get('passed')] + passed.sort(key=lambda r: r['duration']) + best = passed[0] if passed else None + if best: + best_kernels.append({'shape': shape, 'kernel': best}) + + print("=" * 95) + print("测试结果统计") + print("=" * 95) + total = summary.get('total_tests', len(all_results)) + p = summary.get('total_passed', len(passed)) + f = summary.get('total_failed', total - p) + print(f"总共测试kernel数量: {total}") + print(f"通过测试kernel数量: {p}") + print(f"失败测试kernel数量: {f}") + print(f"通过率: {p/total*100:.2f}%" if total else "") + print() + + print("=" * 95) + print("最优Kernel") + print("=" * 95) + print(f"{'Rank':<6} {'File':<70} {'Duration(us)':<15}") + print("=" * 95) + for i, r in enumerate(passed[:1], 1): + print(f"{i:<6} {r['file']:<70} {r['duration']:<15.2f}") + print("=" * 95) + else: + for shape in summary.get('shapes', []): + if shape is None: + continue + ss = get_shape_summary(ssh, shape) + if not ss: + continue + results = ss.get('results', []) + passed = [r for r in results if r.get('passed')] + if not passed: + continue + passed.sort(key=lambda r: r['duration']) + best_kernels.append({'shape': shape, 'kernel': passed[0]}) + + print("=" * 95) + print("多Shape测试结果分析") + print("=" * 95) + print(f"总Shape数: {len(best_kernels)}") + print(f"总测试数: {summary.get('total_tests','?')}") + print(f"总通过数: {summary.get('total_passed','?')}") + print(f"总失败数: {summary.get('total_failed','?')}") + + if not best_kernels: + print("\n没有找到通过测试的kernel") + exit(0) + + print() + print("=" * 95) + print("各Shape最优Kernel") + print("=" * 95) + for item in best_kernels: + s = item['shape'] + k = item['kernel'] + print(f" M={s['M']}, N={s['N']}, K={s['K']} | {k['file']} | {k['duration']:.2f} us") + + print(f"\n正在从服务器下载最优kernel代码到 {OUTPUT_DIR} ...") + for item in best_kernels: + local = download_kernel(ssh, item['kernel']['file']) + if local: + print(f" ✓ {item['kernel']['file']}") + else: + print(f" ✗ {item['kernel']['file']} (failed)") + + excel = generate_excel(best_kernels) + print(f"\n✓ Excel报告: {excel}") + print("✓ 完成") + +finally: + ssh.close() diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/batch_verification.py b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/batch_verification.py new file mode 100644 index 00000000..5f557ca2 --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/batch_verification.py @@ -0,0 +1,341 @@ +#!/usr/bin/env python3 +""" +Batch verification script - runs on server +Scans multiple kernel code files and executes msprof op for each one +""" + +import os +import subprocess +import json +import re +import ast +import argparse +import glob +import difflib +from datetime import datetime + + +def extract_kernel_name(code_content, filename): + """Extract kernel name from code using AST parsing""" + kernel_names = [] + try: + tree = ast.parse(code_content) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef): + has_triton_jit = False + for decorator in node.decorator_list: + if isinstance(decorator, ast.Attribute): + if decorator.attr == 'jit': + if isinstance(decorator.value, ast.Name) and decorator.value.id == 'triton': + has_triton_jit = True + elif isinstance(decorator, ast.Name) and decorator.id == 'jit': + has_triton_jit = True + elif isinstance(decorator, ast.Call): + if isinstance(decorator.func, ast.Attribute): + if decorator.func.attr == 'jit': + if isinstance(decorator.func.value, ast.Name) and decorator.func.value.id == 'triton': + has_triton_jit = True + if has_triton_jit: + kernel_names.append(node.name) + except SyntaxError as e: + print(f'AST parse error: {e}') + + if not kernel_names: + return None + if len(kernel_names) == 1: + return kernel_names[0] + + base_filename = os.path.splitext(os.path.basename(filename))[0] + best_name = None + best_ratio = -1 + for name in kernel_names: + ratio = difflib.SequenceMatcher(None, base_filename, name).ratio() + if ratio > best_ratio: + best_ratio = ratio + best_name = name + return best_name + + +def replace_shape_in_code(code_content, M, N, K): + """Replace M, N, K values in the main() function""" + patterns = [ + (r'(\n M = )\d+', rf'\g<1>{M}'), + (r'(\n N = )\d+', rf'\g<1>{N}'), + (r'(\n K = )\d+', rf'\g<1>{K}'), + ] + + modified_code = code_content + for pattern, replacement in patterns: + modified_code = re.sub(pattern, replacement, modified_code) + + return modified_code + + +def run_single_kernel(code_file, work_dir, output_dir, npu_device, shape=None): + """Run verification for a single kernel code file""" + filename = os.path.basename(code_file) + kernel_name = os.path.splitext(filename)[0] + + if shape: + M, N, K = shape + result_dirname = f"{kernel_name}_M{M}_N{N}_K{K}" + else: + result_dirname = kernel_name + + result_dir = os.path.join(output_dir, result_dirname) + os.makedirs(result_dir, exist_ok=True) + + stdout_file = os.path.join(result_dir, 'stdout.txt') + stderr_file = os.path.join(result_dir, 'stderr.txt') + aaaa_file = os.path.join(result_dir, 'AAAA.json') + + with open(code_file, 'r') as f: + original_code = f.read() + + extracted_kernel = extract_kernel_name(original_code, code_file) + if extracted_kernel: + kernel_name = extracted_kernel + + code_modified = original_code + + if shape: + M, N, K = shape + code_modified = replace_shape_in_code(code_modified, M, N, K) + + device_setting = f"import torch\ntorch.npu.set_device({npu_device})\n" + code_with_device = code_modified.replace('import torch\n', device_setting, 1) + + temp_code_file = os.path.join(work_dir, f'temp_{filename}') + with open(temp_code_file, 'w') as f: + f.write(code_with_device) + + start_time = datetime.now().strftime('%H:%M:%S') + + with open(aaaa_file, 'w') as f: + aaaa_data = {'status': 'in_progress', 'passed': None, 'start_time': start_time} + if shape: + aaaa_data['shape'] = {'M': shape[0], 'N': shape[1], 'K': shape[2]} + json.dump(aaaa_data, f) + + env = os.environ.copy() + env['ASCEND_VISIBLE_DEVICES'] = npu_device + + cmd = f'PS1=dummy && source /root/.bashrc && export ASCEND_VISIBLE_DEVICES={npu_device} && msprof op --output={result_dir} --kernel-name={kernel_name} python {temp_code_file}' + + shape_str = f" (M={shape[0]}, N={shape[1]}, K={shape[2]})" if shape else "" + print(f'[{start_time}] Running: {filename}{shape_str} (kernel: {kernel_name})') + result = subprocess.run(cmd, shell=True, capture_output=True, text=True, executable='/bin/bash', env=env) + + with open(stdout_file, 'w') as f: + f.write(result.stdout) + with open(stderr_file, 'w') as f: + f.write(result.stderr) + + stdout_content = result.stdout + stderr_content = result.stderr + + passed = 'Test passed!' in stdout_content + + task_duration_match = re.search(r'Task Duration\(us\):\s*([0-9.]+)', stdout_content) + task_duration = float(task_duration_match.group(1)) if task_duration_match else 999999 + + end_time = datetime.now().strftime('%H:%M:%S') + + data = { + 'status': 'completed', + 'passed': passed, + 'start_time': start_time, + 'end_time': end_time, + 'task_duration': task_duration, + 'kernel_name': kernel_name, + 'code_file': filename, + 'stdout': stdout_content, + 'stderr': stderr_content + } + + if shape: + data['shape'] = {'M': shape[0], 'N': shape[1], 'K': shape[2]} + + with open(aaaa_file, 'w') as f: + json.dump(data, f, indent=2) + + status = 'PASSED' if passed else 'FAILED' + print(f'[{end_time}] {status}: {filename}{shape_str} (duration: {task_duration} us)') + + os.remove(temp_code_file) + + return passed, task_duration, kernel_name + + +def parse_shape(shape_str): + """Parse shape string like '128 4096 7168' (or '128,4096,7168') to (M, N, K)""" + if not shape_str: + return None + # Try splitting by space first, then comma for backward compatibility + parts = shape_str.split() + if len(parts) != 3: + parts = shape_str.split(',') + if len(parts) != 3: + raise ValueError(f"Invalid shape format: {shape_str}. Expected 'M N K' or 'M,N,K'") + return (int(parts[0]), int(parts[1]), int(parts[2])) + + +def parse_shape_list_file(file_path): + """Parse shape list from a file, returns list of (M, N, K) tuples""" + shapes = [] + with open(file_path, 'r') as f: + for line in f: + line = line.strip() + # Skip empty lines and comments + if not line or line.startswith('#') or line.startswith('```'): + continue + try: + shape = parse_shape(line) + shapes.append(shape) + except: + continue + return shapes + + +def main(): + parser = argparse.ArgumentParser(description='Batch verification script for multiple kernel codes') + parser.add_argument('--code-dir', required=True, help='Directory containing all kernel code subdirectories') + parser.add_argument('--code-subdir', required=True, help='Which subdirectory to test (e.g., code-fp16, code-fp16-fuse)') + parser.add_argument('--work-dir', default='/root/MyAICode/batch_verification', help='Working directory') + parser.add_argument('--output-dir', default='/root/MyAICode/batch_verification/results', help='Output directory') + parser.add_argument('--npu-device', default='0', help='NPU device ID') + parser.add_argument('--pattern', default='*.py', help='File pattern to match (default: *.py)') + parser.add_argument('--shape', default=None, help='Matrix shape as M N K (e.g., "128 4096 7168")') + parser.add_argument('--shape-list', default=None, help='File containing list of shapes (one per line)') + args = parser.parse_args() + + CODE_DIR = os.path.join(args.code_dir, args.code_subdir) + WORK_DIR = args.work_dir + OUTPUT_DIR = args.output_dir + NPU_DEVICE = args.npu_device + PATTERN = args.pattern + + # Parse shapes: either from --shape or from --shape-list + SHAPES = [] + if args.shape: + SHAPES = [parse_shape(args.shape)] + elif args.shape_list: + SHAPES = parse_shape_list_file(args.shape_list) + + # If no shapes specified, use None (use shape from code) + if not SHAPES: + SHAPES = [None] + + os.makedirs(WORK_DIR, exist_ok=True) + os.makedirs(OUTPUT_DIR, exist_ok=True) + + code_files = sorted(glob.glob(os.path.join(CODE_DIR, PATTERN))) + + if not code_files: + print(f'No code files found in {CODE_DIR} with pattern {PATTERN}') + return + + print(f'Found {len(code_files)} code files in {CODE_DIR}') + print(f'NPU Device: {NPU_DEVICE}') + print(f'Number of shapes to test: {len(SHAPES)}') + for i, shape in enumerate(SHAPES, 1): + if shape: + print(f' Shape {i}: M={shape[0]}, N={shape[1]}, K={shape[2]}') + else: + print(f' Shape {i}: (using shape from code)') + print(f'Output directory: {OUTPUT_DIR}') + print('=' * 60) + + all_results = [] + total_passed = 0 + total_failed = 0 + + for shape_idx, shape in enumerate(SHAPES, 1): + print(f'\n{"=" * 60}') + if shape: + print(f'TESTING SHAPE {shape_idx}/{len(SHAPES)}: M={shape[0]}, N={shape[1]}, K={shape[2]}') + else: + print(f'TESTING SHAPE {shape_idx}/{len(SHAPES)}: (using shape from code)') + print('=' * 60) + + shape_results = [] + shape_passed = 0 + shape_failed = 0 + + for code_idx, code_file in enumerate(code_files, 1): + print(f'\n[{shape_idx}/{len(SHAPES)}][{code_idx}/{len(code_files)}] Processing: {os.path.basename(code_file)}') + passed, duration, kernel_name = run_single_kernel(code_file, WORK_DIR, OUTPUT_DIR, NPU_DEVICE, shape) + + result_entry = { + 'file': os.path.basename(code_file), + 'kernel': kernel_name, + 'passed': passed, + 'duration': duration + } + if shape: + result_entry['shape'] = {'M': shape[0], 'N': shape[1], 'K': shape[2]} + shape_results.append(result_entry) + all_results.append(result_entry) + + if passed: + shape_passed += 1 + total_passed += 1 + else: + shape_failed += 1 + total_failed += 1 + + # Save per-shape summary + if shape: + shape_summary_file = os.path.join(OUTPUT_DIR, f'summary_M{shape[0]}_N{shape[1]}_K{shape[2]}.json') + else: + shape_summary_file = os.path.join(OUTPUT_DIR, 'summary_default.json') + + shape_summary = { + 'total': len(code_files), + 'passed': shape_passed, + 'failed': shape_failed, + 'npu_device': NPU_DEVICE, + 'timestamp': datetime.now().isoformat(), + 'results': shape_results + } + + if shape: + shape_summary['shape'] = {'M': shape[0], 'N': shape[1], 'K': shape[2]} + + with open(shape_summary_file, 'w') as f: + json.dump(shape_summary, f, indent=2) + + print(f'\nShape {shape_idx} summary:') + print(f' Total: {len(code_files)}') + print(f' Passed: {shape_passed}') + print(f' Failed: {shape_failed}') + print(f' Summary saved to: {shape_summary_file}') + + # Save overall summary + summary_file = os.path.join(OUTPUT_DIR, 'summary.json') + summary = { + 'total_shapes': len(SHAPES), + 'shapes': [{'M': s[0], 'N': s[1], 'K': s[2]} if s else None for s in SHAPES], + 'total_tests': len(SHAPES) * len(code_files), + 'total_passed': total_passed, + 'total_failed': total_failed, + 'npu_device': NPU_DEVICE, + 'timestamp': datetime.now().isoformat(), + 'results': all_results + } + + with open(summary_file, 'w') as f: + json.dump(summary, f, indent=2) + + print('\n' + '=' * 60) + print('BATCH VERIFICATION SUMMARY (ALL SHAPES)') + print('=' * 60) + print(f'Total shapes: {len(SHAPES)}') + print(f'Total tests: {len(SHAPES) * len(code_files)}') + print(f'Total passed: {total_passed}') + print(f'Total failed: {total_failed}') + print(f'Overall summary saved to: {summary_file}') + + +if __name__ == '__main__': + main() diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/check_config.py b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/check_config.py new file mode 100644 index 00000000..45ac7dad --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/check_config.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +""" +配置检查脚本 - 检测服务器配置文件是否存在且有效 + +使用方法: +python script/check_config.py # 检查配置文件 +python script/check_config.py --test-connection # 测试 SSH 连接 +""" + +import json +import sys +import os +import argparse +from pathlib import Path + +REQUIRED_FIELDS = ['ip', 'username', 'auth_method', 'docker_container', 'npu_device'] + +def check_config_exists(config_path): + """检查配置文件是否存在""" + return os.path.exists(config_path) + +def validate_config(config): + """验证配置文件内容""" + errors = [] + + for field in REQUIRED_FIELDS: + if field not in config: + errors.append(f"缺少必填字段: {field}") + + if 'auth_method' in config: + auth_method = config['auth_method'] + if auth_method == 'password': + if not config.get('password'): + errors.append("密码认证方式需要填写 'password' 字段") + elif auth_method == 'key': + if not config.get('ssh_key_path'): + errors.append("密钥认证方式需要填写 'ssh_key_path' 字段") + elif not os.path.exists(config['ssh_key_path']): + errors.append(f"密钥文件不存在: {config['ssh_key_path']}") + else: + errors.append(f"不支持的认证方式: {auth_method} (应为 'password' 或 'key')") + + return errors + +def test_connection(config): + """测试 SSH 连接""" + try: + import paramiko + except ImportError: + print("❌ 错误: 未安装 paramiko 库,请运行: pip install paramiko") + return False + + ssh = paramiko.SSHClient() + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + try: + auth_method = config.get('auth_method', 'key') + username = config.get('username', 'root') + + print(f"正在连接 {config['ip']}...") + + if auth_method == 'password': + ssh.connect( + hostname=config['ip'], + username=username, + password=config.get('password'), + port=22, + timeout=15 + ) + else: + ssh.connect( + hostname=config['ip'], + username=username, + key_filename=config['ssh_key_path'], + port=22, + timeout=15 + ) + + print("✅ SSH 连接成功!") + + if config.get('docker_container'): + stdin, stdout, stderr = ssh.exec_command( + f"docker inspect {config['docker_container']}", + timeout=10 + ) + if stdout.read(): + print(f"✅ Docker 容器 '{config['docker_container']}' 存在") + else: + print(f"⚠️ 警告: Docker 容器 '{config['docker_container']}' 不存在或无法访问") + + ssh.close() + return True + + except paramiko.AuthenticationException: + print("❌ 认证失败: 请检查用户名和密码/密钥") + return False + except paramiko.SSHException as e: + print(f"❌ SSH 连接错误: {e}") + return False + except Exception as e: + print(f"❌ 连接失败: {e}") + return False + +def print_config_template(): + """打印配置模板""" + print("\n📋 配置模板 (保存到 script/server_config.json):") + print(""" +{ + "ip": "你的服务器IP", + "username": "root", + "auth_method": "password", + "password": "你的密码", + "docker_container": "容器名称", + "npu_device": "NPU设备ID" +} + +或使用密钥认证: +{ + "ip": "你的服务器IP", + "username": "root", + "auth_method": "key", + "password": "", + "ssh_key_path": "密钥文件路径", + "docker_container": "容器名称", + "npu_device": "NPU设备ID" +} +""") + +def main(): + parser = argparse.ArgumentParser(description='检查服务器配置文件') + parser.add_argument('--test-connection', action='store_true', + help='测试 SSH 连接') + parser.add_argument('--config', default='script/server_config.json', + help='配置文件路径 (默认: script/server_config.json)') + args = parser.parse_args() + + config_path = args.config + + print("=" * 50) + print("🔍 配置检查") + print("=" * 50) + + if not check_config_exists(config_path): + print(f"❌ 配置文件不存在: {config_path}") + print_config_template() + sys.exit(1) + + print(f"✅ 配置文件存在: {config_path}") + + try: + with open(config_path, 'r', encoding='utf-8') as f: + config = json.load(f) + except json.JSONDecodeError as e: + print(f"❌ 配置文件格式错误: {e}") + sys.exit(1) + + errors = validate_config(config) + if errors: + print("\n❌ 配置验证失败:") + for error in errors: + print(f" - {error}") + print_config_template() + sys.exit(1) + + print("✅ 配置验证通过") + print(f" - 服务器: {config['ip']}") + print(f" - 用户: {config['username']}") + print(f" - 认证方式: {config['auth_method']}") + print(f" - Docker 容器: {config['docker_container']}") + print(f" - NPU 设备: {config['npu_device']}") + + if args.test_connection: + print("\n" + "=" * 50) + print("🔌 测试连接") + print("=" * 50) + if not test_connection(config): + sys.exit(1) + + print("\n✅ 所有检查通过,可以开始使用!") + +if __name__ == '__main__': + main() diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/check_file_exist.py b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/check_file_exist.py new file mode 100644 index 00000000..f5610f92 --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/check_file_exist.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +import paramiko +import json +import os +import sys +import argparse + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +STATUS_FILE = os.path.join(PROJECT_ROOT, 'script', '.upload_status.json') + +WORK_DIR = '/root/MyAICode/batch_verification' +CODE_REMOTE_DIR = WORK_DIR + '/codes' +SCRIPT_PATH = WORK_DIR + '/batch_verification.py' + +CODE_SUBDIRS = ['code-fp16', 'code-fp16-fuse', 'code-int8', 'code-int8-fuse'] + + +def create_ssh_client(config): + ssh = paramiko.SSHClient() + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + ssh.connect( + hostname=config['ip'], + username='root', + key_filename=config['ssh_key_path'], + port=22, + timeout=15 + ) + return ssh + + +def run_command(ssh, container, command): + full_cmd = f'docker exec {container} bash -c "{command}"' + stdin, stdout, stderr = ssh.exec_command(full_cmd, timeout=30) + return stdout.read().decode('utf-8'), stderr.read().decode('utf-8') + + +def write_status(value): + os.makedirs(os.path.dirname(STATUS_FILE), exist_ok=True) + with open(STATUS_FILE, 'w', encoding='utf-8') as f: + json.dump({"status": value}, f) + + +def main(): + parser = argparse.ArgumentParser(description='Check if code files and batch script exist on server') + parser.add_argument('--code-subdir', required=True, choices=CODE_SUBDIRS, + help='Code subdirectory to check') + parser.add_argument('--server-config', default=os.path.join(PROJECT_ROOT, 'script', 'server_config.json'), + help='Path to server_config.json') + args = parser.parse_args() + + code_subdir = args.code_subdir + code_subdir_path = f'{CODE_REMOTE_DIR}/{code_subdir}' + + config_path = args.server_config + try: + with open(config_path, 'r', encoding='utf-8') as f: + config = json.load(f) + except FileNotFoundError: + print(f'check_file_exist: config not found -> status=false') + write_status("false") + sys.exit(1) + + container = config['docker_container'] + ssh = create_ssh_client(config) + + try: + py_count_output, _ = run_command(ssh, container, + f'ls {code_subdir_path}/*.py 2>/dev/null | wc -l') + py_count = int(py_count_output.strip()) if py_count_output.strip().isdigit() else 0 + + script_output, _ = run_command(ssh, container, + f'test -f {SCRIPT_PATH} && echo EXISTS') + script_exists = 'EXISTS' in script_output + + if py_count >= 1 and script_exists: + print(f'check_file_exist: {code_subdir}({py_count} .py files) + batch_verification.py -> status=true') + write_status("true") + else: + print(f'check_file_exist: {code_subdir}={py_count} .py scripts={"YES" if script_exists else "NO"} -> status=false') + write_status("false") + + except Exception as e: + print(f'check_file_exist: error -> status=false ({e})') + write_status("false") + finally: + ssh.close() + + +if __name__ == '__main__': + main() diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/reorder_excel.py b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/reorder_excel.py new file mode 100644 index 00000000..fab91572 --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/reorder_excel.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +import os +from openpyxl import load_workbook, Workbook +from openpyxl.styles import Font, PatternFill, Alignment + +# Read merged_results_all.xlsx to get shape order +print("Reading merged_results_all.xlsx...") +wb1 = load_workbook(r'D:\AI-Triton\AICoding\matmul_skill_for_any_shape\output\merged_results_all.xlsx') +ws1 = wb1.active + +merged_shapes = [] +for row_idx, row in enumerate(ws1.iter_rows(min_row=2, values_only=True), 2): + if len(row) >= 3: + m, n, k = row[0], row[1], row[2] + merged_shapes.append((m, n, k)) + print(" Shape %d: M=%s, N=%s, K=%s" % (len(merged_shapes), m, n, k)) + +# Read best_kernels_report.xlsx +print("\nReading best_kernels_report.xlsx...") +wb2 = load_workbook(r'D:\AI-Triton\AICoding\matmul_skill_for_any_shape\output\best_kernels_report.xlsx') +ws2 = wb2.active + +best_kernels_data = {} +for row_idx, row in enumerate(ws2.iter_rows(min_row=2, values_only=True), 2): + if len(row) >= 7: + m, n, k = row[1], row[2], row[3] + key = (m, n, k) + best_kernels_data[key] = row + print(" Found: M=%s, N=%s, K=%s" % (m, n, k)) + +# Create new Excel file +print("\nCreating reordered Excel...") +wb_new = Workbook() +ws_new = wb_new.active +ws_new.title = "Best Kernels" + +# Add header +headers = ["Shape", "M", "N", "K", "Task Duration (us)", "Passed", "Kernel File", "Kernel Name"] +ws_new.append(headers) + +# Apply header styling +header_fill = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid") +header_font = Font(bold=True, color="FFFFFF") +for col_idx, cell in enumerate(ws_new[1], 1): + cell.fill = header_fill + cell.font = header_font + cell.alignment = Alignment(horizontal="center", vertical="center") + +# Add data in merged_results_all order +count = 0 +for (m, n, k) in merged_shapes: + key = (m, n, k) + if key in best_kernels_data: + row = best_kernels_data[key] + ws_new.append(row) + count += 1 + print(" Added: M=%s, N=%s, K=%s" % (m, n, k)) + else: + print(" Warning: M=%s, N=%s, K=%s not found in best_kernels_report" % (m, n, k)) + +# Adjust column widths +for col in ws_new.columns: + max_length = 0 + column = col[0].column_letter + for cell in col: + try: + if len(str(cell.value)) > max_length: + max_length = len(str(cell.value)) + except: + pass + adjusted_width = (max_length + 2) + ws_new.column_dimensions[column].width = adjusted_width + +# Save +output_path = r'D:\AI-Triton\AICoding\matmul_skill_for_any_shape\output\best_kernels_report_reordered.xlsx' +wb_new.save(output_path) +print("\nSuccess! Reordered Excel saved to: %s" % output_path) +print("Total shapes added: %d/%d" % (count, len(merged_shapes))) diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/server_check.py b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/server_check.py new file mode 100644 index 00000000..171dde10 --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/server_check.py @@ -0,0 +1,275 @@ +#!/usr/bin/env python3 +""" +通用服务器检查脚本 - 用于检查测试进度、状态和结果 + +使用方法: +python script/server_check.py status # 检查测试状态 +python script/server_check.py progress # 检查测试进度 +python script/server_check.py results # 检查测试结果 +python script/server_check.py summary # 查看汇总结果 +python script/server_check.py clean # 清理服务器 +python script/server_check.py kill # 清除指定NPU上的进程 +""" + +import paramiko +import json +import sys + +def create_ssh_client(config): + ssh = paramiko.SSHClient() + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + auth_method = config.get('auth_method', 'key') + username = config.get('username', 'root') + + if auth_method == 'password': + ssh.connect( + hostname=config['ip'], + username=username, + password=config.get('password'), + port=22, + timeout=15 + ) + else: + ssh.connect( + hostname=config['ip'], + username=username, + key_filename=config['ssh_key_path'], + port=22, + timeout=15 + ) + + return ssh + +def run_command(ssh, container, command): + full_cmd = f'docker exec {container} bash -c "{command}"' + stdin, stdout, stderr = ssh.exec_command(full_cmd, timeout=30) + return stdout.read().decode('utf-8'), stderr.read().decode('utf-8') + +def check_status(config): + """检查服务器上的测试状态""" + ssh = create_ssh_client(config) + try: + output, error = run_command(ssh, config['docker_container'], + "ps aux | grep -E 'python.*batch_verification|msprof' | grep -v grep || echo 'No running verification processes'") + print("=== 运行状态 ===") + print(output) + if error: + print(f"Error: {error}") + finally: + ssh.close() + +def check_progress(config): + """检查测试进度""" + ssh = create_ssh_client(config) + try: + # 检查日志 + output, error = run_command(ssh, config['docker_container'], + 'if [ -f /root/MyAICode/batch_verification/batch_verification.log ]; then tail -30 /root/MyAICode/batch_verification/batch_verification.log; else echo "Log file not found"; fi') + print("=== 测试进度 ===") + print(output) + + # 检查结果目录 + output, error = run_command(ssh, config['docker_container'], + 'ls -la /root/MyAICode/batch_verification/results/ 2>/dev/null | wc -l') + print(f"\n结果目录文件数: {output.strip()}") + + finally: + ssh.close() + +def check_results(config): + """检查测试结果""" + ssh = create_ssh_client(config) + try: + output, error = run_command(ssh, config['docker_container'], + 'ls -la /root/MyAICode/batch_verification/results/ 2>/dev/null | head -20') + print("=== 结果目录 ===") + print(output) + finally: + ssh.close() + +def check_summary(config): + """查看汇总结果""" + ssh = create_ssh_client(config) + try: + output, error = run_command(ssh, config['docker_container'], + 'cat /root/MyAICode/batch_verification/results/summary.json 2>/dev/null || echo "summary.json not found"') + print("=== 汇总结果 ===") + print(output) + finally: + ssh.close() + +def check_all_durations(config): + """查看所有kernel的task_duration(从summary.json读取)""" + import json as json_module + ssh = create_ssh_client(config) + try: + output, error = run_command(ssh, config['docker_container'], + 'cat /root/MyAICode/batch_verification/results/summary.json 2>/dev/null || echo "SUMMARY_NOT_FOUND"') + + if not output or "SUMMARY_NOT_FOUND" in output: + print("summary.json not found") + return + + try: + summary = json_module.loads(output) + except json_module.JSONDecodeError as e: + print(f"Error parsing summary.json: {e}") + return + + durations = [] + for item in summary.get('results', []): + durations.append({ + 'kernel': item.get('kernel', 'unknown'), + 'file': item.get('file', 'unknown'), + 'duration': item.get('duration', 999999), + 'passed': item.get('passed', False), + 'shape': item.get('shape', {}) + }) + + if durations: + print("=== 所有Kernel性能 ===") + print(f"{'Kernel':<60} {'Shape':<25} {'Duration(us)':<15} {'Passed'}") + print("=" * 115) + for item in sorted(durations, key=lambda x: x['duration']): + status = "✓" if item['passed'] else "✗" + shape = item.get('shape', {}) + shape_str = f"M={shape.get('M','?')} N={shape.get('N','?')} K={shape.get('K','?')}" + print(f"{item['kernel']:<60} {shape_str:<25} {item['duration']:<15.2f} {status}") + print("=" * 115) + print(f"Total: {len(durations)} kernels") + else: + print("No results found in summary.json") + + finally: + ssh.close() + +def check_detail(config, kernel_name): + """查看单个kernel的详细结果""" + ssh = create_ssh_client(config) + try: + output, error = run_command(ssh, config['docker_container'], + f'cat /root/MyAICode/batch_verification/results/{kernel_name}/AAAA.json 2>/dev/null || echo "AAAA.json not found"') + print("=== 详细结果 ===") + print(output) + finally: + ssh.close() + +def clean_server(config): + """清理服务器(删除整个batch_verification工作目录)""" + ssh = create_ssh_client(config) + try: + output, error = run_command(ssh, config['docker_container'], + 'rm -rf /root/MyAICode/batch_verification && echo "Cleanup completed: /root/MyAICode/batch_verification removed"') + print("=== 清理结果 ===") + print(output) + if error: + print(f"Error: {error}") + finally: + ssh.close() + +def kill_npu_processes(config): + """清除指定NPU上的相关进程""" + npu_device = config.get('npu_device', 'unknown') + ssh = create_ssh_client(config) + try: + print(f"=== 正在清除 NPU {npu_device} 上的进程 ===") + + # 步骤1: 找到指定NPU的batch_verification主进程PID + find_main_cmd = f"ps aux | grep 'python.*batch_verification.*--npu-device {npu_device}' | grep -v grep | awk '{{print $2}}'" + main_pids_output, _ = run_command(ssh, config['docker_container'], find_main_cmd) + main_pids = [pid.strip() for pid in main_pids_output.split('\n') if pid.strip()] + + if not main_pids: + print(f"未找到 NPU {npu_device} 上的运行进程") + return + + print(f"找到主进程 PID: {', '.join(main_pids)}") + + # 步骤2: 显示将要清除的进程 + print("\n将要清除的进程:") + for pid in main_pids: + show_cmd = f"ps aux | grep -E 'PID|{pid}' | head -5" + output, _ = run_command(ssh, config['docker_container'], show_cmd) + print(output) + + # 步骤3: 找到这些主进程的所有子进程(msprof, msopprof等) + all_pids_to_kill = main_pids.copy() + for main_pid in main_pids: + find_children_cmd = f"pgrep -P {main_pid}" + children_output, _ = run_command(ssh, config['docker_container'], find_children_cmd) + child_pids = [pid.strip() for pid in children_output.split('\n') if pid.strip()] + all_pids_to_kill.extend(child_pids) + + # 递归查找孙子进程 + for child_pid in child_pids: + find_grandchildren_cmd = f"pgrep -P {child_pid}" + grandchildren_output, _ = run_command(ssh, config['docker_container'], find_grandchildren_cmd) + grandchild_pids = [pid.strip() for pid in grandchildren_output.split('\n') if pid.strip()] + all_pids_to_kill.extend(grandchild_pids) + + if len(all_pids_to_kill) > len(main_pids): + print(f"\n包含子进程 PID: {', '.join(all_pids_to_kill[len(main_pids):])}") + + # 步骤4: 执行kill + pids_str = ' '.join(all_pids_to_kill) + kill_cmd = f"kill -9 {pids_str} 2>/dev/null || true" + run_command(ssh, config['docker_container'], kill_cmd) + + # 步骤5: 检查是否清除成功 + import time + time.sleep(1) + + check_cmd = f"ps aux | grep 'python.*batch_verification.*--npu-device {npu_device}' | grep -v grep || echo '已清除'" + output, _ = run_command(ssh, config['docker_container'], check_cmd) + print("\n当前进程状态:") + print(output) + + if '已清除' in output: + print(f"\n✅ NPU {npu_device} 进程清除完成!") + else: + print(f"\n⚠️ 部分进程可能未清除,请手动检查") + + finally: + ssh.close() + +def main(): + if len(sys.argv) < 2: + print(__doc__) + return + + try: + with open('script/server_config.json', 'r', encoding='utf-8') as f: + config = json.load(f) + except FileNotFoundError: + print("Error: script/server_config.json not found") + return + + action = sys.argv[1].lower() + + if action == 'status': + check_status(config) + elif action == 'progress': + check_progress(config) + elif action == 'results': + check_results(config) + elif action == 'summary': + check_summary(config) + elif action == 'clean': + clean_server(config) + elif action == 'kill': + kill_npu_processes(config) + elif action == 'detail': + if len(sys.argv) < 3: + print("Usage: python script/server_check.py detail ") + return + kernel_name = sys.argv[2] + check_detail(config, kernel_name) + elif action == 'durations': + check_all_durations(config) + else: + print(f"Unknown action: {action}") + print(__doc__) + +if __name__ == '__main__': + main() diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/server_config.json b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/server_config.json new file mode 100644 index 00000000..7b769c8e --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/server_config.json @@ -0,0 +1,11 @@ +{ + "ip": "你的服务器IP", + "username": "root", + "auth_method": "key", + "password": "", + "ssh_key_path": "你的密钥文件路径", + "docker_container": "你的容器名称", + "host_temp_dir": "/tmp/triton_upload", + "docker_working_dir": "/root/MyAICode", + "npu_device": "0" +} diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/verify_code.py b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/verify_code.py new file mode 100644 index 00000000..ddcdbd21 --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/script/verify_code.py @@ -0,0 +1,275 @@ +import os +import sys +import time +import argparse +import json +import paramiko +import posixpath +from datetime import datetime + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +UPLOAD_TIMEOUT = 120 +UPLOAD_STATUS_FILE = os.path.join(PROJECT_ROOT, 'script', '.upload_status.json') + + +def create_ssh_client(host, port, user, password=None, key_path=None): + client = paramiko.SSHClient() + client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + if password: + client.connect(host, port=port, username=user, password=password, timeout=30) + elif key_path and os.path.exists(key_path): + client.connect(host, port=port, username=user, key_filename=key_path, timeout=30) + else: + raise ValueError("Either password or key_path must be provided") + + return client + + +def run_ssh_command(client, command, max_retries=3, timeout=300): + for attempt in range(max_retries): + try: + stdin, stdout, stderr = client.exec_command(command, timeout=timeout) + stdout_content = stdout.read().decode('utf-8', errors='replace') + stderr_content = stderr.read().decode('utf-8', errors='replace') + returncode = stdout.channel.recv_exit_status() + return returncode, stdout_content, stderr_content + except Exception as e: + if attempt < max_retries - 1: + time.sleep(2) + else: + return -1, '', str(e) + + +def poll_verify(client, container, check_cmd, label, interval=5): + """Poll until check_cmd succeeds or UPLOAD_TIMEOUT expires""" + deadline = time.time() + UPLOAD_TIMEOUT + while time.time() < deadline: + rc, out, _ = run_ssh_command(client, + f'docker exec {container} bash -c "{check_cmd}"', timeout=15) + if rc == 0 and out.strip(): + print(f' [{label}] verified OK after {int(time.time() - (deadline - UPLOAD_TIMEOUT))}s', flush=True) + return True + time.sleep(interval) + print(f' [{label}] verify FAILED after {UPLOAD_TIMEOUT}s', flush=True) + return False + + +def upload_file_to_docker(client, container, local_path, remote_tmp_name, docker_dest, label): + """Upload a file via SFTP, then poll-verify until it appears or timeout""" + remote_tmp = posixpath.join('/tmp', remote_tmp_name) + + deadline = time.time() + UPLOAD_TIMEOUT + attempt = 0 + while time.time() < deadline: + attempt += 1 + print(f' [{label}] upload attempt {attempt} ...', flush=True) + + sftp = client.open_sftp() + sftp.put(local_path, remote_tmp) + sftp.close() + + rc, _, err = run_ssh_command(client, + f'docker cp {remote_tmp} {container}:{docker_dest}', timeout=60) + run_ssh_command(client, f'rm -f {remote_tmp}', timeout=10) + + if rc != 0: + print(f' [{label}] docker cp failed: {err[:200]}', flush=True) + continue + + print(f' [{label}] docker cp done, polling for {docker_dest} ...', flush=True) + if poll_verify(client, container, + f'test -e {docker_dest} && echo OK', label): + return True + + print(f' [{label}] retrying upload ...', flush=True) + + print(f' [{label}] FATAL after {UPLOAD_TIMEOUT}s', flush=True) + return False + + +def upload_code_lib(client, container, zip_path, remote_dir): + """Upload code_lib.zip, extract, poll-verify .py files exist""" + if not os.path.isfile(zip_path): + print(f'code_lib.zip not found: {zip_path}', flush=True) + return False + + zip_name = os.path.basename(zip_path) + deadline = time.time() + UPLOAD_TIMEOUT + attempt = 0 + + while time.time() < deadline: + attempt += 1 + print(f' [code_lib.zip] upload attempt {attempt} ...', flush=True) + + remote_tmp = posixpath.join('/tmp', zip_name) + sftp = client.open_sftp() + sftp.put(zip_path, remote_tmp) + sftp.close() + + unzip_cmd = ( + f'docker exec {container} mkdir -p {remote_dir} && ' + f'docker cp {remote_tmp} {container}:{remote_dir}/ && ' + f'docker exec {container} bash -c "cd {remote_dir} && unzip -o {zip_name} && rm -f {zip_name}"' + ) + rc, out, err = run_ssh_command(client, unzip_cmd, timeout=120) + run_ssh_command(client, f'rm -f {remote_tmp}', timeout=10) + + if rc != 0: + print(f' [code_lib.zip] unzip error (rc={rc}): {err[:500]}', flush=True) + time.sleep(5) + continue + + print(f' [code_lib.zip] unzip output: {out[:200]}', flush=True) + + print(f' [code_lib.zip] extract done, polling for .py files ...', flush=True) + if poll_verify(client, container, + f'ls {remote_dir}/code-fp16/*.py 2>/dev/null | wc -l', + 'code_lib.zip'): + return True + + print(f' [code_lib.zip] retrying ...', flush=True) + + print(f' [code_lib.zip] FATAL after {UPLOAD_TIMEOUT}s', flush=True) + return False + + +def main(): + os.makedirs(os.path.dirname(UPLOAD_STATUS_FILE), exist_ok=True) + with open(UPLOAD_STATUS_FILE, 'w', encoding='utf-8') as f: + json.dump({"status": "false"}, f) + + parser = argparse.ArgumentParser(description='Generate and test Triton matmul kernels') + parser.add_argument('--code-subdir', required=True, + choices=['code-fp16', 'code-fp16-fuse', 'code-int8', 'code-int8-fuse']) + parser.add_argument('--server-config', required=True) + parser.add_argument('--npu-device', default=None) + parser.add_argument('--pattern', default='*.py') + parser.add_argument('--shape', default=None) + parser.add_argument('--shape-list', default=None) + args = parser.parse_args() + + code_lib_path = os.path.join(PROJECT_ROOT, 'code_lib.zip') + if not os.path.isfile(code_lib_path): + print(f'Error: code_lib.zip not found at {code_lib_path}') + sys.exit(1) + + with open(args.server_config, 'r', encoding='utf-8') as f: + server_config = json.load(f) + + host = server_config['ip'] + username = server_config.get('username', 'root') + auth_method = server_config.get('auth_method', 'key') + key_path = server_config.get('ssh_key_path') + password = server_config.get('password') + container = server_config['docker_container'] + npu_device = args.npu_device or server_config.get('npu_device', '0') + + if not all([host, container]): + print('Error: Missing server config fields') + sys.exit(1) + + work_dir = '/root/MyAICode/batch_verification' + code_remote_dir = posixpath.join(work_dir, 'codes') + output_dir = posixpath.join(work_dir, 'results') + script_path = posixpath.join(work_dir, 'batch_verification.py') + log_file = posixpath.join(work_dir, 'batch_verification.log') + + shape_list_remote = None + if args.shape_list: + if not os.path.isfile(args.shape_list): + print(f'Error: Shape list file not found: {args.shape_list}') + sys.exit(1) + shape_list_remote = posixpath.join(work_dir, 'shape_list.txt') + + print(f'Connecting to server: {host}') + print(f'Auth method: {auth_method}') + print(f'Username: {username}') + print(f'Docker container: {container}') + print(f'Code subdirectory: {args.code_subdir}') + print(f'NPU device: {npu_device}') + if args.shape: + print(f'Shape: {args.shape}') + if args.shape_list: + print(f'Shape list: {args.shape_list}') + print() + + if auth_method == 'password': + client = create_ssh_client(host, 22, username, password=password) + else: + client = create_ssh_client(host, 22, username, key_path=key_path) + + try: + print('=== Step 1: Clean server ===') + run_ssh_command(client, + f'docker exec {container} bash -c "rm -rf {work_dir}"', timeout=30) + run_ssh_command(client, + f'docker exec {container} mkdir -p {work_dir} {code_remote_dir} {output_dir}', + timeout=15) + print('Server cleanup OK') + print() + + print('=== Step 2a: Upload code_lib.zip ===') + if not upload_code_lib(client, container, code_lib_path, code_remote_dir): + print('FATAL: Failed to upload code_lib.zip') + sys.exit(1) + print() + + script_dir = os.path.dirname(os.path.abspath(__file__)) + local_batch = os.path.join(script_dir, 'batch_verification.py') + + print('=== Step 2b: Upload batch_verification.py ===') + if not upload_file_to_docker(client, container, local_batch, + 'batch_verification_tmp.py', script_path, + 'batch_verification.py'): + print('FATAL: Failed to upload batch_verification.py') + sys.exit(1) + run_ssh_command(client, + f'docker exec {container} chmod +x {script_path}', timeout=10) + print() + + if args.shape_list: + print('=== Step 2c: Upload shape_list ===') + if not upload_file_to_docker(client, container, args.shape_list, + 'shape_list_tmp.txt', shape_list_remote, + 'shape_list'): + print('FATAL: Failed to upload shape_list') + sys.exit(1) + print() + + print('=== Step 3: Start batch verification ===') + start_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S') + + shape_arg = f"--shape '{args.shape}'" if args.shape else '' + shape_list_arg = f'--shape-list {shape_list_remote}' if args.shape_list else '' + + exec_cmd = ( + f'docker exec {container} bash -c ' + f'"nohup python3 {script_path} ' + f'--code-dir {code_remote_dir} ' + f'--code-subdir {args.code_subdir} ' + f'--work-dir {work_dir} ' + f'--output-dir {output_dir} ' + f'--npu-device {npu_device} ' + f'--pattern {args.pattern} ' + f'{shape_arg} {shape_list_arg} ' + f'> {log_file} 2>&1 &"' + ) + run_ssh_command(client, exec_cmd, timeout=30) + + print(f'Batch verification started at {start_time}') + print(f'Log: {log_file}') + print(f'Results: {output_dir}') + print() + print('Monitor: python script/server_check.py status') + print('Monitor: python script/server_check.py progress') + + with open(UPLOAD_STATUS_FILE, 'w', encoding='utf-8') as f: + json.dump({"status": "true"}, f) + + finally: + client.close() + + +if __name__ == '__main__': + main() diff --git a/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/shape_list.md b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/shape_list.md new file mode 100644 index 00000000..0ea44b16 --- /dev/null +++ b/skills/triton/matmul-related-gen/matmul_skill_for_any_shape/shape_list.md @@ -0,0 +1,94 @@ +# Shape List for Batch Testing + +这个文件用于批量测试多个shape的矩阵乘法kernel。 + +## 格式说明 + +每行一个shape,格式为:`M N K` + +## Shape 列表 + +``` +128 4096 7168 +1024 4096 7168 +128 7168 2048 +1024 7168 2048 +1 131072 6144 +1024 36 6144 +1024 768 6144 +1024 96 6144 +16384 48 1536 +16384 1024 2048 +16384 6144 2048 +16384 36 6144 +16384 96 6144 +128 8192 6144 +64 384 1024 +64 8192 6144 +8 768 6144 +1 5120 1280 +1 13824 5120 +1 1792 5120 +1 37936 5120 +1 5120 6912 +5 5120 1280 +5 13824 5120 +5 1792 5120 +5 5120 6912 +140236 128 6912 +15 5120 10240 +240 9696 5120 +720 9696 5120 +15 128 5120 +45 128 5120 +17940 1280 1280 +17940 1280 3420 +17940 3840 1280 +17940 6480 1280 +4485 3584 5120 +4485 5120 5120 +17940 1280 1176 +75 152064 3584 +76 152064 3584 +77 152064 3584 +78 152064 3584 +79 152064 3584 +80 152064 3584 +4525 3584 18944 +4525 3584 3584 +4525 37888 3584 +4525 4608 3584 +4826 3584 18944 +4826 3584 3584 +4826 37888 3584 +4826 4608 3584 +15 5120 2048 +15 12800 5120 +15 2560 5120 +15 5120 6400 +18 5120 2048 +18 12800 5120 +18 2560 5120 +18 5120 6400 +19 5120 2048 +19 12800 5120 +19 2560 5120 +19 5120 6400 +21 5120 2048 +21 12800 5120 +21 2560 5120 +21 5120 6400 +3000 5120 2048 +3000 12800 5120 +3000 2560 5120 +3000 5120 6400 +6 37984 5120 +7 37984 5120 +8 37984 5120 +``` + +## 说明 + +- 每行一个shape,使用空格分隔M、N、K +- 支持空行和注释(以#开头) +- 会自动跳过无效行 diff --git a/skills/triton/matmul-related-gen/skills_overview.md b/skills/triton/matmul-related-gen/skills_overview.md new file mode 100644 index 00000000..cc4a31ce --- /dev/null +++ b/skills/triton/matmul-related-gen/skills_overview.md @@ -0,0 +1,175 @@ +# Triton Matmul 相关 Skill 体系说明 + +本目录包含 4 个 Skill,它们构成了一套完整的 Triton 矩阵乘法(GEMM)算子自动生成体系。 + +特别注意,需要环境中事先安装好 `msprof op`(强调不是 `msprof`)。 + +## 架构概览 + +``` +matmul_skill_for_any_shape (终端用户入口 Skill) + │ + │ 离线阶段:通过以下 3 个原子 Skill 生成 code_lib.zip + │ + ├── triton_w_nd2nz (原子 Skill ①: B 矩阵伪NZ格式优化) + ├── tiling_diagonal (原子 Skill ②: 对角 Tiling 核心映射) + └── k_axis_offset (原子 Skill ③: K 轴偏移优化) +``` + +## 1. matmul_skill_for_any_shape — 终端用户入口 Skill + +### 定位 + +这是一个 **可直接生成 matmul 算子的完整 Skill**,面向最终用户。用户无需手动组合各种优化策略,只需传入 shape 参数即可自动从预生成的代码库中筛选最优 Kernel 并完成测试验证。 + +### 使用前准备:服务器配置 + +该 Skill 的运行需要远程 NPU 服务器来编译和性能测试 Kernel,因此 **必须先配置服务器连接信息**。 + +编辑 `script/server_config.json`: + +```json +{ + "ip": "你的服务器IP", + "username": "root", + "auth_method": "key", + "password": "", + "ssh_key_path": "你的密钥文件路径", + "docker_container": "容器名称", + "npu_device": "NPU设备ID" +} +``` + +| 字段 | 说明 | +|------|------| +| `ip` | NPU 服务器 IP 地址 | +| `auth_method` | `"password"` 密码认证 或 `"key"` 密钥认证 | +| `password` | 密码认证时填写 | +| `ssh_key_path` | 密钥认证时填写 .pem 文件路径 | +| `docker_container` | Docker 容器名称 | +| `npu_device` | NPU 设备 ID(0-7) | + +配置完成后运行检测: + +```bash +python script/check_config.py +``` + +### 使用方式 + +**单 shape 测试**(例如 M=128 N=7168 K=4096): + +``` +请按照 SKILL.md 生成 M=128 N=7168 K=4096 的 FP16 高性能算子 +或 请按照 SKILL.md 生成 M=128 N=7168 K=4096 的 FP16 ND格式算子 +或 请按照 SKILL.md 生成 M=128 N=7168 K=4096 的 INT8 高性能算子 +或 请按照 SKILL.md 生成 M=128 N=7168 K=4096 的 INT8 ND格式算子 +``` + +**多 shape 批量测试**(使用 `shape_list.md` 中预定义的 shape): + +``` +请按照 SKILL.md 生成 FP16 高性能算子,且请使用 shape_list 中的 shape 进行测试 +或 请按照 SKILL.md 生成 FP16 ND格式算子,且请使用 shape_list 中的 shape 进行测试 +或 请按照 SKILL.md 生成 INT8 高性能算子,且请使用 shape_list 中的 shape 进行测试 +或 请按照 SKILL.md 生成 INT8 ND格式算子,且请使用 shape_list 中的 shape 进行测试 +``` + +### 支持的任务类型 + +| 用户需求 | `--code-subdir` | 精度 | 输入格式 | Kernel 数量 | +|----------|-----------------|------|----------|-------------| +| FP16 高性能 matmul | `code-fp16` | FP16 | 伪NZ | 36 | +| FP16 ND 格式 matmul | `code-fp16-fuse` | FP16 | ND + 融合 | 72 | +| INT8 高性能 matmul | `code-int8` | INT8 | 伪NZ | 36 | +| INT8 ND 格式 matmul | `code-int8-fuse` | INT8 | ND + 融合 | 72 | + +### 完整工作流 + +``` +Step 0: CONFIG_CHECK — 自动检测服务器配置是否有效 + ↓ +Step 1: CLEAN — 清理服务器端旧测试数据 + ↓ +Step 2: UPLOAD + START — 上传 code_lib.zip + 启动远程批量测试 + ↓ +Step 3: VERIFY — 确认远程测试进程已启动 + ↓ +用户自行等待并查看进度 / 下载结果 +``` + +### 输出 + +测试完成后,通过 `python script/analyze_results.py` 下载结果到本地 `output/` 目录: +- `best_kernels_report.xlsx` — 每个 shape 的最优 Kernel 性能报告 +- `*.py` — 所有最优 Kernel 的源代码 + +--- + +## 2. 三个原子 Skill + +以下三个 Skill 是构成 matmul_skill_for_any_shape 的底层基础。它们各自独立可用,专注于单一的优化维度。**matmul_skill_for_any_shape 中的 `code_lib.zip` 正是通过组合这三个原子 Skill 离线生成的。** + +### 2.1 triton_w_nd2nz — B 矩阵伪NZ格式优化 + +**文件**: [triton_w_nd2nz/SKILL.md](file:///d:/AscendICT/AscendOpGenAgent/skills/triton/matmul-related-gen/triton_w_nd2nz/SKILL.md) + +--- + +### 2.2 tiling_diagonal — 对角 Tiling 核心映射 + +**文件**: [tiling_diagonal/skill.md](file:///d:/AscendICT/AscendOpGenAgent/skills/triton/matmul-related-gen/tiling_diagonal/skill.md) + +--- + +### 2.3 k_axis_offset — K 轴偏移优化 + +**文件**: [k_axis_offset/skill.md](file:///d:/AscendICT/AscendOpGenAgent/skills/triton/matmul-related-gen/k_axis_offset/skill.md) + +--- + +## 3. 整体关系总结 + +``` +┌────────────────────────────────────────────────────────────────┐ +│ matmul_skill_for_any_shape │ +│ ┌──────────────────────────────────────────────────────────┐ │ +│ │ code_lib.zip │ │ +│ │ ┌─────────────┐ ┌──────────────┐ ┌───────────────┐ │ │ +│ │ │ code-fp16/ │ │ code-fp16- │ │ code-int8/ │ │ │ +│ │ │ (36 kernels) │ │ fuse/ (72) │ │ code-int8-fuse/│ │ │ +│ │ └─────────────┘ └──────────────┘ └───────────────┘ │ │ +│ └──────────────────────────────────────────────────────────┘ │ +│ │ │ +│ 离线阶段:通过组合 3 个原子 Skill 生成 code_lib.zip │ +│ │ │ +│ ┌─────────────────────────┼─────────────────────────┐ │ +│ │ │ │ │ +│ ▼ ▼ ▼ │ +│ triton_w_nd2nz tiling_diagonal k_axis_offset │ +│ (B矩阵格式优化) (对角Tiling映射) (K轴偏移) │ +│ │ │ +│ ┌────────────────────────────────────────────────────┘ │ +│ │ 运行时:根据用户 shape 从 code_lib 中筛选最优 Kernel │ +│ │ 上传 → 远程编译 → 性能测试 → 输出最优 Kernel 和报告 │ +│ └───────────────────────────────────────────────────────── │ +└────────────────────────────────────────────────────────────────┘ +``` + +### 三个原子 Skill 的职责分工 + +| Skill | 优化维度 | 影响范围 | +|-------|---------|---------| +| **triton_w_nd2nz** | 内存布局 | B 矩阵数据排布,Host 端预处理 + Kernel 端加载模式 | +| **tiling_diagonal** | Core 调度 | 计算块的分配策略,多核负载均衡 | +| **k_axis_offset** | 内存访问 | K 轴遍历顺序,减少多核内存访问冲突 | + +### 运行时流程 + +1. 用户指定 shape 和精度要求 +2. 系统从 `code_lib.zip` 中提取对应子目录的所有 Kernel +3. 上传至 NPU 服务器 +4. 远程逐一编译并执行性能测试(使用 `msprof`) +5. 收集性能数据,输出每个 shape 的最优 Kernel 和对应的源代码 + +生成的代码和测试报告存储在本地 `output/` 文件夹下。 diff --git a/skills/triton/matmul-related-gen/tiling_diagonal/skill.md b/skills/triton/matmul-related-gen/tiling_diagonal/skill.md new file mode 100644 index 00000000..d3561df0 --- /dev/null +++ b/skills/triton/matmul-related-gen/tiling_diagonal/skill.md @@ -0,0 +1,104 @@ +--- +name: Matrix Multiplication (matmul) Diagonal Tiling Core Logic Algorithm +description: Achieve efficient memory access through diagonal core mapping rules +--- + +## Code Generation Requirements: +1. Core mapping rule: Ensure adjacent block_id access different M, N blocks; tl.swizzle2d is not allowed +2. Generated code must satisfy CORE_Load_Balancing +3. Generated code must satisfy Coding_Rules +4. The algorithm must include at least 4 autotune parameters: BLOCK_M, BLOCK_N, BLOCK_K, and GROUP_SIZE + +## Autotune Config Generation Rules: +1. Keep all original BLOCK_M, BLOCK_N, BLOCK_K configurations from the base code +2. For each original configuration, generate variants with GROUP_SIZE=4 and GROUP_SIZE=8 +3. Total configs = original_configs × len(GROUP_SIZE_values) +4. GROUP_SIZE values: [4, 8] + +### Example: +Original configs (3): +- Config A: {'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 256} +- Config B: {'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 256} +- Config C: {'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 256} + +After applying diagonal tiling (6 configs): +- Config A + GROUP_SIZE=4: {'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 256, 'GROUP_SIZE': 4} +- Config A + GROUP_SIZE=8: {'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 256, 'GROUP_SIZE': 8} +- Config B + GROUP_SIZE=4: {'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 256, 'GROUP_SIZE': 4} +- Config B + GROUP_SIZE=8: {'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 256, 'GROUP_SIZE': 8} +- Config C + GROUP_SIZE=4: {'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 256, 'GROUP_SIZE': 4} +- Config C + GROUP_SIZE=8: {'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 256, 'GROUP_SIZE': 8} + +## Related Variable Definitions: + +### block_id Mapping Rule +- block_id=idx means block_id is assigned to a Core ID, for example: block_id=tl.program_id(0) +- (M,N) represents the 2D Tile block index in M and N directions +- block_id=0 → (0,1) means Core 0 processes the Tile block at position (0,1) + +### GROUP_2D Definition +- GROUP_2D contains GROUP_SIZE rows, each row has GROUP_SIZE data blocks +- The total number of data blocks in a GROUP_2D is: tiles_per_group = GROUP_SIZE * GROUP_SIZE + +## Algorithm Flow +1. Each group has GROUP_2D small data blocks; diagonal tiling within each group, taking GROUP_SIZE=4 as an example: + - First 4*4 diagonal tiling group: + block_id=0 → (0,0), block_id=1 → (1,1), block_id=2 → (2,2), block_id=3 → (3,3) + block_id=4 → (0,1), block_id=5 → (1,2), block_id=6 → (2,3), block_id=7 → (3,0) + block_id=8 → (0,2), block_id=9 → (1,3), block_id=10 → (2,0), block_id=11 → (3,1) + block_id=12 → (0,3), block_id=13 → (1,0), block_id=14 → (2,1), block_id=15 → (3,2) + - Second 4*4 diagonal tiling group: + block_id=16 → (0,4), block_id=17 → (1,5), block_id=18 → (2,6), block_id=19 → (3,7) + +## Reference Implementation +```python +# Get current program ID (Core ID) +pid = tl.program_id(0) + +# Calculate grid dimensions +num_pid_m = tl.cdiv(M, BLOCK_M) +num_pid_n = tl.cdiv(N, BLOCK_N) + +# Calculate total virtual blocks, padded by Group +# To implement diagonal tiling, we divide the grid into GROUP_SIZE x GROUP_SIZE SuperBlocks +groups_m = tl.cdiv(num_pid_m, GROUP_SIZE) +groups_n = tl.cdiv(num_pid_n, GROUP_SIZE) +num_groups = groups_m * groups_n + +# Number of blocks per Group +tiles_per_group = GROUP_SIZE * GROUP_SIZE + +# Total virtual tasks (including padded blocks) +total_virtual_tiles = num_groups * tiles_per_group + +# Loop through tasks assigned to current Core +for v_idx in range(pid, total_virtual_tiles, CORE_NUM): + # 1. Decode v_idx into Group ID and in-Group ID + group_idx = v_idx // tiles_per_group + in_group_idx = v_idx % tiles_per_group + + # 2. Calculate Group coordinates (grid of groups) + group_m = group_idx // groups_n + group_n = group_idx % groups_n + + # 3. Calculate in-Group coordinates (Diagonal Tiling) + # Mapping rule: + # block_id=0 -> (0,0), block_id=1 -> (1,1) ... + # i = local_id % GROUP_SIZE + # j = local_id // GROUP_SIZE + # m = i + # n = (i + j) % GROUP_SIZE + + i = in_group_idx % GROUP_SIZE + j = in_group_idx // GROUP_SIZE + local_m = i + local_n = (i + j) % GROUP_SIZE + + # 4. Calculate global coordinates + block_m = group_m * GROUP_SIZE + local_m + block_n = group_n * GROUP_SIZE + local_n + + # 5. Boundary check: use if instead of continue, as Triton compiler may not support continue + if block_m < num_pid_m and block_n < num_pid_n: + # --- Start matrix multiplication calculation --- +``` diff --git a/skills/triton/matmul-related-gen/triton_w_nd2nz/SKILL.md b/skills/triton/matmul-related-gen/triton_w_nd2nz/SKILL.md new file mode 100644 index 00000000..7b65e639 --- /dev/null +++ b/skills/triton/matmul-related-gen/triton_w_nd2nz/SKILL.md @@ -0,0 +1,361 @@ +--- +name: "b_matrix_format_optimization" +description: "Optimizes B matrix memory layout for GEMM operations with dot instructions. Invoke when code contains tl.dot() and needs matrix multiplication performance optimization." +--- + +# B Matrix Format Optimization + +## 1. Skill Overview + +### 1.1 What This Skill Does +This skill optimizes the memory layout of the B matrix in GEMM (General Matrix Multiply) operations by reorganizing data into a block-based format. This optimization improves memory access patterns and cache utilization, particularly effective on NPU hardware. + +### 1.2 When to Invoke This Skill +Invoke this skill when: +- The code contains `tl.dot()` instructions (this is a **prerequisite**) +- You need to optimize matrix multiplication performance +- The B matrix dimensions are compatible with the tile sizes used in the kernel + +## 2. Prerequisites + +**Step 1: Check for dot instruction** + +Before applying this optimization, verify that the code contains a dot instruction: + +```python +# Example of dot instruction in Triton +acc += tl.dot(a_val, b_val, out_dtype=tl.int32) +``` + +If no dot instruction is present, this optimization is not applicable. + +## 3. Implementation Steps + +### Step 2: Identify Tile Variables + +Identify the N-axis and K-axis tile variables in the code. These are typically defined as constants in the kernel function signature: + +```python +@triton.jit +def gemm_kernel( + # ... + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, # This is the N-axis tile variable + BLOCK_SIZE_K: tl.constexpr, # This is the K-axis tile variable + # ... +): + # ... +``` + +Common naming conventions: +- N-axis: `BLOCK_SIZE_N`, `BLOCK_N`, `TILE_N` +- K-axis: `BLOCK_SIZE_K`, `BLOCK_K`, `TILE_K` + +### Step 3: Convert B Matrix Format + +Use the identified tile variables to convert the B matrix: + +```python +def convert_b_matrix(b: torch.Tensor, block_k: int, block_n: int): + """ + Convert B matrix to optimized block format. + + Args: + b: Original B matrix with shape (K, N) + block_k: K-axis tile size (e.g., BLOCK_SIZE_K) + block_n: N-axis tile size (e.g., BLOCK_SIZE_N) + + Returns: + Optimized B matrix with shape (k_div_block_k, n_div_block_n, block_k, block_n) + """ + K, N = b.shape + assert K % block_k == 0, f'K ({K}) must be divisible by block_k ({block_k})' + assert N % block_n == 0, f'N ({N}) must be divisible by block_n ({block_n})' + + k_div_block_k = K // block_k + n_div_block_n = N // block_n + + # Reshape and permute for optimized memory layout + b_reshaped = b.view(k_div_block_k, block_k, n_div_block_n, block_n) + b_optimized = b_reshaped.permute(0, 2, 1, 3).contiguous() + + return b_optimized +``` + +### Step 4: Update Kernel Access Pattern + +Modify the kernel to access the optimized B matrix: + +**⚠️ CRITICAL: Load A MUST be executed before Load B in the kernel loop.** + +```python +# Inside the kernel +k_div_block_k = k // block_k +n_div_block_n = (block_n * BLOCK_SIZE_N) // block_n + +# Calculate offset for the current block +b_offset = k_div_block_k * stride_b_k_div_block_k + n_div_block_n * stride_b_n_div_block_n + +# IMPORTANT: Load A first (before loading B) +a_val = tl.load(a_ptr, boundary_check=(0, 1)) + +# Then load B +block_elements = block_k * block_n +b_raw = tl.load(b + b_offset + tl.arange(0, block_elements)) + +# Reshape to the required dimensions +b_val = b_raw.reshape(block_k, block_n) +``` + +## 4. Complete Example + +### 4.1 Kernel Definition + +```python +@triton.jit +def gemm_kernel( + a, b, c, M, N, K, + stride_am, stride_ak, + stride_b_k_div_block_k, stride_b_n_div_block_n, + stride_cm, stride_cn, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + pid = tl.program_id(0) + + # ... (block calculation logic) + + for k in range(0, K, BLOCK_SIZE_K): + # ⚠️ CRITICAL: Load A MUST be before Load B + a_val = tl.load(a_ptr, boundary_check=(0, 1)) + + # Optimized B matrix access (after loading A) + k_div_block_k = k // BLOCK_SIZE_K + n_div_block_n = (block_n * BLOCK_SIZE_N) // BLOCK_SIZE_N + + b_offset = k_div_block_k * stride_b_k_div_block_k + n_div_block_n * stride_b_n_div_block_n + b_raw = tl.load(b + b_offset + tl.arange(0, BLOCK_SIZE_K * BLOCK_SIZE_N)) + b_val = b_raw.reshape(BLOCK_SIZE_K, BLOCK_SIZE_N) + + # Dot instruction (prerequisite) + acc += tl.dot(a_val, b_val, out_dtype=tl.int32) + + a_ptr = tl.advance(a_ptr, (0, BLOCK_SIZE_K)) +``` + +### 4.2 Host Code with Autotune Config + +**⚠️ CRITICAL: When calling `convert_b_matrix`, you MUST add a comment immediately before the function call to specify the actual tile variable names from the autotune config.** + +```python +AUTOTUNE_CONFIGS = [ + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 256}), +] + +def gemm(a: torch.Tensor, b: torch.Tensor): + M, K = a.shape + K_from_b, N = b.shape + + # ========== CRITICAL: Must use this exact config structure ========== + # Get the selected config from autotune + # In practice, this would be determined by the autotuner + config = AUTOTUNE_CONFIGS[0] # Example: using first config + BLOCK_SIZE_N = config.kwargs['BLOCK_SIZE_N'] # TileN = 256 + BLOCK_SIZE_K = config.kwargs['BLOCK_SIZE_K'] # TileK = 256 + # ================================================================ + + # ========== CRITICAL: Add comment before convert_b_matrix ========== + b_optimized = convert_b_matrix(b, block_k=BLOCK_SIZE_K, block_n=BLOCK_SIZE_N) + # ==================================================================== + + # Calculate strides for the optimized format + k_div_block_k = K // BLOCK_SIZE_K + n_div_block_n = N // BLOCK_SIZE_N + + stride_b_n_div_block_n = BLOCK_SIZE_K * BLOCK_SIZE_N + stride_b_k_div_block_k = n_div_block_n * stride_b_n_div_block_n + + # Launch kernel + grid = (CORE_NUM,) + gemm_kernel[grid]( + a, b_optimized, c, M, N, K, + a.stride(0), a.stride(1), + stride_b_k_div_block_k, stride_b_n_div_block_n, + c.stride(0), c.stride(1), + BLOCK_SIZE_M=config.kwargs['BLOCK_SIZE_M'], + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + ) + + return c +``` + +### 4.3 Example with Multiple Autotune Configs + +When using multiple autotune configurations, the comment becomes even more important: + +```python +AUTOTUNE_CONFIGS = [ + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 256}), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128}), + triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 512, 'BLOCK_SIZE_K': 256}), +] + +@triton.autotune(configs=AUTOTUNE_CONFIGS, key=['M', 'N', 'K']) +@triton.jit +def gemm_kernel(...): + # ... kernel implementation + +def gemm(a: torch.Tensor, b: torch.Tensor): + M, K = a.shape + K_from_b, N = b.shape + + # ========== CRITICAL: Must use this exact config structure ========== + # After autotune selects a config, extract the tile sizes + # Note: In practice, you would get the selected config from the autotuner + # For demonstration, we use a specific config + config = AUTOTUNE_CONFIGS[0] + BLOCK_SIZE_N = config.kwargs['BLOCK_SIZE_N'] # TileN = 256 + BLOCK_SIZE_K = config.kwargs['BLOCK_SIZE_K'] # TileK = 256 + # ================================================================ + + # ========== CRITICAL: Comment format ========== + b_optimized = convert_b_matrix(b, block_k=BLOCK_SIZE_K, block_n=BLOCK_SIZE_N) + # ============================================== + + # ... rest of the implementation +``` + +## 5. Key Benefits + +### 5.1 Memory Access Efficiency +- **Contiguous memory access**: Loads entire blocks in a single operation +- **Improved cache utilization**: Block-based layout matches access patterns +- **Reduced memory transactions**: Fewer load operations for the same data + +### 5.2 Hardware Optimization +- **NPU-friendly**: Matches hardware memory access patterns +- **Reduced addressing overhead**: Simplified offset calculations +- **Better parallelism**: Efficient data loading supports higher computational throughput + +### 5.3 Flexibility +- **Dynamic tile sizes**: Adapts to any tile size defined in the kernel +- **No hardcoded constants**: Works with various block configurations +- **AutoTune compatible**: Can be combined with Triton's autotuning + +## 6. Important Considerations + +### 6.1 **⚠️ CRITICAL: Load Order Requirement** +**Load A MUST be executed before Load B in the kernel loop.** + +This ordering is essential for: +- **Memory access optimization**: Ensures proper memory access patterns +- **Hardware efficiency**: Matches NPU hardware expectations +- **Correct execution**: Prevents potential race conditions or undefined behavior + +**Correct order:** +```python +# ✓ CORRECT: Load A first +a_val = tl.load(a_ptr, boundary_check=(0, 1)) +# Then load B +b_raw = tl.load(b + b_offset + tl.arange(0, BLOCK_SIZE_K * BLOCK_SIZE_N)) +``` + +**Incorrect order:** +```python +# ✗ INCORRECT: Load B before A +b_raw = tl.load(b + b_offset + tl.arange(0, BLOCK_SIZE_K * BLOCK_SIZE_N)) +a_val = tl.load(a_ptr, boundary_check=(0, 1)) # WRONG! +``` + +### 6.2 Dimension Requirements +- K must be divisible by the K-axis tile size (block_k) +- N must be divisible by the N-axis tile size (block_n) +- If dimensions don't meet requirements, consider padding + +### 6.3 Memory Overhead +- The converted B matrix requires the same amount of memory +- Additional memory is needed during the conversion process +- Consider caching converted matrices for repeated use + +### 6.4 Performance Trade-offs +- **One-time conversion cost**: Initial conversion takes time +- **Best for repeated operations**: Most beneficial when B matrix is reused +- **Batch processing**: Ideal for scenarios with multiple GEMM operations + +### 6.5 **⚠️ CRITICAL: Comment Requirement** +**You MUST add a comment immediately before calling `convert_b_matrix` to specify:** +1. The actual variable name for `block_k` from the autotune config (e.g., `BLOCK_SIZE_K`) +2. The actual variable name for `block_n` from the autotune config (e.g., `BLOCK_SIZE_N`) + +**Format:** +```python +# ========== CRITICAL: Comment format ========== +b_optimized = convert_b_matrix(b, block_k=, block_n=) +``` + +**Why this is important:** +- Ensures traceability between the conversion and the kernel configuration +- Makes it clear which autotune config is being used +- Helps with debugging and maintenance +- Prevents mismatched tile sizes between conversion and kernel execution + +## 7. Troubleshooting + +### 7.1 Common Issues + +**Issue**: "K must be divisible by block_k" +- **Cause**: Matrix dimension doesn't match tile size +- **Solution**: Pad the matrix or adjust tile size + +**Issue**: Incorrect results after optimization +- **Cause**: Incorrect stride calculations or mismatched tile sizes +- **Solution**: Verify stride formulas and ensure tile sizes match between conversion and kernel + +**Issue**: No performance improvement +- **Cause**: Matrix too small or conversion overhead dominates +- **Solution**: Profile to identify bottleneck; consider matrix size + +**Issue**: Missing comment before convert_b_matrix call +- **Cause**: Forgot to add the required comment +- **Solution**: Add comment specifying the autotune config variable names + +**Issue**: Load B before Load A in kernel +- **Cause**: Incorrect load order in the kernel loop +- **Solution**: Ensure `tl.load(a_ptr)` is executed BEFORE `tl.load(b + b_offset + ...)` + +### 7.2 Validation + +Always validate the optimization produces correct results: + +```python +# Test with small example +c_optimized = gemm(a, b_optimized, block_k=BLOCK_K, block_n=BLOCK_N) +c_reference = torch.matmul(a.float(), b.float()).half() + +assert torch.allclose(c_optimized, c_reference, rtol=1e-3), "Results mismatch!" +``` + +## 8. Summary + +This skill optimizes B matrix memory layout for GEMM operations by: +1. **Checking prerequisite**: Ensures code contains dot instruction +2. **Identifying tile variables**: Finds BLOCK_K and BLOCK_N from kernel definition +3. **Converting format**: Reorganizes B matrix to block-based layout +4. **Adding required comment**: Specifies autotune config variable names before conversion +5. **Updating access pattern**: Modifies kernel to use optimized loading +6. **Ensuring correct load order**: Load A MUST be executed before Load B in the kernel loop + +The optimization is most effective when: +- Code contains `tl.dot()` instructions +- Matrix dimensions align with tile sizes +- B matrix is reused across multiple operations +- Running on hardware with specific memory access patterns (e.g., NPU) +- **Load A is executed before Load B in the kernel** + +## 9. References + +- [Triton Documentation](https://triton-lang.org/docs/) +- [Matrix Multiplication Optimization Guide](https://triton-lang.org/programming-guide.html) +- [NPU Programming Best Practices](https://developer.huawei.com/ict/en/site-type/docs)