From fe085f8a6954710a2750e382203b3bd1735c4fc8 Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti <113389104+yraparti@users.noreply.github.com> Date: Fri, 29 May 2026 17:09:29 +0000 Subject: [PATCH 001/143] [rocm-libraries] ROCm/rocm-libraries#7761 (commit 237b766) [CK][CK TILE] Clean up tile_engine grouped_conv harness (#7761) ## Motivation Tile_engine grouped_conv contains ML heuristic validation scripts which cause confusion to new developers. So, this PR is intended to relocate the scripts into dispatcher/heuristic directory to maintain separation of concern. ## Technical Details The grouped_conv tile_engine directory is a benchmarking harness for grouped convolution kernels; ML-heuristic content does not belong there. - Move compare_ml_vs_oracle.py and validate_ml_vs_oracle.py from tile_engine/ops/grouped_conv/ to dispatcher/heuristics/validation/grouped_conv/, and rebase their sys.path / oracle CSV / model dir lookups for the new location (CSV path is now an --oracle-csv flag instead of a hard-coded sibling). - Move GROUPED_CONV_HEURISTIC_REPORT.md (system-level ML report) into dispatcher/heuristics/ where the rest of the heuristic docs live. - Rewrite tile_engine/ops/grouped_conv/README.md as a pure benchmarking / dispatcher-sweep doc (kernel enumeration, JIT pipeline, CSV schema, problem registry), in the style of tile_engine/ops/fmha/README.md. All ML training / model-efficiency content is removed and replaced with a pointer to dispatcher/heuristics/. ## Test Plan Validation scripts are re-wired and tested locally ## Test Result Tests passed on local machine. ## Submission Checklist - [x ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- .../grouped_conv/compare_ml_vs_oracle.py | 21 +- .../grouped_conv/validate_ml_vs_oracle.py | 44 ++- tile_engine/ops/grouped_conv/README.md | 345 ++++++------------ 3 files changed, 161 insertions(+), 249 deletions(-) rename {tile_engine/ops => dispatcher/heuristics/validation}/grouped_conv/compare_ml_vs_oracle.py (94%) rename {tile_engine/ops => dispatcher/heuristics/validation}/grouped_conv/validate_ml_vs_oracle.py (87%) diff --git a/tile_engine/ops/grouped_conv/compare_ml_vs_oracle.py b/dispatcher/heuristics/validation/grouped_conv/compare_ml_vs_oracle.py similarity index 94% rename from tile_engine/ops/grouped_conv/compare_ml_vs_oracle.py rename to dispatcher/heuristics/validation/grouped_conv/compare_ml_vs_oracle.py index 974b85e4f8..ce8dca980b 100644 --- a/tile_engine/ops/grouped_conv/compare_ml_vs_oracle.py +++ b/dispatcher/heuristics/validation/grouped_conv/compare_ml_vs_oracle.py @@ -1,4 +1,7 @@ #!/usr/bin/env python3 +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + """ Compare ML heuristic predictions against oracle benchmark results. @@ -114,7 +117,15 @@ def run_end_to_end_workflow(args): elif args.problem_set: print(f"Problem set: {args.problem_set}") # Import problem set dynamically - sys.path.insert(0, str(Path(__file__).parent / "problems")) + # Problem sets live with the benchmarking harness in tile_engine. + _THIS_DIR = Path(__file__).parent + _TILE_ENGINE_GROUPED_CONV = ( + _THIS_DIR.parent.parent.parent.parent + / "tile_engine" + / "ops" + / "grouped_conv" + ) + sys.path.insert(0, str(_TILE_ENGINE_GROUPED_CONV / "problems")) try: problem_module = __import__(args.problem_set) problem_attr = ( @@ -165,15 +176,15 @@ def run_end_to_end_workflow(args): print() print("Please use the manual workflow documented in README.md:") print() - print(" 1. Create problem set file in problems/") + print(" 1. Create problem set file in tile_engine/ops/grouped_conv/problems/") print( - " 2. Run: python grouped_conv_full_benchmark.py --problems --csv oracle.csv" + " 2. Run: cd tile_engine/ops/grouped_conv && python grouped_conv_full_benchmark.py --problems --csv oracle.csv" ) print( - " 3. Run: cd ../../dispatcher/heuristics && python predict_cli.py --problem-module --output ml.csv" + " 3. Run: cd dispatcher/heuristics && python predict_cli.py --problem-module --output ml.csv" ) print( - " 4. Run: cd ../../tile_engine/ops/grouped_conv && python compare_ml_vs_oracle.py --oracle-csv oracle.csv --ml-csv ml.csv --plot result.png" + " 4. Run: cd dispatcher/heuristics/validation/grouped_conv && python compare_ml_vs_oracle.py --oracle-csv oracle.csv --ml-csv ml.csv --plot result.png" ) print() diff --git a/tile_engine/ops/grouped_conv/validate_ml_vs_oracle.py b/dispatcher/heuristics/validation/grouped_conv/validate_ml_vs_oracle.py similarity index 87% rename from tile_engine/ops/grouped_conv/validate_ml_vs_oracle.py rename to dispatcher/heuristics/validation/grouped_conv/validate_ml_vs_oracle.py index 9e5124caf8..0da88839ae 100755 --- a/tile_engine/ops/grouped_conv/validate_ml_vs_oracle.py +++ b/dispatcher/heuristics/validation/grouped_conv/validate_ml_vs_oracle.py @@ -12,18 +12,24 @@ 4. Reports efficiency metrics """ +import argparse import sys from pathlib import Path import pandas as pd import numpy as np _THIS_DIR = Path(__file__).parent -_DISPATCHER_ROOT = _THIS_DIR.parent.parent.parent / "dispatcher" +# This file lives at: /projects/composablekernel/dispatcher/heuristics/validation/grouped_conv/ +# Walk up three levels (validation -> heuristics -> dispatcher) to find the dispatcher root. +_DISPATCHER_ROOT = _THIS_DIR.parent.parent.parent +_CK_ROOT = _DISPATCHER_ROOT.parent +# Problem definitions still live with the benchmarking harness in tile_engine. +_TILE_ENGINE_GROUPED_CONV = _CK_ROOT / "tile_engine" / "ops" / "grouped_conv" sys.path.insert(0, str(_DISPATCHER_ROOT / "python")) sys.path.insert(0, str(_DISPATCHER_ROOT / "heuristics")) sys.path.insert(0, str(_DISPATCHER_ROOT / "codegen")) -sys.path.insert(0, str(_THIS_DIR / "problems")) +sys.path.insert(0, str(_TILE_ENGINE_GROUPED_CONV / "problems")) from validation_holdout import VALIDATION_PROBLEMS # noqa: E402 from predict import Predictor # noqa: E402 @@ -81,11 +87,31 @@ def _build_kernel_name(kconf, ndim): ) -# Load model -model_dir = ( - _DISPATCHER_ROOT - / "heuristics/models/grouped_conv_forward_bf16_gfx950_2d_3d_no_compv5" +# Parse CLI args +_parser = argparse.ArgumentParser(description=__doc__) +_parser.add_argument( + "--oracle-csv", + type=Path, + default=_TILE_ENGINE_GROUPED_CONV / "validation_oracle_results.csv", + help="Oracle benchmark CSV (produced by tile_engine/ops/grouped_conv/grouped_conv_full_benchmark.py)", +) +_parser.add_argument( + "--model-dir", + type=Path, + default=_DISPATCHER_ROOT + / "heuristics/models/grouped_conv_forward_bf16_gfx950_2d_3d_no_compv5", + help="Trained LightGBM model directory.", +) +_parser.add_argument( + "--output", + type=Path, + default=_THIS_DIR / "validation_heuristic_vs_oracle.csv", + help="Where to write the per-problem comparison CSV.", ) +_args = _parser.parse_args() + +# Load model +model_dir = _args.model_dir feature_engine = GroupedConvFeatureEngine() predictor = Predictor(model_dir, feature_engine=feature_engine) @@ -98,7 +124,7 @@ def _build_kernel_name(kconf, ndim): print() # Load oracle benchmark results -oracle_df = pd.read_csv(_THIS_DIR / "validation_oracle_results.csv") +oracle_df = pd.read_csv(_args.oracle_csv) print(f"Oracle measurements: {len(oracle_df)}") print() @@ -281,7 +307,7 @@ def _build_kernel_name(kconf, ndim): print() # Save detailed results - results_df.to_csv(_THIS_DIR / "validation_heuristic_vs_oracle.csv", index=False) - print("Detailed results saved to: validation_heuristic_vs_oracle.csv") + results_df.to_csv(_args.output, index=False) + print(f"Detailed results saved to: {_args.output}") else: print("ERROR: No predictions could be compared with oracle data") diff --git a/tile_engine/ops/grouped_conv/README.md b/tile_engine/ops/grouped_conv/README.md index 71a5ecacdc..73434b4224 100644 --- a/tile_engine/ops/grouped_conv/README.md +++ b/tile_engine/ops/grouped_conv/README.md @@ -1,294 +1,169 @@ -# Grouped Convolution ML Heuristics & Benchmarking +# Grouped Convolution Tile Engine -Training data collection and validation utilities for ML-based kernel selection in grouped convolution operations. +Benchmarking harness for grouped convolution kernels via the CK dispatcher's pipelined JIT compilation. -## Overview +Covers all three variants -- **forward**, **backward-data**, **backward-weight** -- across the suffix-aware pipeline pool (compv3 / compv4 / compv5 / mem, intrawave / interwave, optional `dsb` / `si` suffixes) for 2D and 3D shapes. -This directory supports the **ML heuristic system** for grouped convolution kernel selection. The system achieves **99.67% efficiency** on unseen production workloads by predicting optimal kernels without exhaustive GPU search. +This directory is purely a benchmarking and sweep tool. ML kernel-selection heuristics, training, and validation live in `dispatcher/heuristics/` (see [Related Documentation](#related-documentation)). -**Key Results:** -- Forward pass: 99.67% mean efficiency (validated on 10 unseen MIOpen shapes) -- 70% perfect oracle matches (selected exact best kernel) -- <1ms selection latency (30,000-60,000× faster than exhaustive search) +## Directory Layout -See [dispatcher/heuristics/GROUPED_CONV_ML_SUMMARY.md](../../dispatcher/heuristics/GROUPED_CONV_ML_SUMMARY.md) for full technical details. - ---- - -## Files - -### Benchmarking & Data Collection -- **`grouped_conv_full_benchmark.py`** - Systematic sweep for training data (kernels × problems) -- **`run_one_grouped_conv_kernel.py`** - Subprocess worker for isolated GPU execution -- **`test_batch_benchmark.py`** - Quick integration test (2 kernels × small problems) -- **`grouped_conv_instance_builder.py`** - Kernel configuration generator from JSON - -### ML Validation -- **`validate_ml_vs_oracle.py`** - Compare ML predictions vs exhaustive GPU search -- **`compare_ml_vs_oracle.py`** - Analysis of ML vs oracle performance - -### Configuration -- **`configs/*.json`** - Kernel trait configurations (forward, bwd_data, bwd_weight) -- **`problems/*.py`** - Problem datasets (training, validation, MIOpen production shapes) - ---- - -## ML Heuristic Workflow - -### 1. Training Data Collection - -Already completed. Training datasets: -- **Forward**: 48,845 samples (1,372 unique shapes) - Tier-1 extended -- **Bwd Data**: 14,562 samples (701 unique shapes) -- **Bwd Weight**: 18,150 samples (921 unique shapes) - -If you need to collect new data: - -```bash -# Full benchmark sweep (all kernels × all problems) -python grouped_conv_full_benchmark.py \ - --variant forward \ - --category full \ - --workers 256 \ - --output training_data_forward_bf16.csv ``` - -### 2. Training Models - -Models are located in `dispatcher/heuristics/models/`: -- `grouped_conv_forward_bf16_gfx950/` - **Production-ready** (99.67% efficiency) -- `grouped_conv_bwd_data_bf16_gfx950/` - Trained, needs hardware validation -- `grouped_conv_bwd_weight_bf16_gfx950/` - Trained, needs hardware validation - -To train new models, see [dispatcher/heuristics/README.md](../../dispatcher/heuristics/README.md). - -### 3. Validation - -Validate ML model performance on unseen shapes: - -```bash -cd ../../dispatcher/heuristics/validation/grouped_conv - -# Quick sanity check on training shapes (hardware) -python validate_training_shapes.py --direction forward - -# Backward models validation (no GPU) -python validate_backward_models.py +grouped_conv/ + grouped_conv_full_benchmark.py Orchestrator: enumerate kernels x problems, JIT compile, benchmark + grouped_conv_instance_builder.py Kernel enumeration from JSON trait config + run_one_grouped_conv_kernel.py Subprocess worker (one kernel, fresh GPU context) + README.md This file + configs/ Kernel trait configurations + forward_bf16.json Forward bf16 (compv3/v4/v5) + bwd_data.json Backward data (compv3 / mem) + bwd_weight.json Backward weight (compv3 / mem) + problems/ Problem datasets (registry keys consumed by --problems) + forward_2d.py / forward_3d.py + bwd_data_2d.py / bwd_data_3d.py + bwd_weight_2d.py / bwd_weight_3d.py + *_test_validation.py Small unseen-shape subsets + validation_holdout.py VALIDATION_PROBLEMS (300 forward shapes) ``` -See [dispatcher/heuristics/validation/README.md](../../dispatcher/heuristics/validation/README.md) for details. - ---- - -## Problem Datasets - -Located in `problems/`: - -### Training Sets -- **`forward_training.py`** - 2,630 shapes (300 MIOpen + 2,330 synthetic) -- **`forward_training_miopen.py`** - 300 MIOpen production shapes -- **`bwd_data_synthetic_extended.py`** - Backward data training set -- **`bwd_weight_synthetic_extended.py`** - Backward weight training set - -### Validation Sets (Unseen) -- **`bwd_data_test_validation.py`** - 10 unseen backward data shapes -- **`bwd_weight_test_validation.py`** - 10 unseen backward weight shapes - -### Dataset Generator -- **`create_miopen_training_set.py`** - Extract shapes from MIOpen ALL_CONFIGS_FULL.txt - ---- - -## Benchmarking Usage - -### Quick Test (2 Kernels × Few Problems) +## Quick Start ```bash -# Test benchmark pipeline -python test_batch_benchmark.py -``` +# Count kernels matching a trait config without compiling +python grouped_conv_instance_builder.py configs/forward_bf16.json --arch gfx950 --count-only -### Full Sweep (All Kernels × All Problems) +# List kernel names +python grouped_conv_instance_builder.py configs/forward_bf16.json --arch gfx950 --list -```bash -# Forward: 20 kernels × 200 problems = 4,000 measurements +# Smoke benchmark: forward 2D on the validation subset python grouped_conv_full_benchmark.py \ --variant forward \ - --category full \ + --problems forward_2d_test_validation \ --workers 256 \ - --output sweep_forward.csv + --output sweep_forward_smoke.csv -# Backward data +# Full sweep: all forward kernels x all forward-2D problems python grouped_conv_full_benchmark.py \ - --variant bwd_data \ - --category full \ - --workers 256 + --variant forward \ + --problems forward_2d \ + --workers 256 \ + --output sweep_forward_2d.csv -# Backward weight -python grouped_conv_full_benchmark.py \ - --variant bwd_weight \ - --category full \ - --workers 256 +# Backward data / weight sweeps +python grouped_conv_full_benchmark.py --variant bwd_data --problems bwd_data_2d --output sweep_bwd_data.csv +python grouped_conv_full_benchmark.py --variant bwd_weight --problems bwd_weight_2d --output sweep_bwd_weight.csv ``` -**Output**: CSV with columns: -``` -kernel,problem_idx,N,C,K,G,Hi,Wi,Y,X,stride_h,stride_w,pad_h,pad_w,latency_ms,tflops,non_zero -``` +The benchmark always starts fresh and overwrites `--output`. Move or rename the file beforehand if you need to keep prior results. -**Note**: The benchmark always starts fresh and overwrites the output CSV file. If you need to preserve previous results, rename or move the CSV file before running a new benchmark. +## How It Works ---- +### Kernel Enumeration -## Instance Builder +``` +JSON trait config (variant + allowed pipelines / wave modes / suffixes) + --> grouped_conv_instance_builder.py + --> dispatcher/codegen/grouped_config_rules.py (tile + suffix-aware pool) + --> list of GroupedConvKernelConfig + --> optional --filter expression +``` -Generate kernel configurations from JSON trait files: +The pipeline rules in `dispatcher/codegen/grouped_config_rules.py` are the single source of truth for the kernel pool (tile sizes, wave modes, pipeline variants, `dsb` / `si` suffixes). The instance builder reads a JSON trait allow-list and produces the cartesian product of legal configurations. -```bash -# List all kernels matching config -python grouped_conv_instance_builder.py configs/forward_bf16.json --arch gfx950 --list +### Benchmark Pipeline -# Count kernels -python grouped_conv_instance_builder.py configs/forward_bf16.json --count-only +``` +grouped_conv_full_benchmark.py (orchestrator) + |-- grouped_conv_instance_builder.py enumerate kernel configs + |-- Build phase codegen -> hipcc -> link .so (serial; avoids fork + GPU init issues) + '-- Benchmark phase one subprocess per kernel batch + '-- run_one_grouped_conv_kernel.py + '-- GpuGroupedConvRunner fresh HIP context per problem +``` -# Apply filter -python grouped_conv_instance_builder.py configs/forward_bf16.json \ - --filter "c.tile_n >= 128 and c.pipeline == 'compv5'" --list +Key design choices: -# Export to JSON -python grouped_conv_instance_builder.py configs/forward_bf16.json \ - --export-json kernels.json -``` +1. **Subprocess isolation** -- a fresh HIP context per kernel batch avoids cumulative driver/device leaks during long sweeps. +2. **Serial GPU access** -- accurate timing, no contention. +3. **Path-only build in the main process** -- the orchestrator never initializes the GPU runtime, so `fork()` after codegen is safe. +4. **Batch size ~20 kernels/subprocess** -- empirically a good throughput/overhead tradeoff. -### Config Files +> The `--workers` flag controls codegen/compile parallelism for the build phase. Benchmarking itself is serial per device. -- **`forward_bf16.json`** - Forward BF16 (compv3/v4/v5, 30 kernels) -- **`bwd_data.json`** - Backward data (compv3/mem, 20 kernels) -- **`bwd_weight.json`** - Backward weight (compv3/mem, 20 kernels) +## JSON Config Format -**Trait filtering** (see configs for examples): ```json { "variant": "forward", "trait_config": { - "data_type": {"values": ["bf16"]}, - "pipeline": {"values": ["compv3", "compv4", "compv5"]}, - "ndim_spatial": {"values": [2]} + "data_type": {"values": ["bf16"]}, + "pipeline": {"values": ["compv3", "compv4", "compv5"]}, + "wave_mode": {"values": ["intrawave", "interwave"]}, + "ndim_spatial": {"values": [2, 3]} } } ``` ---- +Allowed keys mirror `GroupedConvKernelConfig` fields. See `dispatcher/codegen/grouped_config_rules.py` for the full schema. -## Architecture +### Filtering examples -Based on FMHA tile engine design with subprocess isolation: +```bash +# Only large tiles on compv5 +python grouped_conv_instance_builder.py configs/forward_bf16.json \ + --arch gfx950 \ + --filter "c.tile_n >= 128 and c.pipeline == 'compv5'" --list +# Export the resolved kernel list to JSON +python grouped_conv_instance_builder.py configs/forward_bf16.json \ + --arch gfx950 --export-json kernels.json ``` -grouped_conv_full_benchmark.py (orchestrator) - ├─> grouped_conv_instance_builder.py (generate kernel configs) - ├─> Build phase: JIT compile all kernels (serial, avoids fork/GPU issues) - └─> Benchmark phase: subprocess workers (serial GPU access) - └─> run_one_grouped_conv_kernel.py (subprocess) - └─> GpuGroupedConvRunner (fresh GPU context per problem) -``` - -**Key design decisions:** -1. **Subprocess isolation** - Fresh GPU context prevents memory leaks -2. **Batch size 20** - Optimal kernels per subprocess -3. **Path-only build** - Main process never initializes GPU -4. **Serial GPU access** - Accurate timing, no contention -5. **Serial codegen/compile** - Avoids ProcessPoolExecutor + GPU fork() issues - -**Note**: The `--workers` flag is accepted for API compatibility but currently ignored. -Codegen and compilation run serially to avoid GPU context issues with process forking. - -**Success rate**: 99.5% (3,760/3,780 measurements succeeded) ---- +## Problem Registry -## Example Workflow: New Data Collection +`--problems` accepts **only registry keys**, not file paths. The keys are wired in `grouped_conv_full_benchmark.py`. Current keys: -```bash -# 1. Generate problem set -cd problems/ -python create_miopen_training_set.py \ - --input /path/to/ALL_CONFIGS_FULL.txt \ - --output forward_training_new.py \ - --count 500 - -# 2. Collect training data -cd .. -python grouped_conv_full_benchmark.py \ - --variant forward \ - --category full \ - --workers 256 \ - --output new_training_data.csv - -# 3. Convert to parquet -cd ../../dispatcher/heuristics -python convert_csv_to_parquet.py \ - --input ../../tile_engine/ops/grouped_conv/new_training_data.csv \ - --output data/grouped_conv_forward_bf16_gfx950/new_data.parquet - -# 4. Train model -python train.py \ - --data_dir data/ \ - --out_dir models/grouped_conv_forward_bf16_gfx950_v2 \ - --op grouped_conv \ - --variant forward - -# 5. Validate (sanity check on training shapes) -cd validation/grouped_conv -python validate_training_shapes.py --direction forward -``` +| Key | Direction | Notes | +|----------------------------------|----------------|------------------------------------------| +| `forward_2d` / `forward_3d` | forward | Full training-grade problem sets | +| `bwd_data_2d` / `bwd_data_3d` | backward data | Full training-grade problem sets | +| `bwd_weight_2d` / `bwd_weight_3d`| backward wgt | Full training-grade problem sets | +| `*_test_validation` | per direction | Small unseen-shape subsets | +| `validation_holdout` | forward | 300 shapes (250 2D + 50 3D) | ---- +Adding a new subset requires both a `problems/.py` file and a registry entry in `grouped_conv_full_benchmark.py`. -## Performance Results +Each problem module exposes a list of dataclasses with fields `N, C, K, G, Hi, Wi[, Di], Y, X[, Z], stride_h, stride_w[, stride_d], pad_h, pad_w[, pad_d]` and optional `dilation_*`. -### Forward Pass (Production-Ready) -- **Mean efficiency**: 99.67% on 10 unseen MIOpen shapes -- **Perfect matches**: 70% (7/10 selected exact oracle best) -- **Min efficiency**: 98.4% (even on edge case: 1×491 spatial) -- **Selection time**: <1ms (vs 30-60s exhaustive search) +## Output CSV Schema -### Backward Passes (Prediction-Validated) -- **Bwd Data**: 14,562 samples, prediction quality tested -- **Bwd Weight**: 18,150 samples, prediction quality tested -- **Status**: Models trained, hardware validation pending +``` +kernel, problem_idx, N, C, K, G, [Di,] Hi, Wi, [Z,] Y, X, + [stride_d,] stride_h, stride_w, + [pad_d,] pad_h, pad_w, + latency_ms, tflops, non_zero +``` -See [dispatcher/heuristics/GROUPED_CONV_ML_SUMMARY.md](../../dispatcher/heuristics/GROUPED_CONV_ML_SUMMARY.md) for full metrics. +`non_zero` is a sanity flag (output checksum != 0). Failed launches are written with `latency_ms=N/A` and `tflops=0`. ---- +## Hardware -## Hardware Tested +- Validated on AMD Instinct MI355X (gfx950). +- Datatypes: bf16 (primary), fp16, fp32. +- Pipelines: compv3 / compv4 / compv5 (forward), compv3 / mem (backward). +- Schedulers: intrawave, interwave (with optional `dsb`, `si` suffixes). -- **GPU**: AMD MI300 (gfx950) -- **Datatypes**: BF16 (primary), FP16, FP32 -- **Pipelines**: CompV3, CompV4, CompV5 (forward), CompV3/Mem (backward) -- **Schedulers**: Intrawave, Interwave -- **Tile sizes**: 16×64×64, 32×64×64, 64×64×64, 128×128×64, etc. +### GPU access caveat (this host) ---- +On the dev host the device files have non-default GIDs (`/dev/kfd` GID 506, `/dev/dri/renderD144` GID 109). If `hipMalloc` returns code 100 (`hipErrorOutOfMemory`) on every allocation, it is a permissions issue, not VRAM exhaustion. Launch the benchmark via `sudo -u sshuser bash -lc '...'` so the process tree picks up `kfdhost`, `renderhost`, and `video` groups. ## Related Documentation -- **ML System Overview**: [dispatcher/heuristics/GROUPED_CONV_ML_SUMMARY.md](../../dispatcher/heuristics/GROUPED_CONV_ML_SUMMARY.md) -- **Training Pipeline**: [dispatcher/heuristics/README.md](../../dispatcher/heuristics/README.md) -- **Validation Framework**: [dispatcher/heuristics/validation/README.md](../../dispatcher/heuristics/validation/README.md) -- **Python Examples**: [dispatcher/examples/grouped_conv/python/README_ML_HEURISTIC.md](../../dispatcher/examples/grouped_conv/python/README_ML_HEURISTIC.md) - ---- - -## Next Steps - -**For Forward Pass**: Production-ready, integrate into runtime dispatcher - -**For Backward Passes**: Run prediction-quality check -```bash -cd ../../dispatcher/heuristics/validation/grouped_conv -python validate_backward_models.py -``` +Anything ML-heuristic-related has been moved out of this directory: -Target: >85% mean efficiency on unseen shapes before production deployment. +- **ML training pipeline & models**: `dispatcher/heuristics/README.md` +- **ML vs oracle comparison & validation**: `dispatcher/heuristics/validation/grouped_conv/` + - `validate_ml_vs_oracle.py` -- run trained predictor over a problem set and compare against oracle CSVs produced by this harness. + - `compare_ml_vs_oracle.py` -- post-hoc comparison of oracle + ML prediction CSVs (efficiency, top-k, scatter plot). +- **Dispatcher Python API**: `dispatcher/python/` +- **End-to-end examples**: `dispatcher/examples/grouped_conv/` \ No newline at end of file From 15c904b46077f7c396a04981cb082feab53de02a Mon Sep 17 00:00:00 2001 From: Aviral Goel <191153937+AviralGoelAMD@users.noreply.github.com> Date: Fri, 29 May 2026 18:45:13 +0000 Subject: [PATCH 002/143] [rocm-libraries] ROCm/rocm-libraries#7724 (commit 4cb149a) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ck_tile: add FillUniformScaleDistribution and fix MX GEMM scale init (#7724) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Problem MX GEMM pipeline tests were passing vacuously: scale bytes were drawn from a fixed range (40–60) which, for e8m0, maps to scales ≈ 10⁻²⁷ — far below FP16 min denorm. Both GPU and CPU produced all-zero outputs, so numerical checks passed without exercising the GEMM. ### Changes **`include/ck_tile/host/fill.hpp`** — new `FillUniformScaleDistribution` functor - Accepts human-readable float bounds and maps them to the raw byte range of any ExMy scale type (e8m0, e4m3, e5m3) by re-centering the IEEE 754 exponent into the type's bias space - Sampling is uniform over raw bytes → uniform over representable values - Fixes left-shift UB: uses multiplication instead of `<< mant_bits` to avoid shifting negative signed integers (C++17 UB) - Adds `assert(min_r <= max_r)` to catch inverted-range UB when both bounds exceed the type's representable range - Provides default member values (0.125f, 2.0f) and `std::optional` seed consistent with sibling fillers - `/** */` Doxygen style with `@note` on snapping asymmetry **`test/ck_tile/gemm_mx/test_mx_gemm_pipeline_util.hpp`** — fix scale initialization - Replace manual byte-range distribution with `FillUniformScaleDistribution<>{0.125f, 2.0f}` - Use distinct seeds for scale_a (11941) and scale_b (11943) to avoid correlated scale tensors that were causing 60 test failures for fp4+e5m3/e4m3 combinations **`test/ck_tile/utility/test_fill.cpp`** — new unit tests for `FillUniformScaleDistribution` - 16 typed tests across e8m0, e4m3, e5m3: validity, range, reproducibility, coverage, snapping, stress, nullopt seed, and range overload - Test helper `expected_raw_range` mirrors implementation clamping exactly --- include/ck_tile/host/fill.hpp | 88 +++++ .../gemm_mx/test_mx_gemm_pipeline_util.hpp | 23 +- test/ck_tile/utility/test_fill.cpp | 314 ++++++++++++++++++ 3 files changed, 414 insertions(+), 11 deletions(-) diff --git a/include/ck_tile/host/fill.hpp b/include/ck_tile/host/fill.hpp index 82c5e1185b..e24fb95164 100644 --- a/include/ck_tile/host/fill.hpp +++ b/include/ck_tile/host/fill.hpp @@ -4,6 +4,7 @@ #pragma once #include +#include #include #include #include @@ -479,6 +480,93 @@ struct FillConstant } }; +/** + * @brief Fills a range with uniformly distributed random ExMy (exponent-mantissa) scale values. + * + * Accepts human-readable float bounds and maps them to the raw byte range of the + * target scale type by re-centering the IEEE 754 exponent into the type's own + * bias space. Sampling is then uniform over raw bytes, which is uniform over + * representable values of the type. + * + * @tparam ScaleType An ExMy scale type (e.g. e8m0_t, e4m3_t, e5m3_t). + * + * @note Both bounds snap down to the nearest representable power-of-two in ScaleType space. + * If min_scale_ is not an exact power-of-two, the effective lower bound is lower than + * requested; if max_scale_ is not an exact power-of-two, the effective upper bound is + * also lower than requested. + * + * Fields: min_scale_ (lower float bound), max_scale_ (upper float bound, no value + * generated exceeds it), seed_ (RNG seed; nullopt for random device, default 11939). + * Precondition: min_scale_ <= max_scale_. Violating this is undefined behavior. + * Usage: + * FillUniformScaleDistribution{0.125f, 2.0f, 42}(scale_tensor); + * FillUniformScaleDistribution{0.125f, 2.0f}(buf.begin(), buf.end()); + */ +template +struct FillUniformScaleDistribution +{ + float min_scale_{0.125f}; + float max_scale_{2.0f}; + std::optional seed_{11939}; + + template + void operator()(ForwardIter first, ForwardIter last) const + { + using RawType = typename ScaleType::type; // uint8_t for all current ExMy types + + // Bias and mantissa layout for the target type, resolved at compile time. + constexpr int float_bias = 127; // IEEE 754 single-precision bias + constexpr int type_bias = + numeric_traits::bias; // e.g. 127 (e8m0), 7 (e4m3), 15 (e5m3) + constexpr int mant_bits = + numeric_traits::mant; // mantissa bits: 0 (e8m0), 3 (e4m3/e5m3) + + // Extract the biased IEEE 754 exponent byte from each float bound. + // get_exponent(f) == (bit_cast(f) >> 23) & 0xFF - the raw 8-bit exponent field. + // Non-power-of-two values snap down: get_exponent(0.1f) == get_exponent(0.0625f) == 123. + const int ieee_min = static_cast(numeric_utils::get_exponent(min_scale_)); + const int ieee_max = static_cast(numeric_utils::get_exponent(max_scale_)); + + // Absolute limits of the raw byte space for this type. + // raw=0 is reserved: denorm-zero for e4m3/e5m3 (decodes to 0.0), and subnormal + // territory for e8m0 (2^-127) - excluded to keep all generated values usable as scales. + // binary_max is the last finite raw value (binary_nan - 1 for all ExMy types). + constexpr int raw_min = 1; + constexpr int raw_max = static_cast(numeric::binary_max); + + // Re-center the IEEE 754 exponent offset into the target type's bias space and pack into + // raw bytes. (ieee_exp - float_bias) gives the true power: e.g. 123-127 = -4 for 0.0625. + // Adding type_bias maps into the target encoding: e.g. -4+7 = 3 for e4m3. + // Left-shift by mant_bits places the exponent field: e.g. 3<<3 = 24 for e4m3. + // max_r uses mant=0 (not | mant_mask) so it decodes to exactly max_scale - the + // power-of-two itself. This ensures no generated value exceeds max_scale_ in float space. + // std::max/min clamp to the valid byte range, preventing out-of-range or NaN raw values. + const int min_r = + std::max(((ieee_min - float_bias) + type_bias) * (1 << mant_bits), raw_min); + const int max_r = + std::min(((ieee_max - float_bias) + type_bias) * (1 << mant_bits), raw_max); + + // Precondition: clamping must not invert the range. This can happen when both bounds + // exceed the type's representable range in the same direction (both too large or too + // small). If triggered, use bounds within the type's representable range. + assert(min_r <= max_r); + + // Sample raw bytes uniformly in [min_r, max_r], then construct ScaleType directly + // from the raw byte - bypassing the float ctor which would discard mantissa bits. + std::mt19937 gen(seed_.has_value() ? *seed_ : std::random_device{}()); + std::uniform_int_distribution dist(min_r, max_r); + std::generate(first, last, [&]() { return ScaleType(static_cast(dist(gen))); }); + } + + // Range overload: accepts any container or HostTensor with begin()/end(). + template + void operator()(ForwardRange&& range) const + { + (*this)(std::begin(std::forward(range)), + std::end(std::forward(range))); + } +}; + //---------------------------------------------------------------------------------------------- /// @brief Transforms given input to fit 2:4 structured sparsity pattern so /// every subgroup of 4 elements contain at most 2 non-zero elements diff --git a/test/ck_tile/gemm_mx/test_mx_gemm_pipeline_util.hpp b/test/ck_tile/gemm_mx/test_mx_gemm_pipeline_util.hpp index ea0fa174b2..981d4c1d33 100644 --- a/test/ck_tile/gemm_mx/test_mx_gemm_pipeline_util.hpp +++ b/test/ck_tile/gemm_mx/test_mx_gemm_pipeline_util.hpp @@ -1,7 +1,6 @@ // Copyright (c) Advanced Micro Devices, Inc., or its affiliates. // SPDX-License-Identifier: MIT #pragma once -#include #include #include @@ -417,16 +416,18 @@ class TestCkTileMxGemmPipeline : public ::testing::Test } { - std::mt19937 gen(std::chrono::steady_clock::now().time_since_epoch().count()); - std::uniform_int_distribution dist(40, 60); - for(auto& s : scale_a.mData) - { - s = AScaleDataType(static_cast(dist(gen))); - } - for(auto& s : scale_b.mData) - { - s = BScaleDataType(static_cast(dist(gen))); - } + // Fill scale tensors with values uniformly drawn from [0.125, 2.0] = [2^-3, 2^1]. + // This spans 5 exponent bands centred around 1.0, keeping scales numerically + // well-behaved without saturating the accumulator. + // + // Per-type raw byte ranges produced (raw bytes sampled uniformly within each): + // e8m0_t (bias=127, mant=0): raw in [124, 128] -> floats {0.125, 0.25, 0.5, 1.0, 2.0} + // e4m3_t (bias=7, mant=3): raw in [32, 64] -> floats 0.125 .. 2.0 + // e5m3_t (bias=15, mant=3): raw in [96, 128] -> floats 0.125 .. 2.0 + // No generated value exceeds 2.0 for any type. + // A and B use different seeds so their scale values are uncorrelated. + ck_tile::FillUniformScaleDistribution{0.125f, 2.0f, 11941}(scale_a); + ck_tile::FillUniformScaleDistribution{0.125f, 2.0f, 11943}(scale_b); } // Pre-shuffle scale buffers for the hardware diff --git a/test/ck_tile/utility/test_fill.cpp b/test/ck_tile/utility/test_fill.cpp index f67dee9757..8ea80a2ae9 100644 --- a/test/ck_tile/utility/test_fill.cpp +++ b/test/ck_tile/utility/test_fill.cpp @@ -2,10 +2,18 @@ // SPDX-License-Identifier: MIT #include "ck_tile/host/fill.hpp" +#include "ck_tile/host/host_tensor.hpp" #include "ck_tile/host/joinable_thread.hpp" +#include "ck_tile/core/numeric/e4m3.hpp" +#include "ck_tile/core/numeric/e5m3.hpp" +#include "ck_tile/core/numeric/e8m0.hpp" + #include +#include #include #include +#include +#include #include using namespace ck_tile; @@ -156,3 +164,309 @@ TYPED_TEST(FillUniformDistributionTest, EdgeCases) } } } // namespace test + +// ============================================================ +// FillUniformScaleDistribution tests +// ============================================================ + +namespace test_scale { + +// Returns true if f is a finite, non-NaN float. +bool is_valid_float(float f) { return std::isfinite(f); } + +// Returns true if f is an exact power of two (positive). +bool is_power_of_two(float f) +{ + if(f <= 0.f || !is_valid_float(f)) + return false; + uint32_t bits; + std::memcpy(&bits, &f, sizeof(bits)); + return (bits & 0x007fffffu) == 0u; // mantissa bits all zero +} + +// Compute the expected raw range for a given ScaleType and float bounds. +template +static std::pair expected_raw_range(float min_f, float max_f) +{ + constexpr int float_bias = 127; + constexpr int type_bias = ck_tile::numeric_traits::bias; + constexpr int mant_bits = ck_tile::numeric_traits::mant; + const int ieee_min = static_cast(ck_tile::numeric_utils::get_exponent(min_f)); + const int ieee_max = static_cast(ck_tile::numeric_utils::get_exponent(max_f)); + // raw=0 excluded: decodes to 0.0 for e4m3/e5m3 and to 2^-127 for e8m0 - same + // assumption as the implementation in FillUniformScaleDistribution. + constexpr int raw_min = 1; + constexpr int raw_max = static_cast(ck_tile::numeric::binary_max); + const int scale = 1 << mant_bits; + const int min_r = std::max(((ieee_min - float_bias) + type_bias) * scale, raw_min); + const int max_r = std::min(((ieee_max - float_bias) + type_bias) * scale, raw_max); + return {min_r, max_r}; +} + +// ---- typed fixture ------------------------------------------------- +template +class FillUniformScaleDistributionTest : public ::testing::Test +{ +}; + +using ScaleTypes = ::testing::Types; +TYPED_TEST_SUITE(FillUniformScaleDistributionTest, ScaleTypes); + +// 1. No garbage: all generated values are finite (not NaN, not Inf). +TYPED_TEST(FillUniformScaleDistributionTest, NoGarbageValues) +{ + using S = TypeParam; + ck_tile::HostTensor buf({10000}); + ck_tile::FillUniformScaleDistribution{0.0625f, 4.0f, 42}(buf.begin(), buf.end()); + std::size_t i = 0; + for(const S& v : buf) + { + float f = static_cast(v); + EXPECT_TRUE(is_valid_float(f)) + << "NaN/Inf at index " << i + << " raw=" << static_cast(static_cast(v)); + ++i; + } +} + +// 2. All generated raw bytes are within [min_r, max_r]. +TYPED_TEST(FillUniformScaleDistributionTest, RawValuesInExpectedRange) +{ + using S = TypeParam; + constexpr float min_scale = 0.0625f; + constexpr float max_scale = 4.0f; + auto [min_r, max_r] = expected_raw_range(min_scale, max_scale); + ck_tile::HostTensor buf({10000}); + ck_tile::FillUniformScaleDistribution{min_scale, max_scale, 7}(buf.begin(), buf.end()); + std::size_t i = 0; + for(const S& v : buf) + { + int raw = static_cast(static_cast(v)); + EXPECT_GE(raw, min_r) << "raw below min at index " << i; + EXPECT_LE(raw, max_r) << "raw above max at index " << i; + ++i; + } +} + +// 3. Reproducibility: identical seed -> identical output. +TYPED_TEST(FillUniformScaleDistributionTest, SameSeedSameOutput) +{ + using S = TypeParam; + ck_tile::HostTensor a({1000}), b({1000}); + ck_tile::FillUniformScaleDistribution{0.125f, 2.0f, 99}(a.begin(), a.end()); + ck_tile::FillUniformScaleDistribution{0.125f, 2.0f, 99}(b.begin(), b.end()); + EXPECT_EQ(0, std::memcmp(a.data(), b.data(), a.size() * sizeof(S))); +} + +// 4. Different seeds produce different outputs (with overwhelming probability). +TYPED_TEST(FillUniformScaleDistributionTest, DifferentSeedsDifferentOutput) +{ + using S = TypeParam; + ck_tile::HostTensor a({1000}), b({1000}); + ck_tile::FillUniformScaleDistribution{0.125f, 2.0f, 1}(a.begin(), a.end()); + ck_tile::FillUniformScaleDistribution{0.125f, 2.0f, 2}(b.begin(), b.end()); + EXPECT_NE(0, std::memcmp(a.data(), b.data(), a.size() * sizeof(S))); +} + +// 5. Single-value range: [v, v] -> all generated raw bytes fall in that exponent band. +TYPED_TEST(FillUniformScaleDistributionTest, SingleValueRange) +{ + using S = TypeParam; + constexpr float pivot = 1.0f; + auto [min_r, max_r] = expected_raw_range(pivot, pivot); + ck_tile::HostTensor buf({2000}); + ck_tile::FillUniformScaleDistribution{pivot, pivot, 5}(buf.begin(), buf.end()); + std::size_t i = 0; + for(const S& v : buf) + { + int raw = static_cast(static_cast(v)); + EXPECT_GE(raw, min_r) << "index " << i; + EXPECT_LE(raw, max_r) << "index " << i; + EXPECT_TRUE(is_valid_float(static_cast(v))) << "index " << i; + ++i; + } +} + +// 6. Non-power-of-two bounds snap to the nearest lower power-of-two exponent. +// Verify outputs are still within the snapped raw range. +TYPED_TEST(FillUniformScaleDistributionTest, NonPowerOfTwoBoundsSnap) +{ + using S = TypeParam; + // 0.1 snaps to 0.0625 (2^-4); 3.5 snaps to 2.0 (2^1) + auto [min_r, max_r] = expected_raw_range(0.0625f, 2.0f); // snapped bounds + ck_tile::HostTensor buf({5000}); + ck_tile::FillUniformScaleDistribution{0.1f, 3.5f, 13}(buf.begin(), buf.end()); + std::size_t i = 0; + for(const S& v : buf) + { + int raw = static_cast(static_cast(v)); + EXPECT_GE(raw, min_r) << "index " << i; + EXPECT_LE(raw, max_r) << "index " << i; + ++i; + } +} + +// 7. Coverage: for a small range, every possible raw value appears at least once +// after enough samples (probabilistic - extremely unlikely to fail with 100k draws). +TYPED_TEST(FillUniformScaleDistributionTest, AllRawValuesGenerated) +{ + using S = TypeParam; + constexpr float min_f = 0.5f; + constexpr float max_f = 2.0f; + auto [min_r, max_r] = expected_raw_range(min_f, max_f); + const int range_size = max_r - min_r + 1; + const std::size_t draws = static_cast(range_size) * 5000; + + ck_tile::HostTensor buf({draws}); + ck_tile::FillUniformScaleDistribution{min_f, max_f, 77}(buf.begin(), buf.end()); + + std::unordered_set seen; + for(auto& v : buf) + seen.insert(static_cast(static_cast(v))); + + EXPECT_EQ(static_cast(seen.size()), range_size) + << "Expected " << range_size << " distinct raw values, got " << seen.size(); +} + +// 8. e8m0 specific: all generated values must be exact powers of two. +TEST(FillUniformScaleDistributionE8M0, AllValuesPowersOfTwo) +{ + using S = ck_tile::e8m0_t; + ck_tile::HostTensor buf({10000}); + ck_tile::FillUniformScaleDistribution{0.0625f, 4.0f, 33}(buf.begin(), buf.end()); + std::size_t i = 0; + for(const S& v : buf) + { + float f = static_cast(v); + EXPECT_TRUE(is_power_of_two(f)) << "Non-power-of-two at index " << i << " value=" << f; + ++i; + } +} + +// 9. Wide range stress: large tensor, wide float range, no garbage. +TYPED_TEST(FillUniformScaleDistributionTest, WideRangeStress) +{ + using S = TypeParam; + ck_tile::HostTensor buf({50000}); + ck_tile::FillUniformScaleDistribution{1.f / 1024, 1024.f, 0}(buf.begin(), buf.end()); + std::size_t i = 0; + for(const S& v : buf) + { + float f = static_cast(v); + EXPECT_TRUE(is_valid_float(f)) << "Bad value at index " << i; + EXPECT_GT(f, 0.f) << "Non-positive scale at index " << i; + ++i; + } +} + +// 10. Empty range does not crash. +TYPED_TEST(FillUniformScaleDistributionTest, EmptyRangeNoCrash) +{ + using S = TypeParam; + ck_tile::HostTensor buf({0}); + EXPECT_NO_THROW( + (ck_tile::FillUniformScaleDistribution{1.0f, 1.0f, 0}(buf.begin(), buf.end()))); +} + +// 11. For e8m0 (mant=0), every generated value is exactly within [min_scale, max_scale]. +// Each exponent band has exactly one value so no overshoot is possible. +TEST(FillUniformScaleDistributionE8M0, StrictFloatBounds) +{ + using S = ck_tile::e8m0_t; + constexpr float min_f = 0.0625f, max_f = 4.0f; + ck_tile::HostTensor buf({10000}); + ck_tile::FillUniformScaleDistribution{min_f, max_f, 11}(buf.begin(), buf.end()); + std::size_t i = 0; + for(const S& v : buf) + { + float f = static_cast(v); + EXPECT_GE(f, min_f) << "index " << i; + EXPECT_LE(f, max_f) << "index " << i; + ++i; + } +} + +// 12. Unlike test 11 (e8m0 only, both bounds strict), this test covers all ExMy types and +// checks only the upper bound. The lower bound is not strict for types with non-zero +// mantissa bits (e4m3/e5m3): mantissa bits allow values between consecutive +// power-of-two exponents, so some generated values can fall below min_scale (test 13 +// verifies this). The upper bound IS strict for all types because max_r is set to the +// exact power-of-two raw encoding (mant=0), so the highest output is exactly max_scale_. +TYPED_TEST(FillUniformScaleDistributionTest, StrictFloatUpperBound) +{ + using S = TypeParam; + constexpr float min_f = 0.0625f, max_f = 4.0f; + ck_tile::HostTensor buf({10000}); + ck_tile::FillUniformScaleDistribution{min_f, max_f, 22}(buf.begin(), buf.end()); + std::size_t i = 0; + for(const S& v : buf) + { + float f = static_cast(v); + EXPECT_LE(f, max_f) << "value " << f << " exceeds max_scale at index " << i; + ++i; + } +} + +// 13. When min_scale is not an exact power of two it snaps down to the nearest lower +// power-of-two exponent, so some generated values will be below min_scale. +TYPED_TEST(FillUniformScaleDistributionTest, NonPowerOfTwoMinSnapsBelow) +{ + using S = TypeParam; + // 0.1 is not a power of two; get_exponent snaps it down to 0.0625 (2^-4). + // Values in [0.0625, 0.1) are therefore reachable. + constexpr float min_f = 0.1f; + ck_tile::HostTensor buf({10000}); + ck_tile::FillUniformScaleDistribution{min_f, 4.0f, 33}(buf.begin(), buf.end()); + bool found_below = false; + for(auto& v : buf) + if(static_cast(v) < min_f) + found_below = true; + EXPECT_TRUE(found_below) + << "Expected some values below non-power-of-two min_scale due to exponent snapping"; +} + +// 14. Extreme bounds that exceed the type's representable range clamp safely +// and still produce only finite, positive values - no NaN, no crash. +TYPED_TEST(FillUniformScaleDistributionTest, ExtremeOutOfRangeBoundsClampSafely) +{ + using S = TypeParam; + ck_tile::HostTensor buf({5000}); + ck_tile::FillUniformScaleDistribution{1e-38f, 1e38f, 55}(buf.begin(), buf.end()); + std::size_t i = 0; + for(const S& v : buf) + { + float f = static_cast(v); + EXPECT_TRUE(std::isfinite(f)) << "index " << i; + EXPECT_GT(f, 0.f) << "index " << i; + ++i; + } +} + +// 15. nullopt seed: two calls produce different outputs (random device seeding). +TYPED_TEST(FillUniformScaleDistributionTest, NulloptSeedProducesRandomOutput) +{ + using S = TypeParam; + ck_tile::HostTensor a({500}), b({500}); + ck_tile::FillUniformScaleDistribution{0.125f, 2.0f, std::nullopt}(a.begin(), a.end()); + ck_tile::FillUniformScaleDistribution{0.125f, 2.0f, std::nullopt}(b.begin(), b.end()); + EXPECT_NE(0, std::memcmp(a.data(), b.data(), a.size() * sizeof(S))); +} + +// 16. Range overload: passing a ck_tile::HostTensor directly compiles and fills correctly. +TYPED_TEST(FillUniformScaleDistributionTest, RangeOverloadFillsHostTensor) +{ + using S = TypeParam; + ck_tile::HostTensor buf({1000}); + ck_tile::FillUniformScaleDistribution{0.125f, 2.0f, 7}(buf); + auto [min_r, max_r] = expected_raw_range(0.125f, 2.0f); + std::size_t i = 0; + for(const S& v : buf) + { + int raw = static_cast(static_cast(v)); + EXPECT_GE(raw, min_r) << "index " << i; + EXPECT_LE(raw, max_r) << "index " << i; + ++i; + } +} + +} // namespace test_scale From 0edfcf06e5db4467282c8e2ffd9e48602dade553 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Fri, 29 May 2026 19:18:57 +0000 Subject: [PATCH 003/143] [rocm-libraries] ROCm/rocm-libraries#7894 (commit 5e66689) [CK] add credentials to docker manifest inspect call ## Motivation This should fix an issue that we recently encountered in CI when we exceeded the limit of accessing docker without authentication: [2026-05-29T16:08:42.447Z] + docker manifest inspect --insecure rocm/composable_kernel:ck_ub24.04_rocm7.13 [2026-05-29T16:08:42.833Z] toomanyrequests: You have reached your unauthenticated pull rate limit. https://www.docker.com/increase-rate-limit ## Technical Details ## Test Plan ## Test Result ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- Jenkinsfile | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/Jenkinsfile b/Jenkinsfile index 0347592405..7cd7a2546a 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -470,7 +470,9 @@ def buildAndPushDockerImage(String install_prefix, String image_name, String doc if(!forceBuild){ try{ echo "Checking for image: ${image_name}" - sh "docker manifest inspect --insecure ${image_name}" + withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { + sh "docker manifest inspect --insecure ${image_name}" + } echo "Image: ${image_name} found! Skipping building image" return image_name } From 95c916369c53449d37372c37df892f191d937e23 Mon Sep 17 00:00:00 2001 From: Emily Martins <65371150+ecamartins@users.noreply.github.com> Date: Fri, 29 May 2026 21:36:49 +0000 Subject: [PATCH 004/143] [rocm-libraries] ROCm/rocm-libraries#7584 (commit 060bad5) [CK_TILE] Fix Stream-K k_size calculation ## Motivation In a recent benchmarking task for CK Tile Stream-K algorithm, we identified that certain instances segfault. This change works to fix the bug and adds necessary regression tests. ## Technical Details The StreamK kernel constructs tensor views using a `k_size` parameter that determines how much of the K dimension to process in each iteration. Previously, this was calculated as: ```cpp index_t k_size = num_loop_sk * TilePartitioner::KPerBlock; ``` This calculation assumes all macro tiles along K are exactly `KPerBlock` in size. However, when `K % KPerBlock != 0`, the final macro tile along K has a remainder size of `K % KPerBlock`, not a full `KPerBlock` (see the figure below): image With the old code, a workgroup working with the `MPerBlock x (K % KPerBlock)` tile in A and B risk accessing illegal memory. Hence, this change ensures that when `K % KPerBlock != 0`, workgroups processing iterations that include the final macro-tile along K calculate the correct `k_size` based on the remainder rather than assuming a full `KPerBlock`. ## Test Plan I added the following tests: 1. Unit tests added for the Stream-K Tile Partitioner: - `StreamKTilePartitionerBaseGetKSize/NoRemainderTiles` - validates full tiles - `StreamKTilePartitionerBaseGetKSize/RemainderTiles` - validates remainder handling 2. Regression tests that test a case where `K % KPerBlock != 0` ## Test Result Tests passed locally on gfx90a, gfx942, and gfx950. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- .../streamk_gemm/streamk_gemm_kernel.hpp | 3 +- .../streamk_gemm_tile_partitioner.hpp | 26 +++++ .../streamk_gemm_tile_partitioner_impl.hpp | 45 ++++++++- test/ck_tile/gemm_streamk/CMakeLists.txt | 3 +- .../gemm_streamk/generate_test_files.py | 5 + .../test_gemm_streamk_regression_cases.inc | 14 +++ .../gemm_streamk/test_gemm_streamk_types.hpp | 5 + .../gemm_streamk/test_generate_test_files.py | 25 +++++ .../test_streamk_tile_partitioner.cpp | 94 ++++++++++++++++++- .../test_streamk_tile_partitioner_common.hpp | 52 +++++++++- 10 files changed, 262 insertions(+), 10 deletions(-) create mode 100644 test/ck_tile/gemm_streamk/test_gemm_streamk_regression_cases.inc diff --git a/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_kernel.hpp index 87fef5089a..f6db545465 100644 --- a/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_kernel.hpp @@ -363,7 +363,8 @@ struct StreamKKernel // Determine the total size along the K dimension the workgroup is using in this // iteration (used to construct tensor views). - index_t k_size = num_loop_sk * TilePartitioner::KPerBlock; + index_t k_size = amd_wave_read_first_lane( + kargs.tile_partitioner.get_k_size(num_loop_sk, local_iter_end)); // Get the K offsets for the A and B tensors auto [i_k_a, i_k_b] = GetKOffsets( diff --git a/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner.hpp b/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner.hpp index 79af955367..6a6ac7ffec 100644 --- a/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner.hpp +++ b/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner.hpp @@ -139,6 +139,19 @@ struct StreamKTilePartitionerBase CK_TILE_DEVICE auto get_output_tile_index(index_t tile_idx) const noexcept -> tuple; + /** + * @brief Calculates the total size along the K dimension the workgroup is using in this + * Stream-K loop iteration + * + * @param num_macro_tiles The number of macro tiles along the K dimension this workgroup is + * assigned. + * @param local_iter_end The workgroup's non-inclusive end iteration that is local to its + * current tile. + * @return index_t The K dimension size for the current Stream-K loop iteration. + */ + CK_TILE_DEVICE index_t get_k_size(index_t num_macro_tiles, + index_t local_iter_end) const noexcept; + /** * @brief Calculates the total space needed for the partials and flags buffers. * @@ -208,6 +221,17 @@ struct StreamKTilePartitionerBase */ CK_TILE_HOST_DEVICE index_t get_n() const noexcept; + /** + * @brief Returns the k dimension for the GEMM problem. + */ + CK_TILE_HOST_DEVICE index_t get_k() const noexcept; + + /** + * @brief Returns the remainder along the k dimension when k is not evenly divisible by + * KPerBlock. + */ + CK_TILE_HOST_DEVICE index_t get_remainder_along_k() const noexcept; + /** * @brief Returns an estimate of the number of workgroups writing to the same macro tile in C. */ @@ -244,6 +268,8 @@ struct StreamKTilePartitionerBase index_t extra_iters_; index_t total_dp_iters_; index_t n_; + index_t k_; + index_t remainder_along_k_; }; /** diff --git a/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner_impl.hpp b/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner_impl.hpp index f2c4d54599..fba70fb9a9 100644 --- a/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner_impl.hpp +++ b/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner_impl.hpp @@ -8,10 +8,11 @@ namespace ck_tile { template StreamKTilePartitionerBase::StreamKTilePartitionerBase( index_t m, index_t n, index_t k, index_t max_active_wgs) - : max_active_wgs_{max_active_wgs}, n_{n} + : max_active_wgs_{max_active_wgs}, n_{n}, k_{k} { - iters_per_tile_ = integer_divide_ceil(k, KPerBlock); - num_tiles_ = integer_divide_ceil(m, MPerBlock) * integer_divide_ceil(n_, NPerBlock); + iters_per_tile_ = integer_divide_ceil(k, KPerBlock); + num_tiles_ = integer_divide_ceil(m, MPerBlock) * integer_divide_ceil(n_, NPerBlock); + remainder_along_k_ = k % KPerBlock; bool big_enough = num_tiles_ > max_active_wgs_; index_t remainder_tiles = num_tiles_ % max_active_wgs_; @@ -250,6 +251,21 @@ StreamKTilePartitionerBase::get_n() c return n_; } +template +CK_TILE_HOST_DEVICE index_t +StreamKTilePartitionerBase::get_k() const noexcept +{ + return k_; +} + +template +CK_TILE_HOST_DEVICE index_t +StreamKTilePartitionerBase::get_remainder_along_k() + const noexcept +{ + return remainder_along_k_; +} + template CK_TILE_HOST index_t StreamKTilePartitionerBase::estimate_num_wgs_per_tile() @@ -334,6 +350,29 @@ StreamKTilePartitionerBase::remap_xcd return block_1d_id; } +template +CK_TILE_DEVICE index_t +StreamKTilePartitionerBase::get_k_size( + index_t num_macro_tiles, index_t local_iter_end) const noexcept +{ + // Determine if this workgroup is responsible for the last macro tile in the K dimension + bool last_tile = get_iters_per_tile() == local_iter_end; + index_t k_size; + // If there is no remainder or if the workgroup was not assigned the last macro tile along K, + // then their k_size will be a multiple of KPerBlock. + if(!remainder_along_k_ || !last_tile) + { + k_size = num_macro_tiles * KPerBlock; + } + // Otherwise, there's a remainder. So, k_size is not a multiple of KPerBlock. + else + { + k_size = (num_macro_tiles - 1) * KPerBlock + remainder_along_k_; + } + + return k_size; +} + template diff --git a/test/ck_tile/gemm_streamk/CMakeLists.txt b/test/ck_tile/gemm_streamk/CMakeLists.txt index 2c5b3bb04c..91e5ea5341 100644 --- a/test/ck_tile/gemm_streamk/CMakeLists.txt +++ b/test/ck_tile/gemm_streamk/CMakeLists.txt @@ -33,12 +33,13 @@ if(GPU_TARGETS MATCHES "gfx90a|gfx942|gfx950") set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${STREAMK_TYPES_HEADER} ${STREAMK_GEN_SCRIPT}) # Define the targets and their corresponding executable names - set(STREAMK_GEN_TARGETS extended atomic_smoke linear_smoke tree_smoke pipelines_smoke) + set(STREAMK_GEN_TARGETS extended atomic_smoke linear_smoke tree_smoke pipelines_smoke regression) set(STREAMK_GEN_EXEC_EXTENDED test_ck_tile_streamk_extended) set(STREAMK_GEN_EXEC_ATOMIC_SMOKE test_ck_tile_streamk_atomic_smoke) set(STREAMK_GEN_EXEC_LINEAR_SMOKE test_ck_tile_streamk_linear_smoke) set(STREAMK_GEN_EXEC_TREE_SMOKE test_ck_tile_streamk_tree_smoke) set(STREAMK_GEN_EXEC_PIPELINES_SMOKE test_ck_tile_streamk_pipelines_smoke) + set(STREAMK_GEN_EXEC_REGRESSION test_ck_tile_streamk_regression) # Collect all test targets for umbrella label set(CK_TILE_GEMM_STREAMK_TEST_TARGETS diff --git a/test/ck_tile/gemm_streamk/generate_test_files.py b/test/ck_tile/gemm_streamk/generate_test_files.py index 61a28c2a46..d1b0e118e7 100644 --- a/test/ck_tile/gemm_streamk/generate_test_files.py +++ b/test/ck_tile/gemm_streamk/generate_test_files.py @@ -80,6 +80,10 @@ class {class_name} : public TestCkTileStreamK "test_gemm_streamk_atomic_cases.inc", ], }, + "regression": { + "filter": lambda suffix: suffix == "Regression", + "inc_files": ["test_gemm_streamk_regression_cases.inc"], + }, } # --------------------------------------------------------------------------- # @@ -97,6 +101,7 @@ class {class_name} : public TestCkTileStreamK ("Tree", "tree"), ("CompV3", "compv3"), ("Pipelines", "pipelines"), + ("Regression", "regression"), ] diff --git a/test/ck_tile/gemm_streamk/test_gemm_streamk_regression_cases.inc b/test/ck_tile/gemm_streamk/test_gemm_streamk_regression_cases.inc new file mode 100644 index 0000000000..d521301f87 --- /dev/null +++ b/test/ck_tile/gemm_streamk/test_gemm_streamk_regression_cases.inc @@ -0,0 +1,14 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +TYPED_TEST(TEST_SUITE_NAME, StreamK_Remainder_Along_K) +{ + + ck_tile::index_t M = 64; + ck_tile::index_t N = 270128; + ck_tile::index_t K = 112; + + this->Run(M, N, K); +} diff --git a/test/ck_tile/gemm_streamk/test_gemm_streamk_types.hpp b/test/ck_tile/gemm_streamk/test_gemm_streamk_types.hpp index cbd3f0f066..6eea3f21c7 100644 --- a/test/ck_tile/gemm_streamk/test_gemm_streamk_types.hpp +++ b/test/ck_tile/gemm_streamk/test_gemm_streamk_types.hpp @@ -197,4 +197,9 @@ using KernelTypesStreamKPipelines = ::testing::Types< std::tuple< Row, Col, Row, F16, F16, F32, F16, I256, I256, I32, I32, I32, I16, NonPersistent, CompV4, Tree>, std::tuple< Col, Col, Row, F16, F16, F32, F16, I256, I256, I32, I32, I32, I16, Persistent, CompV4, Linear> >; + +using KernelTypesStreamKRegression = ::testing::Types< + std::tuple< Row, Col, Row, BF16, BF16, F32, BF16, I256, I256, I32, I32, I32, I16, Persistent, CompV3, Atomic>, + std::tuple< Col, Col, Row, BF16, BF16, F32, BF16, I256, I256, I32, I32, I32, I16, Persistent, CompV3, Atomic> +>; // clang-format on diff --git a/test/ck_tile/gemm_streamk/test_generate_test_files.py b/test/ck_tile/gemm_streamk/test_generate_test_files.py index 7b904da319..7c22a1e929 100644 --- a/test/ck_tile/gemm_streamk/test_generate_test_files.py +++ b/test/ck_tile/gemm_streamk/test_generate_test_files.py @@ -68,6 +68,11 @@ def test_pipelines_token(self): expected_tag = "pipelines" self.assertEqual(suffix_to_file_tag(suffix), expected_tag) + def test_regression_token(self): + suffix = "Regression" + expected_tag = "regression" + self.assertEqual(suffix_to_file_tag(suffix), expected_tag) + def test_unknown_token(self): suffix = "unknown" with self.assertRaises(ValueError): @@ -206,6 +211,26 @@ def test_tree_smoke(self): ] self.validate_entries(entries, expected) + def test_regression(self): + """Test regression target: matches suffix == 'Regression'. + Includes: Regression + """ + mock_content = ( + "using KernelTypesStreamKRegression = ...\n" + "using KernelTypesStreamKFp16Linear = ...\n" + "using KernelTypesStreamKPipelines = ...\n" + ) + with patch("builtins.open", mock_open(read_data=mock_content)): + entries = parse_types_header("fake_path.hpp", "regression") + expected = [ + { + "type_alias": "KernelTypesStreamKRegression", + "class_name": "TestCkTileStreamKRegression", + "file_tag": "regression", + } + ] + self.validate_entries(entries, expected) + class TestOutputPath(unittest.TestCase): def test_output_path(self): diff --git a/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner.cpp b/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner.cpp index 75ba762892..c4c14dde17 100644 --- a/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner.cpp +++ b/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner.cpp @@ -12,7 +12,7 @@ TEST(StreamKTilePartitionerBaseConstructor, SKOnly) Config::M, Config::N, Config::K, Config::MAX_ACTIVE_WGS}; StreamKTilePartitionerBaseExpected expected_values{ - 2, 0, 3, 4, 1, 2, 1, 0, 2, Config::MAX_ACTIVE_WGS, Config::N}; + 2, 0, 3, 4, 1, 2, 1, 0, 2, Config::MAX_ACTIVE_WGS, Config::N, Config::K, 0}; validate_streamk_base_constructor(expected_values, tile_partitioner); } @@ -24,7 +24,7 @@ TEST(StreamKTilePartitionerBaseConstructor, DPOnly) Config::M, Config::N, Config::K, Config::MAX_ACTIVE_WGS}; StreamKTilePartitionerBaseExpected expected_values{ - 0, 6, 0, 0, 0, 2, 0, 12, 6, Config::MAX_ACTIVE_WGS, Config::N}; + 0, 6, 0, 0, 0, 2, 0, 12, 6, Config::MAX_ACTIVE_WGS, Config::N, Config::K, 0}; validate_streamk_base_constructor(expected_values, tile_partitioner); } @@ -36,7 +36,7 @@ TEST(StreamKTilePartitionerBaseConstructor, DP2TileSK) Config::M, Config::N, Config::K, Config::MAX_ACTIVE_WGS}; StreamKTilePartitionerBaseExpected expected_values{ - 4, 3, 3, 8, 2, 2, 2, 6, 7, Config::MAX_ACTIVE_WGS, Config::N}; + 4, 3, 3, 8, 2, 2, 2, 6, 7, Config::MAX_ACTIVE_WGS, Config::N, Config::K, 0}; validate_streamk_base_constructor(expected_values, tile_partitioner); } @@ -48,7 +48,19 @@ TEST(StreamKTilePartitionerBaseConstructor, EdgeCase) Config::M, Config::N, Config::K, Config::MAX_ACTIVE_WGS}; StreamKTilePartitionerBaseExpected expected_values{ - 0, 1, 0, 0, 0, 2, 0, 2, 1, Config::MAX_ACTIVE_WGS, Config::N}; + 0, 1, 0, 0, 0, 2, 0, 2, 1, Config::MAX_ACTIVE_WGS, Config::N, Config::K, 0}; + validate_streamk_base_constructor(expected_values, tile_partitioner); +} + +TEST(StreamKTilePartitionerBaseConstructor, RemainderAlongK) +{ + using Config = StreamKTilePartitionerBaseConfigRemainderAlongK; + + ck_tile::StreamKTilePartitionerBase tile_partitioner{ + Config::M, Config::N, Config::K, Config::MAX_ACTIVE_WGS}; + + StreamKTilePartitionerBaseExpected expected_values{ + 1, 0, 2, 3, 1, 3, 1, 0, 1, Config::MAX_ACTIVE_WGS, Config::N, Config::K, 1}; validate_streamk_base_constructor(expected_values, tile_partitioner); } @@ -567,6 +579,80 @@ TEST(StreamKTilePartitionerBaseGetTileLocalCtaIndex, DP2TileSK) } } +TEST(StreamKTilePartitionerBaseGetKSize, NoRemainderTiles) +{ + // Types + using Config = StreamKTilePartitionerBaseConfigRemainderAlongK; + using TilePartitioner = ck_tile::StreamKTilePartitionerBase; + using Kernel = + KernelWrapperSpecialized; + + // Test parameters + ck_tile::StreamKTilePartitionerBase tile_partitioner{ + Config::M, Config::N, Config::K, Config::MAX_ACTIVE_WGS}; + ck_tile::DeviceMem k_size_dev(sizeof(ck_tile::index_t)); + ck_tile::index_t num_macro_tiles = 2; + ck_tile::index_t local_iter_end = 2; + + // Launch kernel + auto kargs = Kernel::MakeKernelArgs(num_macro_tiles, + local_iter_end, + Config::UNUSED, + k_size_dev.GetDeviceBuffer(), + nullptr, + tile_partitioner); + ck_tile::launch_kernel(ck_tile::stream_config{nullptr, false, 0, 0, 1}, + ck_tile::make_kernel<1>(Kernel{}, 1, 1, 0, kargs)); + + // Validate results + ck_tile::index_t k_size; + k_size_dev.FromDevice(&k_size); + + /* + In the StreamKTilePartitionerBaseConfigRemainderAlongK config, workgroup 0 is assigned the first + 2 macro tile along K. Both of these macro tiles are MPerBlock x KPerBlock. So, the k_size is + K_TILE * 2. (See the struct definition for a detailed diagram.) + */ + EXPECT_EQ(k_size, Config::K_TILE * 2); +} + +TEST(StreamKTilePartitionerBaseGetKSize, RemainderTiles) +{ + // Types + using Config = StreamKTilePartitionerBaseConfigRemainderAlongK; + using TilePartitioner = ck_tile::StreamKTilePartitionerBase; + using Kernel = + KernelWrapperSpecialized; + + // Test parameters + ck_tile::StreamKTilePartitionerBase tile_partitioner{ + Config::M, Config::N, Config::K, Config::MAX_ACTIVE_WGS}; + ck_tile::DeviceMem k_size_dev(sizeof(ck_tile::index_t)); + ck_tile::index_t num_macro_tiles = 1; + ck_tile::index_t local_iter_end = 3; + + // Launch kernel + auto kargs = Kernel::MakeKernelArgs(num_macro_tiles, + local_iter_end, + Config::UNUSED, + k_size_dev.GetDeviceBuffer(), + nullptr, + tile_partitioner); + ck_tile::launch_kernel(ck_tile::stream_config{nullptr, false, 0, 0, 1}, + ck_tile::make_kernel<1>(Kernel{}, 1, 1, 0, kargs)); + + // Validate results + ck_tile::index_t k_size; + k_size_dev.FromDevice(&k_size); + + /* + In the StreamKTilePartitionerBaseConfigRemainderAlongK config, workgroup 1 is assigned the final + macro tile along K. This macro tiles is MPerBlock x (K % K_TILE). So, the k_size is + K % K_TILE. (See the struct definition for a detailed diagram.) + */ + EXPECT_EQ(k_size, Config::K % Config::K_TILE); +} + // Persistent TEST(StreamKTilePartitioner_PersistentConstructor, SKOnly) { diff --git a/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner_common.hpp b/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner_common.hpp index 2276a6b0c3..49b62081c1 100644 --- a/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner_common.hpp +++ b/test/ck_tile/gemm_streamk/test_streamk_tile_partitioner_common.hpp @@ -14,7 +14,8 @@ enum StreamKTilePartitionerBaseMethodId GET_TILE_INDEX, GET_ITER_BOUNDARIES, GET_OUTPUT_TILE_INDEX, - GET_TILE_LOCAL_CTA_INDEX + GET_TILE_LOCAL_CTA_INDEX, + GET_K_SIZE, }; // Base kernel wrapper class to facilitate testing class device functions. @@ -108,6 +109,20 @@ struct KernelWrapperSpecialized +struct KernelWrapperSpecialized + : public KernelWrapper +{ + + using Base = KernelWrapper; + + CK_TILE_DEVICE void operator()(typename Base::KernelArgs kargs) + { + *(static_cast(kargs.result1)) = + kargs.tile_partitioner.get_k_size(kargs.arg1, kargs.arg2); + } +}; + template struct KernelWrapperSpecialized : public KernelWrapper @@ -167,6 +182,8 @@ struct StreamKTilePartitionerBaseExpected ck_tile::index_t num_tiles_; ck_tile::index_t max_active_wgs_; ck_tile::index_t n_; + ck_tile::index_t k_; + ck_tile::index_t remainder_along_k_; }; template @@ -185,6 +202,8 @@ void validate_streamk_base_constructor( EXPECT_EQ(tile_partitioner.get_num_tiles(), expected_values.num_tiles_); EXPECT_EQ(tile_partitioner.get_max_active_wgs(), expected_values.max_active_wgs_); EXPECT_EQ(tile_partitioner.get_n(), expected_values.n_); + EXPECT_EQ(tile_partitioner.get_k(), expected_values.k_); + EXPECT_EQ(tile_partitioner.get_remainder_along_k(), expected_values.remainder_along_k_); } struct StreamKTilePartitionerBaseConfig @@ -318,6 +337,37 @@ struct StreamKTilePartitionerBaseConfigSKOnlyLargeK : public StreamKTilePartitio ck_tile::sequence>; }; +struct StreamKTilePartitionerBaseConfigRemainderAlongK : public StreamKTilePartitionerBaseConfig +{ + /* + Since K % K_Tile <=> 5 % 2 = 1, there will be 2 full macro tiles along K of size MPerBlock x + KPerBlock for A and KPerBlock x NPerBlock for B. The final macro tile along K will be of size + (K % K_Tile) along the K dimension. + + Consider the A tensor as an example: Let R = K % K_TILE + ------------------------------------- + | | | | + MPerBlock | WG0 | WG0 | WG1 | + | | | | + ------------------------------------- + |<-KPerBlock->|<-KPerBlock->|<--R-->| + + */ + + static constexpr ck_tile::index_t M = 2; + static constexpr ck_tile::index_t N = 2; + static constexpr ck_tile::index_t K = 5; + static constexpr ck_tile::index_t MAX_ACTIVE_WGS = 2; + + static constexpr ck_tile::index_t M_TILE = 2; + static constexpr ck_tile::index_t N_TILE = 2; + static constexpr ck_tile::index_t K_TILE = 2; + + using GemmShape = ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; +}; + struct StreamKTilePartitionerBaseConfigEdgeCase : public StreamKTilePartitionerBaseConfig { From e7e8801dc3764c998162cc79764fe44cb84af67a Mon Sep 17 00:00:00 2001 From: Hosang Yoon <156028780+hyoon1@users.noreply.github.com> Date: Sat, 30 May 2026 00:10:26 +0000 Subject: [PATCH 005/143] [rocm-libraries] ROCm/rocm-libraries#7586 (commit c18f2c7) [CK_TILE] Use gfx11 float buffer atomics in FMHA Bwd ## Motivation FlashAttention CK backward on gfx11 can hit out-of-bounds/tail writes in the dQ accumulator atomic-add path when sequence rows are padded at the tile level but not marked invalid in the DQDKDV main tensor view. With the generic global atomic fallback, an incorrectly-valid tail element can issue an actual pointer-based `atomicAdd`. With the buffer atomic path, the write is issued through a buffer resource with bounds information and follows the same backend already used by gfx9/gfx12. This fixes the gfx11 FMHA BWD failure without changing the gfx11 default for unrelated CK Tile kernels. ## Technical Details This PR enables the existing CK Tile AMD buffer float atomic-add path only for generated FMHA BWD gfx11 translation units. gfx11 normally uses the generic global atomic fallback for floating-point `buffer_view::atomic_add`. That fallback performs the atomic through a raw computed pointer and depends on the software validity predicate to avoid invalid elements. In FMHA BWD dQ accumulation, padded tail rows can reach this path, so using the buffer atomic backend is safer: it uses a buffer resource with base pointer, bounds information, and an element offset, matching the backend already used by gfx9/gfx12. Enabling `CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT` globally for gfx11 is too broad and can break unrelated gfx11 CK builds such as GEMM. Instead, `config.hpp` now preserves an explicitly pre-defined `CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT`, while keeping the existing default disabled for gfx11. ## Test Plan Validated the change with the FlashAttention CK full test suite with backward pass enabled on gfx11. pytest -q -s tests/test_flash_attn_ck.py ## Test Result FlashAttention CK gfx11 test result: 260680 passed, 152076 skipped ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. Co-authored-by: Po Yen Chen --- example/ck_tile/01_fmha/CMakeLists.txt | 4 ++-- example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py | 7 +++++++ include/ck_tile/core/config.hpp | 2 ++ 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/example/ck_tile/01_fmha/CMakeLists.txt b/example/ck_tile/01_fmha/CMakeLists.txt index 0650bd3de0..2d44b44996 100644 --- a/example/ck_tile/01_fmha/CMakeLists.txt +++ b/example/ck_tile/01_fmha/CMakeLists.txt @@ -2,8 +2,8 @@ # SPDX-License-Identifier: MIT set(INST_TARGETS ${SUPPORTED_GPU_TARGETS}) -# Currently only gfx9 and gfx12 archs are supported by FMHA -list(FILTER INST_TARGETS INCLUDE REGEX "gfx9|gfx12") +# Currently only gfx9, gfx11, and gfx12 archs are supported by FMHA +list(FILTER INST_TARGETS INCLUDE REGEX "gfx9|gfx1[12]") if(NOT INST_TARGETS) message(WARNING "Skipping Tile Engine FMHA compilation: No supported GPU targets (gfx9, gfx11, gfx12) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") return() diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py index 8079b3d858..dae78e243c 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py @@ -28,6 +28,13 @@ FMHA_BWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.\n // auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && \\ + (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || \\ + defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || \\ + defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__)) +#undef CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT +#define CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT 1 +#endif #include "fmha_bwd.hpp" """ diff --git a/include/ck_tile/core/config.hpp b/include/ck_tile/core/config.hpp index 2656167651..1e76be705f 100644 --- a/include/ck_tile/core/config.hpp +++ b/include/ck_tile/core/config.hpp @@ -173,6 +173,7 @@ #endif // buffer atomic add: floating point +#ifndef CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT #ifndef __HIP_DEVICE_COMPILE__ // for host code #define CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT 1 #elif defined(__gfx9__) || defined(__gfx12__) // for GPU code @@ -180,6 +181,7 @@ #else // for GPU code #define CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT 0 #endif +#endif #if(defined(__gfx90a__) || defined(__gfx94__)) // for GPU code #define CK_TILE_USE_AMD_BUFFER_ATOMIC_MAX_FLOAT64 1 From 8d97265896bfb6b57e3769a8f442bc72468d32b9 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Sat, 30 May 2026 00:15:12 +0000 Subject: [PATCH 006/143] [rocm-libraries] ROCm/rocm-libraries#7863 (commit 0845ce7) [CK] apply the compiler warning suppression flags in cmake files (#7863) ## Motivation Apply the blanket suppression flags for latest clang warnings in staging compiler such as: lifetime-safety-lifetimebound-violation lifetime-safety-intra-tu-suggestions lifetime-safety-cross-tu-suggestions unknown-warning-option ## Technical Details ## Test Plan ## Test Result ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- CMakeLists.txt | 4 ++++ cmake/EnableCompilerWarnings.cmake | 8 ++++++++ dispatcher/CMakeLists.txt | 2 +- 3 files changed, 13 insertions(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ce054255c..1cb4825c85 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -741,6 +741,10 @@ SET(BUILD_DEV ON CACHE BOOL "BUILD_DEV") if(BUILD_DEV) add_compile_options(-Werror) add_compile_options(-Weverything) + add_compile_options(-Wno-lifetime-safety-intra-tu-suggestions) + add_compile_options(-Wno-lifetime-safety-cross-tu-suggestions) + add_compile_options(-Wno-lifetime-safety-lifetimebound-violation) + add_compile_options(-Wno-unknown-warning-option) endif() message(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}") diff --git a/cmake/EnableCompilerWarnings.cmake b/cmake/EnableCompilerWarnings.cmake index 9cc960cc23..2f9a04f485 100644 --- a/cmake/EnableCompilerWarnings.cmake +++ b/cmake/EnableCompilerWarnings.cmake @@ -50,6 +50,10 @@ else() -Wsign-compare -Wno-extra-semi-stmt -Wno-unused-template + -Wno-lifetime-safety-intra-tu-suggestions + -Wno-lifetime-safety-cross-tu-suggestions + -Wno-lifetime-safety-lifetimebound-violation + -Wno-unknown-warning-option ) if (CMAKE_${COMPILER}_COMPILER_ID MATCHES "Clang") list(APPEND CMAKE_COMPILER_WARNINGS @@ -76,6 +80,10 @@ else() -Wno-unsafe-buffer-usage -Wno-unused-lambda-capture -Wno-nvcc-compat + -Wno-lifetime-safety-intra-tu-suggestions + -Wno-lifetime-safety-cross-tu-suggestions + -Wno-lifetime-safety-lifetimebound-violation + -Wno-unknown-warning-option ) if(CK_CXX_STANDARD GREATER_EQUAL 20) list(APPEND CMAKE_COMPILER_WARNINGS -Wno-c++20-compat) diff --git a/dispatcher/CMakeLists.txt b/dispatcher/CMakeLists.txt index ed9b20d33c..79bdde45e8 100644 --- a/dispatcher/CMakeLists.txt +++ b/dispatcher/CMakeLists.txt @@ -59,7 +59,7 @@ endif() # Compiler warnings if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") target_compile_options(ck_tile_dispatcher PRIVATE - -Wall -Wextra -Wpedantic + -Wall -Wextra -Wpedantic -Wno-lifetime-safety-intra-tu-suggestions -Wno-lifetime-safety-cross-tu-suggestions -Wno-lifetime-safety-lifetimebound-violation -Wno-unknown-warning-option ) elseif(CMAKE_CXX_COMPILER_ID MATCHES "MSVC") target_compile_options(ck_tile_dispatcher PRIVATE From 22a99f97e8e7d8eed03e0485035ce5dd9b9695ec Mon Sep 17 00:00:00 2001 From: Tianyuan Wu Date: Sat, 30 May 2026 01:28:48 +0000 Subject: [PATCH 007/143] [rocm-libraries] ROCm/rocm-libraries#7677 (commit 308af93) [CK_Tile] Add scale16 Support for F4 WMMA in CK_Tile ## Motivation This PR adds CK Tile support for the scale16 F4 WMMA path on gfx1250 and improves warp GEMM unit test coverage/structure for gfx1250-specific cases. ## Technical Details - Scale16 support in warp GEMM dispatch and WMMA trait plumbing: added IsScale16 plumbing to warp GEMM dispatcher path - Warp GEMM test restructuring for gfx1250: added Warp GEMM gfx1250 coverage to verify all F4 WMMA paths ## Test Plan Run ./test_ck_tile_wg_32x16x128_fp4. ## Test Result ``` ./test_ck_tile_wg_32x16x128_fp4 [----------] Global test environment tear-down [==========] 3 tests from 1 test suite ran. (1751 ms total) [ PASSED ] 3 tests. ``` ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- include/ck_tile/core/numeric/mxfp_scale.hpp | 94 ++++++++ .../gemm/warp/warp_gemm_attribute_wmma.hpp | 12 +- .../warp/warp_gemm_attribute_wmma_impl.hpp | 20 +- ...p_gemm_attribute_wmma_impl_8bit_traits.hpp | 74 ++++-- .../ops/gemm/warp/warp_gemm_dispatcher.hpp | 48 +++- .../ck_tile/ops/gemm/warp/warp_gemm_impl.hpp | 21 +- .../ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp | 16 +- test/ck_tile/warp_gemm/CMakeLists.txt | 4 + .../warp_gemm/test_f32_16x16x128_fp4.cpp | 192 +-------------- .../warp_gemm/test_f32_32x16x128_fp4.cpp | 39 +++ test/ck_tile/warp_gemm/test_gemm_util.hpp | 223 ++++++++++++++++++ 11 files changed, 508 insertions(+), 235 deletions(-) create mode 100644 test/ck_tile/warp_gemm/test_f32_32x16x128_fp4.cpp create mode 100644 test/ck_tile/warp_gemm/test_gemm_util.hpp diff --git a/include/ck_tile/core/numeric/mxfp_scale.hpp b/include/ck_tile/core/numeric/mxfp_scale.hpp index 1eb8063c02..54687e604f 100644 --- a/include/ck_tile/core/numeric/mxfp_scale.hpp +++ b/include/ck_tile/core/numeric/mxfp_scale.hpp @@ -103,7 +103,101 @@ struct Packed4Scale } }; +template +struct Packed8Scale +{ + using scale_type = ScaleType; + using raw_type = uint64_t; + using raw_scale_type = typename ScaleType::raw_type; + + static constexpr int num_pack = 8; + union + { + raw_type data_; + raw_scale_type scales_[num_pack]; // Direct byte/element access + }; + + // Constructors + CK_TILE_HOST_DEVICE constexpr Packed8Scale() = default; + CK_TILE_HOST_DEVICE constexpr Packed8Scale(raw_type val) : data_(val) {} + CK_TILE_HOST_DEVICE constexpr Packed8Scale( + float s0, float s1, float s2, float s3, float s4, float s5, float s6, float s7) + { + set_scales_from_float(s0, s1, s2, s3, s4, s5, s6, s7); + } + + CK_TILE_HOST_DEVICE constexpr Packed8Scale(ScaleType s0, + ScaleType s1, + ScaleType s2, + ScaleType s3, + ScaleType s4, + ScaleType s5, + ScaleType s6, + ScaleType s7) + { + set_scales(s0, s1, s2, s3, s4, s5, s6, s7); + } + + CK_TILE_HOST_DEVICE constexpr void set_scales_from_float( + float s0, float s1, float s2, float s3, float s4, float s5, float s6, float s7) + { + set_scales(ScaleType(s0), + ScaleType(s1), + ScaleType(s2), + ScaleType(s3), + ScaleType(s4), + ScaleType(s5), + ScaleType(s6), + ScaleType(s7)); + } + + CK_TILE_HOST_DEVICE constexpr void set_scales(ScaleType s0, + ScaleType s1, + ScaleType s2, + ScaleType s3, + ScaleType s4, + ScaleType s5, + ScaleType s6, + ScaleType s7) + { + data_ = 0; + pack_scale(s0, 7); + pack_scale(s1, 6); + pack_scale(s2, 5); + pack_scale(s3, 4); + pack_scale(s4, 3); + pack_scale(s5, 2); + pack_scale(s6, 1); + pack_scale(s7, 0); + } + + CK_TILE_HOST_DEVICE constexpr operator raw_type() const { return data_; } + CK_TILE_HOST_DEVICE constexpr raw_type& data() [[clang::lifetimebound]] { return data_; } + CK_TILE_HOST_DEVICE constexpr raw_type data() const { return data_; } + + CK_TILE_HOST_DEVICE constexpr float unpack_to_float(int i) const + { + return static_cast(unpack_scale(i)); + } + + CK_TILE_HOST_DEVICE constexpr ScaleType unpack_scale(int i) const + { + return ScaleType(scales_[i]); + } + + CK_TILE_HOST_DEVICE constexpr void pack_from_float(float scale, int i) + { + pack_scale(ScaleType(scale), i); + } + + CK_TILE_HOST_DEVICE constexpr void pack_scale(ScaleType scale, int i) + { + scales_[i] = scale.get(); + } +}; + // Type alias for e8m0_t scales using Packed4Scale_E8M0 = Packed4Scale; +using Packed8Scale_E8M0 = Packed8Scale; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma.hpp index 8cbaa9bfc8..9947915cbe 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma.hpp @@ -234,12 +234,12 @@ struct WarpGemmAttributeWmma } // c_vec += a_vec * b_vec - template + template CK_TILE_DEVICE void operator()(CVecType& c_vec, const AVecType& a_vec, - const int32_t& a_scale, + const AScaleType& a_scale, const BVecType& b_vec, - const int32_t& b_scale) const + const BScaleType& b_scale) const { if constexpr(kTransC) { @@ -253,11 +253,11 @@ struct WarpGemmAttributeWmma } // c_vec = a_vec * b_vec - template + template CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, - const int32_t& a_scale, + const AScaleType& a_scale, const BVecType& b_vec, - const int32_t& b_scale) const + const BScaleType& b_scale) const { if constexpr(kTransC) { diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp index 8aa02aba6e..8fd185cb42 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp @@ -19,6 +19,11 @@ template struct WmmaTraits; +// Tag used to select scale16 WMMA traits specializations. +struct WmmaScale16Tag +{ +}; + // Generic WMMA implementation using traits template struct WarpGemmAttributeWmmaImpl @@ -88,22 +93,22 @@ struct WarpGemmAttributeWmmaImpl Traits::template wmma_intrinsic(a_vec, b_vec, CVecType{0.f})); } - template + template CK_TILE_DEVICE void operator()(CVecType& c_vec, const AVecType& a_vec, - const int32_t& a_scale, + const AScaleType& a_scale, const BVecType& b_vec, - const int32_t& b_scale) const + const BScaleType& b_scale) const { c_vec = Traits::template wmma_intrinsic(a_vec, a_scale, b_vec, b_scale, c_vec); } // c_vec = a_vec * b_vec - template + template CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, - const int32_t& a_scale, + const AScaleType& a_scale, const BVecType& b_vec, - const int32_t& b_scale) const + const BScaleType& b_scale) const { return bit_cast(Traits::template wmma_intrinsic( a_vec, a_scale, b_vec, b_scale, CVecType{0.f})); @@ -177,6 +182,9 @@ using WarpGemmAttributeWmmaImpl_f32_32x16x128_f4 = using WarpGemmAttributeWmmaImpl_f32_32x32x128_f4 = WarpGemmAttributeWmmaImpl>; +using WarpGemmAttributeWmmaImpl_f32_32x32x128_f4_scale16 = WarpGemmAttributeWmmaImpl< + WmmaTraits>; + using WarpGemmAttributeWmmaImpl_f16_16x16x64_f8_f8 = WarpGemmAttributeWmmaImpl>; diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp index dcc40304f4..77dafd0956 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp @@ -6,6 +6,9 @@ #include "warp_gemm_attribute_wmma_impl_base_traits.hpp" #include "warp_gemm_params.hpp" namespace ck_tile { + +struct WmmaScale16Tag; + // int8 specialization - GFX11 template <> struct WmmaTraits @@ -528,17 +531,18 @@ struct WmmaTraits } }; -template <> -struct WmmaTraits +template +struct WmmaTraitsGfx125PkFp4F32_32x32x128 : WmmaTraitsBase { - using ArchType = gfx125_t; + using ArchType = gfx125_t; + using ScaleType = std::conditional_t; template CK_TILE_DEVICE static CVecType wmma_intrinsic(const AVecType& a_vec, - const int32_t& a_scale, + const ScaleType& a_scale, const BVecType& b_vec, - const int32_t& b_scale, + const ScaleType& b_scale, const CVecType& c_vec) { #ifdef __gfx125__ @@ -569,19 +573,38 @@ struct WmmaTraits const auto& b_slice = b_buffer.template get_as()[n]; auto& c_slice = c_result.template get_as()[n]; - c_slice = __builtin_amdgcn_wmma_scale_f32_32x16x128_f4( - bit_cast(a_slice), - bit_cast(b_slice), - 0, - c_slice, - 1, // OPSEL[0] - fixed to 1 for F4 - P::scale_a, // OPSEL_HI[0] - scale data type for A - a_scale, - n.value, // OPSEL[1] - select B scale (iterates over N blocks) - P::scale_b, // OPSEL_HI[1] - scale data type for B - b_scale, - 0, // NEG - 0); // NEG_HI + if constexpr(IsScale16) + { + c_slice = __builtin_amdgcn_wmma_scale16_f32_32x16x128_f4( + bit_cast(a_slice), + bit_cast(b_slice), + 0, + c_slice, + 1, // OPSEL[0] - fixed to 1 for F4 + P::scale_a, // OPSEL_HI[0] - scale data type for A + a_scale, + n.value, // OPSEL[1] - select B scale (iterates over N blocks) + P::scale_b, // OPSEL_HI[1] - scale data type for B + b_scale, + 0, // NEG + 0); // NEG_HI + } + else + { + c_slice = __builtin_amdgcn_wmma_scale_f32_32x16x128_f4( + bit_cast(a_slice), + bit_cast(b_slice), + 0, + c_slice, + 1, // OPSEL[0] - fixed to 1 for F4 + P::scale_a, // OPSEL_HI[0] - scale data type for A + a_scale, + n.value, // OPSEL[1] - select B scale (iterates over N blocks) + P::scale_b, // OPSEL_HI[1] - scale data type for B + b_scale, + 0, // NEG + 0); // NEG_HI + } }); return bit_cast(c_result); @@ -602,7 +625,8 @@ struct WmmaTraits #ifdef __gfx125__ // Pass default scale values 1.0f Packed4Scale_E8M0 pkscale(1.0f, 1.0f, 1.0f, 1.0f); - return wmma_intrinsic(a_vec, pkscale, b_vec, pkscale, c_vec); + const auto default_scale = static_cast(pkscale); + return wmma_intrinsic(a_vec, default_scale, b_vec, default_scale, c_vec); #else ck_tile::ignore = a_vec; ck_tile::ignore = b_vec; @@ -612,6 +636,18 @@ struct WmmaTraits } }; +template <> +struct WmmaTraits + : WmmaTraitsGfx125PkFp4F32_32x32x128 +{ +}; + +template <> +struct WmmaTraits + : WmmaTraitsGfx125PkFp4F32_32x32x128 +{ +}; + // f8f6f4 specialization - GFX125 enum F8F6F4OpDataTypeEnum { diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp index 4027b8ed34..2e6fa605ba 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp @@ -33,6 +33,7 @@ template struct Dispatcher; @@ -178,10 +179,10 @@ template<> struct Dispatcher { using Ty #if !defined(__gfx125__) // scale mfma based f8f6f4 -template -struct Dispatcher> { using Type = WarpGemmMfma_f32_16x16x128_f8f6f4; }; -template -struct Dispatcher> { using Type = WarpGemmMfma_f32_16x16x128_f8f6f4_CTransposed; }; +template +struct Dispatcher> { using Type = WarpGemmMfma_f32_16x16x128_f8f6f4; }; +template +struct Dispatcher> { using Type = WarpGemmMfma_f32_16x16x128_f8f6f4_CTransposed; }; #endif template<> struct Dispatcher { using Type = WarpGemmMfma_f32_32x32x64_fp8_fp8<>; }; @@ -224,7 +225,7 @@ template struct Dispatcher struct Dispatcher : WmmaTag { using Type = WarpGemmWmma_f32_16x16x64_bf8_f8; }; template struct Dispatcher : WmmaTag { using Type = WarpGemmWmma_f32_32x16x128_f4; }; -template struct Dispatcher : WmmaTag { using Type = WarpGemmWmma_f32_32x32x128_f4; }; +template struct Dispatcher : WmmaTag { using Type = WarpGemmWmma_f32_32x32x128_f4; }; #if defined(__gfx125__) template struct Dispatcher : WmmaTag { using Type = WarpGemmWmma_f32_16x16x64_f8_f8; }; @@ -244,8 +245,27 @@ template<> struct Dispatcher { using Typ template<> struct Dispatcher { using Type = WarpGemmMfma_f32_16x16x64_bf8_bf8_CTransposed; }; #endif -template -struct Dispatcher : WmmaTag { using Type = WarpGemmWmma_f32_32x32x128_f8f6f4; }; +template +struct Dispatcher : WmmaTag +{ + using Type = WarpGemmWmma_f32_32x32x128_f8f6f4; +}; template struct Dispatcher : WmmaTag { using Type =WarpGemmWmma_f16_16x16x64_f8_f8; }; template struct Dispatcher : WmmaTag { using Type =WarpGemmWmma_f16_16x16x64_bf8_bf8; }; @@ -265,12 +285,12 @@ template struct Dispatcher + bool TransposeC, bool SA, bool SS, bool IsScale16> struct Dispatcher>>> - : Dispatcher {}; + Dispatcher>>> + : Dispatcher {}; // clang-format on } // namespace warp_gemm_dispatcher @@ -286,7 +306,8 @@ template + WGAttrNumAccessEnum AttrNumAccessB = AttrNumAccessA, + bool IsScale16 = false> using WarpGemmDispatcher = typename impl::warp_gemm_dispatcher::Dispatcher< // AType, BType, @@ -298,6 +319,7 @@ using WarpGemmDispatcher = typename impl::warp_gemm_dispatcher::Dispatcher< // SwizzleA, UseStructuredSparsity, AttrNumAccessA, - AttrNumAccessB>::Type; + AttrNumAccessB, + IsScale16>::Type; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_impl.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_impl.hpp index f0353672a0..6801f627c7 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_impl.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_impl.hpp @@ -90,12 +90,17 @@ struct WarpGemmImpl c.get_thread_buffer().template set_as(I0, c_vec); } - template + template CK_TILE_DEVICE void operator()(CTensor& c, const ATensor& a, const BTensor& b, - const int32_t& a_scale, - const int32_t& b_scale) const + const AScaleType& a_scale, + const BScaleType& b_scale) const { static_assert(detail::is_similiar_distributed_tensor_v && detail::is_similiar_distributed_tensor_v && @@ -141,11 +146,15 @@ struct WarpGemmImpl return c; } - template + template CK_TILE_DEVICE auto operator()(const ATensor& a, const BTensor& b, - const int32_t& a_scale, - const int32_t& b_scale) const + const AScaleType& a_scale, + const BScaleType& b_scale) const { using CTensor = CWarpTensor; static_assert(detail::is_similiar_distributed_tensor_v && diff --git a/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp b/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp index e7b601306f..1c522d07c1 100644 --- a/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp @@ -187,12 +187,16 @@ using WarpGemmWmma_f32_32x16x128_f4 = AttrNumAccess, AttrNumAccess>>; -template -using WarpGemmWmma_f32_32x32x128_f4 = - WarpGemmImpl>; +template +using WarpGemmWmma_f32_32x32x128_f4 = WarpGemmImpl< + WarpGemmAttributeWmma, + kTransC, + AttrNumAccess, + AttrNumAccess>>; template -#include "ck_tile/host.hpp" -#include "ck_tile/host/kernel_launch.hpp" -#include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp" - -using namespace ck_tile; - -template -struct WGDispCase -{ - using AType = A; - using BType = B; - using AccType = Acc; - static constexpr index_t MPerWave = M; - static constexpr index_t NPerWave = N; - static constexpr index_t KPerWave = K; - static constexpr bool kTransposeC = TransposeC; - static constexpr bool kSwizzleA = SwizzleA; - static constexpr bool kUSS = UseStructuredSparsity; - static constexpr WGAttrNumAccessEnum kNA = NA; -}; +#include "test_gemm_util.hpp" +#include "gtest/gtest.h" using WGDispatcherTypesList = - ::testing::Types>; - -template -struct WarpGemmKernel -{ - static constexpr int kBlockSize = 64; - __device__ void operator()(void* A, void* B, void* C, void* ScaleA, void* ScaleB) const - { - using WarpGemm = ck_tile::WarpGemmDispatcher; - // A: [M,K] row-major (packed) - const auto a_view = ck_tile::make_naive_tensor_view( - static_cast(A), - ck_tile::make_tuple(M, K), - ck_tile::make_tuple(K, ck_tile::number<1>{}), - ck_tile::number{}, - ck_tile::number<1>{}); - // B: expose as logical [N,K] with strides (1, N) over the original row-major [K,N] buffer - const auto b_view = ck_tile::make_naive_tensor_view( - static_cast(B), - ck_tile::make_tuple(N, K), - ck_tile::make_tuple(K, ck_tile::number<1>{}), - ck_tile::number{}, - ck_tile::number<1>{}); - // C: [M,N] row-major (packed) - const auto c_view = ck_tile::make_naive_tensor_view( - static_cast(C), - ck_tile::make_tuple(M, N), - ck_tile::make_tuple(N, ck_tile::number<1>{}), - ck_tile::number{}, - ck_tile::number<1>{}); - - using AWarpTensor = typename WarpGemm::AWarpTensor; - using BWarpTensor = typename WarpGemm::BWarpTensor; - using CWarpTensor = typename WarpGemm::CWarpTensor; - - constexpr auto a_len = AWarpTensor::get_tile_distribution().get_lengths(); - constexpr auto b_len = BWarpTensor::get_tile_distribution().get_lengths(); - constexpr auto c_len = CWarpTensor::get_tile_distribution().get_lengths(); - - auto a_win = ck_tile::make_tile_window( - a_view, a_len, ck_tile::make_multi_index(0, 0), AWarpTensor::get_tile_distribution()); - auto b_win = ck_tile::make_tile_window( - b_view, b_len, ck_tile::make_multi_index(0, 0), BWarpTensor::get_tile_distribution()); - auto c_win = ck_tile::make_tile_window( - c_view, c_len, ck_tile::make_multi_index(0, 0), CWarpTensor::get_tile_distribution()); - - AWarpTensor a_tile; - BWarpTensor b_tile; - ck_tile::load_tile(a_tile, a_win); - ck_tile::load_tile(b_tile, b_win); - - auto scale_a = static_cast(static_cast(ScaleA)[0].get()); - auto scale_b = static_cast(static_cast(ScaleB)[0].get()); - - auto c_tile = - WarpGemm{}.template operator(), OpSelB<0>>(a_tile, b_tile, scale_a, scale_b); - - ck_tile::store_tile(c_win, c_tile); - } -}; - -template -static void RunWarpGemmCase(const ck_tile::HostTensor& A, - const ck_tile::HostTensor& B, - const ck_tile::HostTensor& ScaleA, - const ck_tile::HostTensor& ScaleB, - ck_tile::HostTensor& C) -{ - ck_tile::DeviceMem Ad(A), Bd(B), Cd(C), SAd(ScaleA), SBd(ScaleB); - dim3 grid(1), block{64}; - - using Kernel = WarpGemmKernel; - - (void)ck_tile::launch_kernel(ck_tile::stream_config{nullptr, true, 0, 0, 1}, - ck_tile::make_kernel(Kernel{}, - grid, - block, - 0, - Ad.GetDeviceBuffer(), - Bd.GetDeviceBuffer(), - Cd.GetDeviceBuffer(), - SAd.GetDeviceBuffer(), - SBd.GetDeviceBuffer())); - - Cd.FromDevice(C.mData.data()); -} - -template + ::testing::Types>; + +template class WGRuntimeTest : public ::testing::Test { }; @@ -156,38 +22,6 @@ TYPED_TEST_SUITE(WGRuntimeTest, WGDispatcherTypesList); TYPED_TEST(WGRuntimeTest, Compare_Dispatcher_MakeWG) { - using Case = TypeParam; - - using AType = typename Case::AType; - using BType = typename Case::BType; - using CType = typename Case::AccType; - using ck_tile::e8m0_t; - - constexpr index_t M = Case::MPerWave; - constexpr index_t N = Case::NPerWave; - constexpr index_t K = Case::KPerWave; - - auto ScaleA = e8m0_t{2.f}; - auto ScaleB = e8m0_t{4.f}; - - ck_tile::HostTensor A({M, K}); - ck_tile::HostTensor B({N, K}); - ck_tile::HostTensor C({M, N}); - ck_tile::HostTensor sA({M, 1}); - ck_tile::HostTensor sB({N, 1}); - - ck_tile::FillUniformDistribution{-5.f, 5.f}(A); - ck_tile::FillUniformDistribution{-5.f, 5.f}(B); - C.SetZero(); - ck_tile::FillConstant{ScaleA}(sA); - ck_tile::FillConstant{ScaleB}(sB); - - RunWarpGemmCase(A, B, sA, sB, C); - - ck_tile::HostTensor C_ref({M, N}); - C_ref.SetZero(); - ck_tile::reference_mx_gemm( - A, B.transpose(), C_ref, sA, sB.transpose()); - - EXPECT_TRUE(ck_tile::check_err(C, C_ref, "Warp gemm result error.")); + ck_tile::test::warp_gemm:: + RunCompareDispatcherAndReference(); } diff --git a/test/ck_tile/warp_gemm/test_f32_32x16x128_fp4.cpp b/test/ck_tile/warp_gemm/test_f32_32x16x128_fp4.cpp new file mode 100644 index 0000000000..be394ee3a3 --- /dev/null +++ b/test/ck_tile/warp_gemm/test_f32_32x16x128_fp4.cpp @@ -0,0 +1,39 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_gemm_util.hpp" +#include "gtest/gtest.h" + +using WGDispatcherTypesList = + ::testing::Types>; + +template +class WGRuntimeTest : public ::testing::Test +{ +}; + +TYPED_TEST_SUITE(WGRuntimeTest, WGDispatcherTypesList); + +TYPED_TEST(WGRuntimeTest, Compare_Dispatcher_MakeWG_NonScaled) +{ + ck_tile::test::warp_gemm:: + RunCompareDispatcherAndReference(); +} + +TYPED_TEST(WGRuntimeTest, Compare_Dispatcher_MakeWG_Scale16) +{ + ck_tile::test::warp_gemm:: + RunCompareDispatcherAndReference(); +} + +TYPED_TEST(WGRuntimeTest, Compare_Dispatcher_MakeWG_Scale32) +{ + ck_tile::test::warp_gemm:: + RunCompareDispatcherAndReference(); +} diff --git a/test/ck_tile/warp_gemm/test_gemm_util.hpp b/test/ck_tile/warp_gemm/test_gemm_util.hpp new file mode 100644 index 0000000000..dcd1d3f342 --- /dev/null +++ b/test/ck_tile/warp_gemm/test_gemm_util.hpp @@ -0,0 +1,223 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/host.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/core/numeric/mxfp_scale.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp" + +namespace ck_tile::test::warp_gemm { + +template +struct WGDispCase +{ + using AType = A; + using BType = B; + using AccType = Acc; + static constexpr bool kTransposeC = TransposeC; + static constexpr bool kSwizzleA = SwizzleA; + static constexpr bool kUSS = UseStructuredSparsity; + static constexpr WGAttrNumAccessEnum kNA = NA; +}; + +template +struct WarpGemmKernel +{ + static constexpr int kBlockSize = 64; + __device__ void operator()(void* A, void* B, void* C, void* ScaleA, void* ScaleB) const + { + using WarpGemm = ck_tile::WarpGemmDispatcher; + + const auto a_view = ck_tile::make_naive_tensor_view( + static_cast(A), + ck_tile::make_tuple(MPerWave, KPerWave), + ck_tile::make_tuple(KPerWave, ck_tile::number<1>{}), + ck_tile::number{}, + ck_tile::number<1>{}); + const auto b_view = ck_tile::make_naive_tensor_view( + static_cast(B), + ck_tile::make_tuple(NPerWave, KPerWave), + ck_tile::make_tuple(KPerWave, ck_tile::number<1>{}), + ck_tile::number{}, + ck_tile::number<1>{}); + const auto c_view = ck_tile::make_naive_tensor_view( + static_cast(C), + ck_tile::make_tuple(MPerWave, NPerWave), + ck_tile::make_tuple(NPerWave, ck_tile::number<1>{}), + ck_tile::number{}, + ck_tile::number<1>{}); + + using AWarpTensor = typename WarpGemm::AWarpTensor; + using BWarpTensor = typename WarpGemm::BWarpTensor; + using CWarpTensor = typename WarpGemm::CWarpTensor; + + constexpr auto a_len = AWarpTensor::get_tile_distribution().get_lengths(); + constexpr auto b_len = BWarpTensor::get_tile_distribution().get_lengths(); + constexpr auto c_len = CWarpTensor::get_tile_distribution().get_lengths(); + + auto a_win = ck_tile::make_tile_window( + a_view, a_len, ck_tile::make_multi_index(0, 0), AWarpTensor::get_tile_distribution()); + auto b_win = ck_tile::make_tile_window( + b_view, b_len, ck_tile::make_multi_index(0, 0), BWarpTensor::get_tile_distribution()); + auto c_win = ck_tile::make_tile_window( + c_view, c_len, ck_tile::make_multi_index(0, 0), CWarpTensor::get_tile_distribution()); + + AWarpTensor a_tile; + BWarpTensor b_tile; + ck_tile::load_tile(a_tile, a_win); + ck_tile::load_tile(b_tile, b_win); + + const auto c_tile = [&]() { + if constexpr(UseScale) + { + using ScaleType = std::conditional_t; + const auto scale_a = static_cast(ScaleA)[0]; + const auto scale_b = static_cast(ScaleB)[0]; + const auto packed_scale_a = [&]() -> ScaleType { + if constexpr(IsScale16) + { + Packed8Scale_E8M0 pkscale( + scale_a, scale_a, scale_a, scale_a, scale_a, scale_a, scale_a, scale_a); + return static_cast(pkscale); + } + else + { + Packed4Scale_E8M0 pkscale(scale_a, scale_a, scale_a, scale_a); + return static_cast(pkscale); + } + }(); + const auto packed_scale_b = [&]() -> ScaleType { + if constexpr(IsScale16) + { + Packed8Scale_E8M0 pkscale( + scale_b, scale_b, scale_b, scale_b, scale_b, scale_b, scale_b, scale_b); + return static_cast(pkscale); + } + else + { + Packed4Scale_E8M0 pkscale(scale_b, scale_b, scale_b, scale_b); + return static_cast(pkscale); + } + }(); + return WarpGemm{}.template operator(), OpSelB<0>>( + a_tile, b_tile, packed_scale_a, packed_scale_b); + } + else + { + ck_tile::ignore = ScaleA; + ck_tile::ignore = ScaleB; + return WarpGemm{}.template operator(), OpSelB<0>>(a_tile, b_tile); + } + }(); + + ck_tile::store_tile(c_win, c_tile); + } +}; + +template +void RunWarpGemmCase(const ck_tile::HostTensor& A, + const ck_tile::HostTensor& B, + const ck_tile::HostTensor& ScaleA, + const ck_tile::HostTensor& ScaleB, + ck_tile::HostTensor& C) +{ + ck_tile::DeviceMem Ad(A), Bd(B), Cd(C), SAd(ScaleA), SBd(ScaleB); + dim3 grid(1), block{64}; + + (void)ck_tile::launch_kernel( + ck_tile::stream_config{nullptr, true, 0, 0, 1}, + ck_tile::make_kernel( + WarpGemmKernel{}, + grid, + block, + 0, + Ad.GetDeviceBuffer(), + Bd.GetDeviceBuffer(), + Cd.GetDeviceBuffer(), + SAd.GetDeviceBuffer(), + SBd.GetDeviceBuffer())); + + Cd.FromDevice(C.mData.data()); +} + +template +void RunCompareDispatcherAndReference() +{ + using AType = typename Case::AType; + using BType = typename Case::BType; + using CType = typename Case::AccType; + + constexpr index_t M = MPerWave; + constexpr index_t N = NPerWave; + constexpr index_t K = KPerWave; + + const auto ScaleA = ck_tile::e8m0_t{2.f}; + const auto ScaleB = ck_tile::e8m0_t{4.f}; + + ck_tile::HostTensor A({M, K}); + ck_tile::HostTensor B({N, K}); + ck_tile::HostTensor C({M, N}); + ck_tile::HostTensor sA({M, 1}); + ck_tile::HostTensor sB({N, 1}); + + ck_tile::FillUniformDistribution{-5.f, 5.f}(A); + ck_tile::FillUniformDistribution{-5.f, 5.f}(B); + C.SetZero(); + ck_tile::FillConstant{ScaleA}(sA); + ck_tile::FillConstant{ScaleB}(sB); + + RunWarpGemmCase(A, B, sA, sB, C); + + ck_tile::HostTensor C_ref({M, N}); + C_ref.SetZero(); + + if constexpr(UseScale) + { + ck_tile::reference_mx_gemm( + A, B.transpose(), C_ref, sA, sB.transpose()); + } + else + { + ck_tile::reference_gemm(A, B.transpose(), C_ref); + } + + EXPECT_TRUE(ck_tile::check_err(C, C_ref, "Warp gemm result error.")); +} + +} // namespace ck_tile::test::warp_gemm From c56c6750d0fc54ed771d532cc92c316423449614 Mon Sep 17 00:00:00 2001 From: Chao Date: Sat, 30 May 2026 10:34:06 +0000 Subject: [PATCH 008/143] [rocm-libraries] ROCm/rocm-libraries#6498 (commit 5961a2e) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [CK_TILE] Fix conditional rescale numerical instability in FMHA forward (#6498) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [CK_TILE] Fix conditional rescale numerical instability in FMHA forward ## Motivation Fix numerical instability in the conditional O-accumulator rescaling optimization for CK-Tile FMHA forward (FlashAttention-4, Algorithm 6, Eq. 6). The conditional rescale optimization skips the expensive O-accumulator rescale when the running row-max shift is within a threshold (tau = log2(256) = 8.0). The original implementation had a bug: attention weights P were computed in the `m_new` reference frame before the skip/rescale decision. In the skip branch, `m` was reverted to `m_old`, but P remained in the `m_new` frame, causing incorrect softmax normalization. This fix introduces a `p_row_correction` factor: in the skip branch, P is multiplied by `exp2(m_new - m_old)` to bring it back to the `m_old` reference frame. - **Correctness:** Fixes broken inference on long sequences where running-max drift causes exp2 overflow (observed as degraded image quality on MI350X Flux2 generation) - **Performance:** Neutral to +4% depending on workload shape ## Technical Details 6 pipeline header files (same pattern in each): - `block_fmha_pipeline_qr_ks_vs.hpp` - `block_fmha_pipeline_qr_ks_vs_async.hpp` - `block_fmha_pipeline_qr_ks_vs_async_trload.hpp` - `block_fmha_pipeline_qr_ks_vs_fp8.hpp` - `block_fmha_pipeline_qr_ks_vs_whole_k_prefetch.hpp` - `block_fmha_pipeline_qs_ks_vs.hpp` In each file: - Lower threshold from 10.0 to 8.0 (tau = log2(256)) - Add `p_row_correction` distributed tensor initialized to 1.0 - Rescale branch: standard rescale of O_acc and l; correction = 1.0 - Skip branch: compute correction = exp2(-acc_scale_log2), update l, revert m, store correction - New `p_spans` sweep applies per-row correction to `p_compute` before P*V GEMM - Move P-to-PDataType cast to after correction sweep ## Dependencies None — this PR is standalone. ## Test Plan - GPU validation on MI300X (gfx942, ROCm 6.4.1): - Command: `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096 -d=128 -prec=bf16 -v=1 -warmup=1 -repeat=3` - GPU validation on MI350X (gfx950, ROCm 7.0): - Command: `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096 -d=128 -prec=bf16 -v=1 -warmup=1 -repeat=3` - Command: `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096 -d=128 -prec=fp16 -v=1 -warmup=1 -repeat=3` ## Test Result Accuracy vs FP32 reference (MI350X, gfx950): | Shape | max_diff | mean_diff | |-------|----------|-----------| | B=1 H=24 M=4096 K=128 bf16 | 9.1e-4 | 4.6e-5 | | B=4 H=32 M=4096 K=128 bf16 | 9.9e-4 | 4.6e-5 | | B=1 H=24 M=4096 K=128 fp16 | 1.2e-4 | 9.0e-6 | Performance (MI350X, gfx950, ROCm 7.0): | Shape | FA4 (TFlops) | Always-rescale (TFlops) | Delta | |-------|-------------|------------------------|-------| | B=1 H=24 M=4096 K=128 bf16 | 425.9 | 428.5 | neutral | | B=2 H=8 M=2048 K=256 bf16 | 513.9 | 509.0 | +1.0% | | B=1 H=64 M=2048 K=64 bf16 | 481.7 | 464.3 | +3.7% | Benchmark results (MI300X, gfx942, ROCm 6.4.1): No regression on MI300X. This correctness fix is performance-neutral. | Config | TFlops / GB/s | Time (ms) | |--------|-------------|-----------| | MHA bf16 b=2 h=8 s=4096 d=128 | 342.49 TFlops | 0.401 | | MHA fp16 b=2 h=8 s=4096 d=128 | 391.70 TFlops | 0.351 | | Causal MHA bf16 b=2 h=8 s=4096 d=128 | 227.07 TFlops | 0.303 | | GQA 4:1 bf16 b=2 h=32 hk=8 s=2048 d=128 | 324.69 TFlops | 0.423 | | GQA 8:1 bf16 b=2 h=64 hk=8 s=2048 d=128 | 348.09 TFlops | 0.790 | | LLaMA-70B prefill b=1 h=64 hk=8 s=4096 d=128 bf16 | 376.71 TFlops | 1.459 | | Long-seq bf16 b=1 h=16 s=16384 d=128 | 383.42 TFlops | 5.735 | | Decode b=64 h=32 hk=8 s_k=4096 d=128 bf16 | 691.64 GB/s | 1.554 | All validation tests pass (`valid:y`) on both MI300X and MI350X. Additional validation: - Uniform scores: softmax output matches FP32 reference (max_diff < 1e-3) - Large seqlen (4096+): no overflow or NaN in O-accumulator - Spike pattern: correct handling of sudden row-max jumps - Multiple spikes: correction applied correctly across multiple skip/rescale transitions - Deterministic: identical outputs across repeated runs - No performance regression on standard workloads --- .../pipeline/block_fmha_pipeline_qr_ks_vs.hpp | 135 ++++++++--- .../block_fmha_pipeline_qr_ks_vs_async.hpp | 123 +++++++--- ...ck_fmha_pipeline_qr_ks_vs_async_trload.hpp | 212 ++++++++++++------ .../block_fmha_pipeline_qr_ks_vs_fp8.hpp | 105 ++++++--- .../pipeline/block_fmha_pipeline_qs_ks_vs.hpp | 140 ++++++++---- 5 files changed, 514 insertions(+), 201 deletions(-) diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp index 6f33fc48a8..15e6e5eb43 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp @@ -851,6 +851,91 @@ struct BlockFmhaPipelineQRKSVS } }; + // Conditional rescaling: skip o_acc rescale when correction factor + // exp2(acc_scale_log2) is negligible (< exp2(-8) ≈ 0.004, below BF16 + // precision). Adapted from FlashAttention-4 (Tri Dao, 2025). + // Eliminates 70-90% of rescale operations in practice. + // + // For skip rows we stabilize P with m_old (the previous max) instead of + // the new max m_j, so P is computed directly in the m_{j-1} frame and no + // post-correction sweep is needed. For rescale rows we use m_j as usual. + // FP8 quant modes (PERTENSOR/BLOCKSCALE/etc.) cast P to FP8 after + // softmax. In the skip branch P is computed with m_old, so P can + // exceed the FP8 representable range and saturate, corrupting the + // P*V GEMM. Disable skip for all FP8 paths (threshold 0). + static constexpr SMPLComputeDataType kRescaleThreshold = + type_convert( + QScaleEnum == BlockAttentionQuantScaleEnum::NO_SCALE ? 8.0f : 0.0f); + + // Per-row stabilizer: m_old for skip rows, m_j for rescale rows. + auto m_stab = + make_static_distributed_tensor(m.get_tile_distribution()); + // Per-row rescale factor (exp2 of acc_scale_log2); only valid when + // needs_rescale[i] is true. + auto rescale_factor = + make_static_distributed_tensor(m.get_tile_distribution()); + auto needs_rescale = make_static_distributed_tensor(m.get_tile_distribution()); + set_tile(needs_rescale, false); + + constexpr auto m_spans = decltype(m)::get_distributed_spans(); + sweep_tile_span(m_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); +#if CK_TILE_FMHA_FWD_FAST_EXP2 + const auto acc_scale_log2 = [&]() { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + if constexpr(kHasLogitsSoftCap) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + auto row_max = scale_s * get_validated_m(m[i_idx]); + return scale_s * m_old[i_idx] - row_max; + } + } + }(); + + const bool need_rescale = + (acc_scale_log2 < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp2(acc_scale_log2); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + // Skip branch: stabilize P with m_old so P is already in + // m_{j-1} frame; restore m to m_old for downstream iterations. + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } +#else + const auto diff = m_old[i_idx] - get_validated_m(m[i_idx]); + const bool need_rescale = + (diff < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp(diff); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } +#endif + }); + constexpr auto p_spans = decltype(p_compute)::get_distributed_spans(); sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); @@ -858,7 +943,7 @@ struct BlockFmhaPipelineQRKSVS // For BLOCKSCALE: precompute (m - shift) once per row // Bias/Alibi/SoftCap: exp2(s - m + shift) = exp2(s - (m - shift)) // else: exp2(scale_s*s - scale_s*m + shift) = exp2(scale_s*s - (scale_s*m - shift)) - auto validated_m = get_validated_m(m[i_idx]); + auto validated_m = get_validated_m(m_stab[i_idx]); auto row_max = scale_s * validated_m; if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE) { @@ -891,7 +976,7 @@ struct BlockFmhaPipelineQRKSVS } } #else - p_compute(i_j_idx) = exp(s[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = exp(s[i_j_idx] - get_validated_m(m_stab[i_idx])); #endif }); }); @@ -904,38 +989,20 @@ struct BlockFmhaPipelineQRKSVS constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); sweep_tile_span(o_spans[number<0>{}], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); -#if CK_TILE_FMHA_FWD_FAST_EXP2 - const auto tmp = [&]() { - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || - BiasEnum == BlockAttentionBiasEnum::ALIBI) - { - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - if constexpr(kHasLogitsSoftCap) - { - - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - auto row_max = scale_s * get_validated_m(m[i_idx]); - return exp2(scale_s * m_old[i_idx] - row_max); - } - } - }(); -#else - const auto tmp = exp(m_old[i_idx] - get_validated_m(m[i_idx])); -#endif - l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; - sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - // FIXME: this use different equation from FA v2 paper, - // but produce correc result. - // Is the equation wrong? - o_acc(i_j_idx) *= tmp; - }); + if(needs_rescale[i_idx]) + { + const auto tmp = rescale_factor[i_idx]; + l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; + sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + o_acc(i_j_idx) *= tmp; + }); + } + else + { + // Skip: P already in m_{j-1} frame, no o_acc rescale needed. + l(i_idx) = l[i_idx] + rowsum_p[i_idx]; + } }); if constexpr(kHasDropout) diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp index 9b28170916..5155794b16 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp @@ -664,6 +664,80 @@ struct BlockFmhaPipelineQRKSVSAsync } }; + // Conditional rescaling (FA4): skip when correction is negligible. + // For skip rows we stabilize P with m_old so P is computed directly in + // the m_{j-1} frame, eliminating the post-correction sweep. + // FP8 quant modes cast P to FP8 after softmax. In the skip + // branch P can exceed the FP8 representable range and saturate, + // corrupting the P*V GEMM. Disable skip for all FP8 paths. + static constexpr SMPLComputeDataType kRescaleThreshold = + type_convert( + QScaleEnum == BlockAttentionQuantScaleEnum::NO_SCALE ? 8.0f : 0.0f); + + auto m_stab = + make_static_distributed_tensor(m.get_tile_distribution()); + auto rescale_factor = + make_static_distributed_tensor(m.get_tile_distribution()); + auto needs_rescale = make_static_distributed_tensor(m.get_tile_distribution()); + set_tile(needs_rescale, false); + + constexpr auto m_spans = decltype(m)::get_distributed_spans(); + sweep_tile_span(m_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); +#if CK_TILE_FMHA_FWD_FAST_EXP2 + const auto acc_scale_log2 = [&]() { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + if constexpr(kHasLogitsSoftCap) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + auto row_max = scale_s * get_validated_m(m[i_idx]); + return scale_s * m_old[i_idx] - row_max; + } + } + }(); + + const bool need_rescale = + (acc_scale_log2 < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp2(acc_scale_log2); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } +#else + const auto diff = m_old[i_idx] - get_validated_m(m[i_idx]); + const bool need_rescale = + (diff < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp(diff); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } +#endif + }); + constexpr auto p_spans = decltype(p_compute)::get_distributed_spans(); sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); @@ -671,7 +745,7 @@ struct BlockFmhaPipelineQRKSVSAsync // For BLOCKSCALE: precompute (m - shift) once per row // Bias/Alibi/SoftCap: exp2(s - m + shift) = exp2(s - (m - shift)) // else: exp2(scale_s*s - scale_s*m + shift) = exp2(scale_s*s - (scale_s*m - shift)) - auto validated_m = get_validated_m(m[i_idx]); + auto validated_m = get_validated_m(m_stab[i_idx]); auto row_max = scale_s * validated_m; if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE) { @@ -704,7 +778,7 @@ struct BlockFmhaPipelineQRKSVSAsync } } #else - p_compute(i_j_idx) = exp(s[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = exp(s[i_j_idx] - get_validated_m(m_stab[i_idx])); #endif }); }); @@ -717,37 +791,20 @@ struct BlockFmhaPipelineQRKSVSAsync constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); sweep_tile_span(o_spans[number<0>{}], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); -#if CK_TILE_FMHA_FWD_FAST_EXP2 - const auto tmp = [&]() { - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || - BiasEnum == BlockAttentionBiasEnum::ALIBI) - { - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - if constexpr(kHasLogitsSoftCap) - { - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - auto row_max = scale_s * get_validated_m(m[i_idx]); - return exp2(scale_s * m_old[i_idx] - row_max); - } - } - }(); -#else - const auto tmp = exp(m_old[i_idx] - get_validated_m(m[i_idx])); -#endif - l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; - sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - // FIXME: this use different equation from FA v2 paper, - // but produce correc result. - // Is the equation wrong? - o_acc(i_j_idx) *= tmp; - }); + if(needs_rescale[i_idx]) + { + const auto tmp = rescale_factor[i_idx]; + l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; + sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + o_acc(i_j_idx) *= tmp; + }); + } + else + { + // Skip: P already in m_{j-1} frame, no o_acc rescale needed. + l(i_idx) = l[i_idx] + rowsum_p[i_idx]; + } }); if constexpr(kHasDropout) diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp index bccb32c546..e03d1f3439 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp @@ -560,22 +560,75 @@ struct BlockFmhaPipelineQRKSVSAsyncTrload } }; + // Conditional rescaling (FA4): skip when correction is negligible. + // For skip rows we stabilize P with m_old so P is computed directly in + // the m_{j-1} frame, eliminating the post-correction sweep. + static constexpr SMPLComputeDataType kRescaleThreshold = + type_convert(8.0f); + + auto m_stab = + make_static_distributed_tensor(m.get_tile_distribution()); + auto rescale_factor = + make_static_distributed_tensor(m.get_tile_distribution()); + auto needs_rescale = make_static_distributed_tensor(m.get_tile_distribution()); + set_tile(needs_rescale, false); + + constexpr auto m_spans = decltype(m)::get_distributed_spans(); + sweep_tile_span(m_spans[I0], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + const auto acc_scale_log2 = [&]() { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + if constexpr(kHasLogitsSoftCap) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + auto row_max = scale_s * get_validated_m(m[i_idx]); + return scale_s * m_old[i_idx] - row_max; + } + } + }(); + + const bool need_rescale = + (acc_scale_log2 < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp2(acc_scale_log2); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } + }); + constexpr auto p_spans = decltype(p_compute)::get_distributed_spans(); sweep_tile_span(p_spans[I0], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); - auto row_max = scale_s * get_validated_m(m[i_idx]); + auto row_max = scale_s * get_validated_m(m_stab[i_idx]); sweep_tile_span(p_spans[I1], [&](auto idx1) { constexpr auto i_j_idx = make_tuple(idx0, idx1); if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || BiasEnum == BlockAttentionBiasEnum::ALIBI) { - p_compute(i_j_idx) = exp2(s_new[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = exp2(s_new[i_j_idx] - get_validated_m(m_stab[i_idx])); } else { if constexpr(kHasLogitsSoftCap) { - p_compute(i_j_idx) = exp2(s_new[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = + exp2(s_new[i_j_idx] - get_validated_m(m_stab[i_idx])); } else { @@ -591,41 +644,30 @@ struct BlockFmhaPipelineQRKSVSAsyncTrload block_tile_reduce_sync( rowsum_p, f_sum, bool_constant{} /*, bool_constant{}*/); - auto p_tile = make_static_distributed_tensor( - Policy::template MakePRegTileDistribution()); - p_tile.get_thread_buffer() = cast_tile(p_compute).get_thread_buffer(); - // l{j}, Oacc{j} constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); sweep_tile_span(o_spans[I0], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); - const auto tmp = [&]() { - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || - BiasEnum == BlockAttentionBiasEnum::ALIBI) - { - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - if constexpr(kHasLogitsSoftCap) - { - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - auto row_max = scale_s * get_validated_m(m[i_idx]); - return exp2(scale_s * m_old[i_idx] - row_max); - } - } - }(); - l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; - sweep_tile_span(o_spans[I1], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - - o_acc(i_j_idx) *= tmp; - }); + if(needs_rescale[i_idx]) + { + const auto tmp = rescale_factor[i_idx]; + l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; + sweep_tile_span(o_spans[I1], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + o_acc(i_j_idx) *= tmp; + }); + } + else + { + // Skip: P already in m_{j-1} frame, no o_acc rescale needed. + l(i_idx) = l[i_idx] + rowsum_p[i_idx]; + } }); + auto p_tile = make_static_distributed_tensor( + Policy::template MakePRegTileDistribution()); + p_tile.get_thread_buffer() = cast_tile(p_compute).get_thread_buffer(); + block_sync_lds_direct_load(); auto v_tile = load_tile_transpose(v_lds_read_window); @@ -1094,22 +1136,75 @@ struct BlockFmhaPipelineQRKSVSAsyncTrload } }; + // Conditional rescaling (FA4): skip when correction is negligible. + // For skip rows we stabilize P with m_old so P is computed directly in + // the m_{j-1} frame, eliminating the post-correction sweep. + static constexpr SMPLComputeDataType kRescaleThreshold = + type_convert(8.0f); + + auto m_stab = + make_static_distributed_tensor(m.get_tile_distribution()); + auto rescale_factor = + make_static_distributed_tensor(m.get_tile_distribution()); + auto needs_rescale = make_static_distributed_tensor(m.get_tile_distribution()); + set_tile(needs_rescale, false); + + constexpr auto m_spans = decltype(m)::get_distributed_spans(); + sweep_tile_span(m_spans[I0], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + const auto acc_scale_log2 = [&]() { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + if constexpr(kHasLogitsSoftCap) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + auto row_max = scale_s * get_validated_m(m[i_idx]); + return scale_s * m_old[i_idx] - row_max; + } + } + }(); + + const bool need_rescale = + (acc_scale_log2 < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp2(acc_scale_log2); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } + }); + constexpr auto p_spans = decltype(p_compute)::get_distributed_spans(); sweep_tile_span(p_spans[I0], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); - auto row_max = scale_s * get_validated_m(m[i_idx]); + auto row_max = scale_s * get_validated_m(m_stab[i_idx]); sweep_tile_span(p_spans[I1], [&](auto idx1) { constexpr auto i_j_idx = make_tuple(idx0, idx1); if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || BiasEnum == BlockAttentionBiasEnum::ALIBI) { - p_compute(i_j_idx) = exp2(s_new[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = exp2(s_new[i_j_idx] - get_validated_m(m_stab[i_idx])); } else { if constexpr(kHasLogitsSoftCap) { - p_compute(i_j_idx) = exp2(s_new[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = + exp2(s_new[i_j_idx] - get_validated_m(m_stab[i_idx])); } else { @@ -1125,41 +1220,30 @@ struct BlockFmhaPipelineQRKSVSAsyncTrload block_tile_reduce_sync( rowsum_p, f_sum, bool_constant{} /*, bool_constant{}*/); - auto p_tile = make_static_distributed_tensor( - Policy::template MakePRegTileDistribution()); - p_tile.get_thread_buffer() = cast_tile(p_compute).get_thread_buffer(); - // l{j}, Oacc{j} constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); sweep_tile_span(o_spans[I0], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); - const auto tmp = [&]() { - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || - BiasEnum == BlockAttentionBiasEnum::ALIBI) - { - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - if constexpr(kHasLogitsSoftCap) - { - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - auto row_max = scale_s * get_validated_m(m[i_idx]); - return exp2(scale_s * m_old[i_idx] - row_max); - } - } - }(); - l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; - sweep_tile_span(o_spans[I1], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - - o_acc(i_j_idx) *= tmp; - }); + if(needs_rescale[i_idx]) + { + const auto tmp = rescale_factor[i_idx]; + l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; + sweep_tile_span(o_spans[I1], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + o_acc(i_j_idx) *= tmp; + }); + } + else + { + // Skip: P already in m_{j-1} frame, no o_acc rescale needed. + l(i_idx) = l[i_idx] + rowsum_p[i_idx]; + } }); + auto p_tile = make_static_distributed_tensor( + Policy::template MakePRegTileDistribution()); + p_tile.get_thread_buffer() = cast_tile(p_compute).get_thread_buffer(); + block_sync_lds(); move_tile_window(k_dram_window, {kN0, 0}); k_lds_write_window.set_bottom_tensor_view_data_ptr(k_lds_write_ptr); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_fp8.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_fp8.hpp index 85526c9e24..4cb62a8ee1 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_fp8.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_fp8.hpp @@ -379,25 +379,87 @@ struct [[deprecated]] BlockFmhaPipelineQRKSVSFp8 } }; + // Conditional rescaling (FA4): skip when correction is negligible. + // For skip rows we stabilize P with m_old so P is computed directly in + // the m_{j-1} frame, eliminating the post-correction sweep. + static constexpr SMPLComputeDataType kRescaleThreshold = + type_convert(8.0f); + + auto m_stab = + make_static_distributed_tensor(m.get_tile_distribution()); + auto rescale_factor = + make_static_distributed_tensor(m.get_tile_distribution()); + auto needs_rescale = make_static_distributed_tensor(m.get_tile_distribution()); + set_tile(needs_rescale, false); + + constexpr auto m_spans = decltype(m)::get_distributed_spans(); + sweep_tile_span(m_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); +#if CK_TILE_FMHA_FWD_FAST_EXP2 + const auto acc_scale_log2 = [&]() { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + auto row_max = scale_s * get_validated_m(m[i_idx]); + return scale_s * m_old[i_idx] - row_max; + } + }(); + + const bool need_rescale = + (acc_scale_log2 < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp2(acc_scale_log2); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } +#else + const auto diff = m_old[i_idx] - get_validated_m(m[i_idx]); + const bool need_rescale = + (diff < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp(diff); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } +#endif + }); + constexpr auto p_spans = decltype(p_compute)::get_distributed_spans(); sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); #if CK_TILE_FMHA_FWD_FAST_EXP2 - auto row_max = scale_s * get_validated_m(m[i_idx]); + auto row_max = scale_s * get_validated_m(m_stab[i_idx]); #endif sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) { constexpr auto i_j_idx = make_tuple(idx0, idx1); #if CK_TILE_FMHA_FWD_FAST_EXP2 if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { - p_compute(i_j_idx) = exp2(s[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = exp2(s[i_j_idx] - get_validated_m(m_stab[i_idx])); } else { p_compute(i_j_idx) = exp2(scale_s * s[i_j_idx] - row_max); } #else - p_compute(i_j_idx) = exp(s[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = exp(s[i_j_idx] - get_validated_m(m_stab[i_idx])); #endif }); }); @@ -410,29 +472,20 @@ struct [[deprecated]] BlockFmhaPipelineQRKSVSFp8 constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); sweep_tile_span(o_spans[number<0>{}], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); -#if CK_TILE_FMHA_FWD_FAST_EXP2 - const auto tmp = [&]() { - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) - { - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - auto row_max = scale_s * get_validated_m(m[i_idx]); - return exp2(scale_s * m_old[i_idx] - row_max); - } - }(); -#else - const auto tmp = exp(m_old[i_idx] - get_validated_m(m[i_idx])); -#endif - l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; - sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - // FIXME: this use different equation from FA v2 paper, - // but produce correc result. - // Is the equation wrong? - o_acc(i_j_idx) *= tmp; - }); + if(needs_rescale[i_idx]) + { + const auto tmp = rescale_factor[i_idx]; + l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; + sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + o_acc(i_j_idx) *= tmp; + }); + } + else + { + // Skip: P already in m_{j-1} frame, no o_acc rescale needed. + l(i_idx) = l[i_idx] + rowsum_p[i_idx]; + } }); block_sync_lds(); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qs_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qs_ks_vs.hpp index 4eb5eb291a..e537a8b98a 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qs_ks_vs.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qs_ks_vs.hpp @@ -476,11 +476,81 @@ struct BlockFmhaPipelineQSKSVS } }; + // Conditional rescaling (FA4): skip when correction is negligible. + // For skip rows we stabilize P with m_old so P is computed directly in + // the m_{j-1} frame, eliminating the post-correction sweep. + static constexpr SMPLComputeDataType kRescaleThreshold = + type_convert(8.0f); + + auto m_stab = + make_static_distributed_tensor(m.get_tile_distribution()); + auto rescale_factor = + make_static_distributed_tensor(m.get_tile_distribution()); + auto needs_rescale = make_static_distributed_tensor(m.get_tile_distribution()); + set_tile(needs_rescale, false); + + constexpr auto m_spans = decltype(m)::get_distributed_spans(); + sweep_tile_span(m_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); +#if CK_TILE_FMHA_FWD_FAST_EXP2 + const auto acc_scale_log2 = [&]() { + if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || + BiasEnum == BlockAttentionBiasEnum::ALIBI) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + if constexpr(kHasLogitsSoftCap) + { + return m_old[i_idx] - get_validated_m(m[i_idx]); + } + else + { + auto row_max = scale_s * get_validated_m(m[i_idx]); + return scale_s * m_old[i_idx] - row_max; + } + } + }(); + + const bool need_rescale = + (acc_scale_log2 < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp2(acc_scale_log2); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } +#else + const auto diff = m_old[i_idx] - get_validated_m(m[i_idx]); + const bool need_rescale = + (diff < type_convert(-kRescaleThreshold)); + + if(need_rescale) + { + rescale_factor(i_idx) = exp(diff); + m_stab(i_idx) = m[i_idx]; + needs_rescale(i_idx) = true; + } + else + { + m_stab(i_idx) = m_old[i_idx]; + m(i_idx) = m_old[i_idx]; + } +#endif + }); + constexpr auto p_spans = decltype(p_compute)::get_distributed_spans(); sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); #if CK_TILE_FMHA_FWD_FAST_EXP2 - auto row_max = scale_s * get_validated_m(m[i_idx]); + auto row_max = scale_s * get_validated_m(m_stab[i_idx]); #endif sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) { constexpr auto i_j_idx = make_tuple(idx0, idx1); @@ -488,13 +558,13 @@ struct BlockFmhaPipelineQSKSVS if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || BiasEnum == BlockAttentionBiasEnum::ALIBI) { - p_compute(i_j_idx) = exp2(s[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = exp2(s[i_j_idx] - get_validated_m(m_stab[i_idx])); } else { if constexpr(kHasLogitsSoftCap) { - p_compute(i_j_idx) = exp2(s[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = exp2(s[i_j_idx] - get_validated_m(m_stab[i_idx])); } else { @@ -502,7 +572,7 @@ struct BlockFmhaPipelineQSKSVS } } #else - p_compute(i_j_idx) = exp(s[i_j_idx] - get_validated_m(m[i_idx])); + p_compute(i_j_idx) = exp(s[i_j_idx] - get_validated_m(m_stab[i_idx])); #endif }); }); @@ -512,6 +582,28 @@ struct BlockFmhaPipelineQSKSVS block_tile_reduce_sync(rowsum_p, f_sum, bool_constant{}); + __builtin_amdgcn_sched_barrier(0); + + // l{j}, Oacc{j} + constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); + sweep_tile_span(o_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + if(needs_rescale[i_idx]) + { + const auto tmp = rescale_factor[i_idx]; + l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; + sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + o_acc(i_j_idx) *= tmp; + }); + } + else + { + // Skip: P already in m_{j-1} frame, no o_acc rescale needed. + l(i_idx) = l[i_idx] + rowsum_p[i_idx]; + } + }); + #if defined(__gfx11__) // gfx11 WMMA uses different lane layouts for GEMM C and GEMM A tiles, so remap // softmax P from GEMM0's C layout into GEMM1's A layout before the PV GEMM. @@ -524,46 +616,6 @@ struct BlockFmhaPipelineQSKSVS cast_tile(tile_elementwise_in(p_compute_element_func, p_compute)); #endif - __builtin_amdgcn_sched_barrier(0); - - // l{j}, Oacc{j} - constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); - sweep_tile_span(o_spans[number<0>{}], [&](auto idx0) { - constexpr auto i_idx = make_tuple(idx0); -#if CK_TILE_FMHA_FWD_FAST_EXP2 - const auto tmp = [&]() { - if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS || - BiasEnum == BlockAttentionBiasEnum::ALIBI) - { - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - if constexpr(kHasLogitsSoftCap) - { - - return exp2(m_old[i_idx] - get_validated_m(m[i_idx])); - } - else - { - auto row_max = scale_s * get_validated_m(m[i_idx]); - return exp2(scale_s * m_old[i_idx] - row_max); - } - } - }(); -#else - const auto tmp = exp(m_old[i_idx] - get_validated_m(m[i_idx])); -#endif - l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx]; - sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - // FIXME: this use different equation from FA v2 paper, - // but produce correc result. - // Is the equation wrong? - o_acc(i_j_idx) *= tmp; - }); - }); - block_sync_lds(); if constexpr(std::is_same_v) { From 843d993835c529439df9d94336f606b951274244 Mon Sep 17 00:00:00 2001 From: Brock Hargreaves <253123018+brockhargreaves-amd@users.noreply.github.com> Date: Mon, 1 Jun 2026 21:10:12 +0000 Subject: [PATCH 009/143] [rocm-libraries] ROCm/rocm-libraries#7743 (commit 15ef85c) [CK] Extract Jenkinsfile helpers into vars/ck.groovy shared library (#7743) ## Motivation The CK Jenkinsfile is a 2,215-line monolith mixing helper function definitions with pipeline stage declarations. This makes it difficult to review, modify, or extend CI stages without wading through unrelated infrastructure code. ## Technical Details Extract all helper functions from the Jenkinsfile into vars/ck.groovy, loaded at runtime via ck = load "vars/ck.groovy" in the first stage. The Jenkinsfile is reduced from 2,215 lines to 810 lines containing only the pipeline structure. - 36 helper functions moved to ck.groovy with no logic changes - 10 new stage-wrapper functions (runBuildCKAndTests, runTileEngineGemmTests, runClangFormat, etc.) extract inline environment{}/steps{} business logic from stages, eliminating the MethodTooLargeException caused by CPS-transformed shell strings exceeding the JVM 64KB bytecode limit - All ck. method calls in steps{} blocks wrapped in script{} as required by Jenkins Declarative Pipeline - rocmnode() remains in the Jenkinsfile (needed for agent{} labels before ck is loaded) - CRON_SETTINGS / POLL_SPEC remain in the Jenkinsfile (triggers{} evaluates at parse time before any workspace is available) - No stage names changed ## Test Plan - Jenkinsfile validated against the Jenkins Pipeline Linter (/pipeline-model-converter/validate) - All 35 shared helper functions diffed line-by-line against develop to verify no regressions - Merge from develop incorporated and verified (gfx1250 stage, ROCm 7.13 default, cmake_build updates) ## Test Result - Linter: passes - Function diff vs develop: all 35 functions match exactly - Awaiting Jenkins run to confirm end-to-end stage execution ## Submission Checklist - [ x ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- Jenkinsfile | 1519 +++--------------------------------------------- vars/ck.groovy | 1414 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 1506 insertions(+), 1427 deletions(-) create mode 100644 vars/ck.groovy diff --git a/Jenkinsfile b/Jenkinsfile index 7cd7a2546a..dfa904fcf8 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -21,1233 +21,24 @@ // - Forces full build if dependency cache stale (>7 days) // - Manual override: set DISABLE_SMART_BUILD=true // -// Benefits: PR builds 5h → 30min (typical), nightly builds unchanged +// Benefits: PR builds 5h -> 30min (typical), nightly builds unchanged // See: script/dependency-parser/README.md for details // - -@NonCPS -String getGitHubCommitHash(def build) -{ - def scmAction = build?.actions.find { action -> - action instanceof jenkins.scm.api.SCMRevisionAction - } - if (scmAction?.revision instanceof org.jenkinsci.plugins.github_branch_source.PullRequestSCMRevision) - { - return scmAction.revision.pullHash - } - else if (scmAction?.revision instanceof jenkins.plugins.git.AbstractGitSCMSource$SCMRevisionImpl) - { - return scmAction.revision.hash - } - return null -} +ck = null def rocmnode(name) { return '(rocmtest || miopen) && (' + name + ')' } -def show_node_info() { - sh """ - echo "NODE_NAME = \$NODE_NAME" - hostname - lsb_release -sd - uname -r - cat /sys/module/amdgpu/version - ls /opt/ -la - """ -} - -def setGithubStatus(String context, String state, String description) { - def sha = env.GIT_COMMIT - def targetUrl = env.RUN_DISPLAY_URL ?: env.BUILD_URL - def statusUrl = "https://api.github.com/repos/ROCm/rocm-libraries/statuses/${sha}" - withCredentials([usernamePassword(credentialsId: 'github-app-miopen', usernameVariable: 'GITHUB_APP', passwordVariable: 'GITHUB_TOKEN')]) { - def code = '0' - try { - retry(3) { - code = sh(returnStdout: true, script: """ - curl -s -w "%{http_code}" -o /dev/null -X POST '${statusUrl}' \\ - -H "Authorization: token \$GITHUB_TOKEN" \\ - -H 'Content-Type: application/json' \\ - -d '{"state":"${state}","context":"${context}","description":"${description}","target_url":"${targetUrl}"}' - """).trim() - if (!code.startsWith('2')) { - error("GitHub status POST returned ${code}") - } - } - } catch (Exception e) { - echo "WARNING: GitHub status POST failed after retries (context=${context}, state=${state}, code=${code})" - } - } -} - -def cloneUpdateRefRepo() { - def refRepoPath = "/var/jenkins/ref-repo/rocm-libraries" - def lockLabel = "git ref repo lock - ${env.NODE_NAME}" - def folderExists = sh( - script: "test -d ${refRepoPath}/refs", - returnStatus: true - ) == 0 - - if (!folderExists) { - echo "rocm-libraries repo does not exist at ${refRepoPath}, creating mirror clone..." - echo "locking on label: ${lockLabel}" - lock(lockLabel) { - def cloneCommand = """ - set -ex - rm -rf ${refRepoPath} && mkdir -p ${refRepoPath} - git clone --mirror https://github.com/ROCm/rocm-libraries.git ${refRepoPath} - """ - sh(script: cloneCommand, label: "clone ref repo") - } - echo "Completed git clone, lock released" - } - echo "rocm-libraries repo exists at ${refRepoPath}, performing git remote update..." - echo "locking on label: ${lockLabel}" - lock(lockLabel) { - def fetchCommand = """ - set -ex - cd ${refRepoPath} - git remote prune origin - git remote update - """ - sh(script: fetchCommand, label: "update ref repo") - } - echo "Completed git ref repo fetch, lock released" -} - -def checkoutComposableKernel() -{ - //update ref repo - cloneUpdateRefRepo() - // checkout project - def scmVars = checkout scm - // getGitHubCommitHash reads SCMRevisionAction recorded before any local merge, - // giving the true PR branch tip (pullHash) or branch HEAD (hash). - // Falls back to ORIG_HEAD (pre-merge HEAD set by git merge) when SCMRevisionAction - // is unavailable, then to HEAD for branch builds where no merge occurred. - env.GIT_COMMIT = getGitHubCommitHash(currentBuild.rawBuild) ?: sh(returnStdout: true, script: ''' - git rev-parse ORIG_HEAD 2>/dev/null || git rev-parse HEAD - ''').trim() -} - -def generateAndArchiveBuildTraceVisualization(String buildTraceFileName) { - try { - checkoutComposableKernel() - - // Retrieve the build trace artifact - def traceFileExists = false - try { - copyArtifacts( - projectName: env.JOB_NAME, - selector: specific(env.BUILD_NUMBER), - filter: buildTraceFileName - ) - traceFileExists = fileExists(buildTraceFileName) - } catch (Exception e) { - echo "Could not copy build trace artifact: ${e.getMessage()}" - traceFileExists = false - return - } - - sh """ - echo "post artifact download:" - ls -la - """ - - // Pull image - def image = "ghcr.io/puppeteer/puppeteer:24.30.0" - echo "Pulling image: ${image}" - def retimage = docker.image("${image}") - retimage.pull() - - // Create a temporary workspace - sh """#!/bin/bash - ls -la - mkdir -p workspace - cp ./projects/composablekernel/script/infra_helper/capture_build_trace.js ./workspace - cp ${buildTraceFileName} ./workspace/${buildTraceFileName} - chmod 777 ./workspace - ls -la ./workspace - """ - - // Run container to get snapshot - def dockerOpts = "--cap-add=SYS_ADMIN -v \"\$(pwd)/workspace:/workspace\" -e NODE_PATH=/home/pptruser/node_modules -e BUILD_TRACE_FILE=${buildTraceFileName}" - // Create unique image name by sanitizing job name - def sanitizedJobName = env.JOB_NAME.replaceAll(/[\/\\:*?"<>| ]/, '_').replaceAll('%2F', '_') - def architectureName = (buildTraceFileName =~ /(gfx[0-9a-zA-Z]+)/)[0][1] - def imageName = "perfetto_snapshot_${sanitizedJobName}_build_${env.BUILD_NUMBER}_${architectureName}.png" - sh """ - docker run --rm ${dockerOpts} ${image} node /workspace/capture_build_trace.js - mv ./workspace/perfetto_snapshot_build.png ./workspace/${imageName} - """ - - // Archive the snapshot - sh """ - mv ./workspace/${imageName} ${imageName} - """ - archiveArtifacts "${imageName}" - - // Notify the channel - withCredentials([string(credentialsId: 'ck_ci_build_perf_webhook_url', variable: 'WEBHOOK_URL')]) { - sh ''' - # Create build trace filename with build number based on the original filename - BUILD_TRACE_WITH_NUMBER=$(echo "''' + buildTraceFileName + '''" | sed 's/.json/_''' + sanitizedJobName + '''_''' + env.BUILD_NUMBER + '''_''' + architectureName + '''.json/') - - # Convert image to base64 - echo "Converting image to base64..." - IMAGE_BASE64=$(base64 -w 0 ''' + imageName + ''') - echo "Image base64 length: ${#IMAGE_BASE64}" - - # Convert build trace to base64 - echo "Converting build trace to base64..." - BUILD_TRACE_BASE64=$(base64 -w 0 ''' + buildTraceFileName + ''') - echo "Build trace base64 length: ${#BUILD_TRACE_BASE64}" - - # Create JSON payload with base64 data - echo "Creating JSON payload..." - { - printf '{\n' - printf ' "jobName": "%s",\n' "''' + env.JOB_NAME + '''" - printf ' "buildNumber": "%s",\n' "''' + env.BUILD_NUMBER + '''" - printf ' "jobUrl": "%s",\n' "''' + env.RUN_DISPLAY_URL + '''" - printf ' "imageName": "%s",\n' "''' + imageName + '''" - printf ' "architecture": "%s",\n' "''' + architectureName + '''" - printf ' "imageData": "%s",\n' "$IMAGE_BASE64" - printf ' "buildTraceName": "%s",\n' "$BUILD_TRACE_WITH_NUMBER" - printf ' "buildTraceData": "%s"\n' "$BUILD_TRACE_BASE64" - printf '}\n' - } > webhook_payload.json - - echo "JSON payload created, size: $(wc -c < webhook_payload.json) bytes" - - curl -X POST "${WEBHOOK_URL}" \ - -H "Content-Type: application/json" \ - -d @webhook_payload.json - - # Clean up temporary file - rm -f webhook_payload.json - ''' - } - } catch (Exception e) { - echo "Throwing error exception while generating build trace visualization" - echo 'Exception occurred: ' + e.toString() - } -} - -class Version { - int major, minor, patch - @Override - String toString() { - return [major, minor, patch].findAll().join('.') - } -} -def parseVersion(String versionString) { - if (!versionString) return null - int[] tokens = versionString.split(/\./).collect { it as int } // Splits the string by '.' and converts each part to an integer. - return new Version( - major: tokens[0], - minor: tokens.length > 1 ? tokens[1] : null, - patch: tokens.length > 2 ? tokens[2] : null, - ) -} - -def nthreads() { - def nproc = sh(returnStdout: true, script: 'nproc') - echo "Number of cores: ${nproc}" - def n = nproc.toInteger() - if (n > 64){ - n = 64 - } - echo "Number of threads used for building: ${n}" - return n -} - -def runShell(String command){ - def responseCode = sh returnStatus: true, script: "${command} > tmp.txt" - def output = readFile(file: "tmp.txt") - return (output != "") -} - -def shouldRunCICheck() { - // File patterns that should not trigger CI - def skipFilePatterns = [ - /^projects\/composablekernel\/\.github\/.*/, // GitHub workflow files - /^projects\/composablekernel\/docs\/.*/, // Documentation files - /^projects\/composablekernel\/LICENSE$/, // License file - /^projects\/composablekernel\/.*\.gitignore$/, // Git ignore files - /^projects\/composablekernel\/.*\.md$/ // Markdown files - ] - - try { - // Always run if this is a base branch build - def baseBranch = "develop" - def isBaseBranchBuild = (env.CHANGE_ID == null && env.BRANCH_NAME == baseBranch) - - if (isBaseBranchBuild) { - echo "Base branch (${baseBranch}) build detected - always running CI for safety" - return true - } - - // Get the list of changed files (all files touched in any commit, even if reverted) - def changedFiles = sh( - returnStdout: true, - script: ''' - BASE_BRANCH="develop" - - if [ "$CHANGE_ID" != "" ]; then - # For PR builds, get all files touched in any commit - echo "PR build detected, checking all touched files against origin/$CHANGE_TARGET" >&2 - git log --name-only --pretty=format: origin/$CHANGE_TARGET..HEAD -- projects/composablekernel/ | sort -u | grep -v '^$' || true - else - # For feature branch builds, compare against merge-base with base branch - MERGE_BASE=$(git merge-base HEAD origin/$BASE_BRANCH 2>/dev/null || echo "HEAD~1") - echo "Branch build detected, checking all touched files since merge-base: $MERGE_BASE" >&2 - git log --name-only --pretty=format: $MERGE_BASE..HEAD -- projects/composablekernel/ | sort -u | grep -v '^$' || true - fi - ''' - ).trim().split('\n') - - if (changedFiles.size() == 1 && changedFiles[0] == '') { - echo "No changed files detected - this might be a manual trigger or merge commit, running CI for safety" - return true - } - - echo "Changed files: ${changedFiles.join(', ')}" - - // Separate files into those requiring CI and those that can be skipped - def filesRequiringCI = [] - def skippedFiles = [] - - changedFiles.each { file -> - def shouldSkip = skipFilePatterns.any { pattern -> - file ==~ pattern - } - - if (shouldSkip) { - skippedFiles.add(file) - } else { - filesRequiringCI.add(file) - } - } - - // Debug output - if (skippedFiles.size() > 0) { - echo "Files that don't require CI (${skippedFiles.size()}):" - skippedFiles.each { echo " - ${it}" } - } - - if (filesRequiringCI.size() > 0) { - echo "Files that require CI (${filesRequiringCI.size()}):" - filesRequiringCI.each { echo " - ${it}" } - return true - } else { - echo "Only non-relevant files changed, skipping CI" - return false - } - } catch (Exception e) { - echo "Error checking changed files: ${e.getMessage()}, running CI by default" - return true - } -} - -def getBaseDockerImageName(){ - def img - if (params.USE_CUSTOM_DOCKER != ""){ - img = "${params.USE_CUSTOM_DOCKER}" - } - else{ - img = "${env.CK_DOCKERHUB}:ck_ub24.04_rocm${params.ROCMVERSION}" - } - return img -} - -def getDockerImageName(){ - def img - def base_name = getBaseDockerImageName() - if (params.USE_CUSTOM_DOCKER != ""){ - img = "${params.USE_CUSTOM_DOCKER}" - } - else{ - if (params.COMPILER_VERSION == "") { - img = "${base_name}" - } - else{ - if (params.COMPILER_COMMIT == ""){ - img = "${base_name}_${params.COMPILER_VERSION}" - } - else{ - def commit = "${params.COMPILER_COMMIT}"[0..6] - img = "${base_name}_${params.COMPILER_VERSION}_${commit}" - } - } - } - return img -} - -def check_host() { - if ("${env.CK_SCCACHE}" != "null"){ - def SCCACHE_SERVER="${env.CK_SCCACHE.split(':')[0]}" - echo "sccache server: ${SCCACHE_SERVER}" - sh "chmod +w -R ${env.WORKSPACE}" - sh '''ping -c 1 -p 6379 "${SCCACHE_SERVER}" | echo $? > tmp.txt''' - def output = readFile(file: "tmp.txt") - echo "tmp.txt contents: \$output" - return (output != "0") - } - else{ - return 1 - } -} - -def check_arch_name(){ - sh 'rocminfo | tee rocminfo.log' - if ( runShell('grep -n "gfx90a" rocminfo.log') ){ - return "gfx90a" - } - else if ( runShell('grep -n "gfx942" rocminfo.log') ) { - return "gfx942" - } - else if ( runShell('grep -n "gfx101" rocminfo.log') ) { - return "gfx101" - } - else if ( runShell('grep -n "gfx103" rocminfo.log') ) { - return "gfx103" - } - else if ( runShell('grep -n "gfx11" rocminfo.log') ) { - return "gfx11" - } - else if ( runShell('grep -n "gfx120" rocminfo.log') ) { - return "gfx12" - } - else if ( runShell('grep -n "gfx908" rocminfo.log') ) { - return "gfx908" - } - else if ( runShell('grep -n "gfx950" rocminfo.log') ) { - return "gfx950" - } - else { - return "" - } -} - -def getDockerImage(Map conf=[:]){ - def image - if ( conf.get("docker_name", "") != "" ){ - image = conf.get("docker_name", "") - echo "Using special docker: ${image}" - } - else{ - image = getDockerImageName() - echo "Using default docker: ${image}" - } - //Check if image exists - def retimage - try - { - echo "Pulling image: ${image}" - retimage = docker.image("${image}") - withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { - retimage.pull() - } - } - catch(Exception ex) - { - error "Unable to locate image: ${image}" - } - return [retimage, image] -} - -// Build and push a docker image, capturing its digest into the specified env var. -// If forceBuild is false, will skip building if the image already exists in the registry. -def buildAndPushDockerImage(String install_prefix, String image_name, String dockerExtraArgs, boolean forceBuild){ - show_node_info() - env.DOCKER_BUILDKIT=1 - checkoutComposableKernel() - def dockerArgs = "--build-arg PREFIX=${install_prefix} --build-arg compiler_version='${params.COMPILER_VERSION}' --build-arg compiler_commit='${params.COMPILER_COMMIT}' --build-arg ROCMVERSION='${params.ROCMVERSION}' " - dockerArgs += " " + dockerExtraArgs - - if(!forceBuild){ - try{ - echo "Checking for image: ${image_name}" - withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { - sh "docker manifest inspect --insecure ${image_name}" - } - echo "Image: ${image_name} found! Skipping building image" - return image_name - } - catch(Exception ex){ - echo "Unable to locate image: ${image_name}. Will attempt to build image now." - } - } - - echo "Building image: ${image_name} with args: ${dockerArgs}" - def retimage = docker.build("${image_name}", dockerArgs) - withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { - retimage.push() - } - def digest = sh(returnStdout: true, script: "docker inspect --format='{{index .RepoDigests 0}}' ${image_name}").trim() - echo "Built image digest: ${digest}" - echo "Pruning dangling Docker images to free disk space on CI agent" - sh "docker image prune -f --filter 'dangling=true' || true" - return digest -} - -def buildDockerBase(install_prefix){ - def image_name = getDockerImageName() - def base_image_name = getBaseDockerImageName() - echo "Building Docker for ${image_name}" - def dockerExtraArgs = " -f projects/composablekernel/Dockerfile . " - if(params.COMPILER_VERSION == "develop" || params.COMPILER_VERSION == "amd-staging" || params.COMPILER_COMMIT != ""){ - dockerExtraArgs = " --no-cache --build-arg BASE_DOCKER='${base_image_name}' -f projects/composablekernel/Dockerfile.compiler . " - } - else if(params.COMPILER_VERSION == "therock"){ - dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile . " - } - env.CK_BASE_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, params.BUILD_DOCKER.toBoolean()) -} - -def buildDockerPytorch(install_prefix){ - def image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_pytorch" - def dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile.pytorch --build-arg CK_PYTORCH_BRANCH='${params.ck_pytorch_branch}' . " - env.CK_PYTORCH_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, true) -} - -def buildDockerAiter(install_prefix){ - def image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_aiter" - def dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile.aiter --build-arg AITER_BRANCH='${params.aiter_branch}' --build-arg CK_AITER_BRANCH='${params.ck_aiter_branch}' . " - env.CK_AITER_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, true) -} - -def buildDockerFa(install_prefix){ - def image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_fa" - def dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile.fa" - dockerExtraArgs += " --build-arg BASE_DOCKER='${params.fa_base_docker}'" - dockerExtraArgs += " --build-arg FA_BRANCH='${params.fa_branch}'" - dockerExtraArgs += " --build-arg CK_FA_BRANCH='${params.ck_fa_branch}'" - dockerExtraArgs += " --build-arg GPU_ARCHS='gfx942;gfx950'" - dockerExtraArgs += " . " - env.CK_FA_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, true) -} - -def buildDocker(install_prefix){ - buildDockerBase(install_prefix) - if (params.RUN_PYTORCH_TESTS.toBoolean()) { - buildDockerPytorch(install_prefix) - } - if (params.RUN_AITER_TESTS.toBoolean()) { - buildDockerAiter(install_prefix) - } - if (params.RUN_FA_TESTS.toBoolean()) { - buildDockerFa(install_prefix) - } -} - -def get_docker_options(){ - def dockerOpts - if ( params.BUILD_INSTANCES_ONLY ){ - dockerOpts = "--network=host --group-add video --group-add render --cap-add=SYS_PTRACE --security-opt seccomp=unconfined" - } - else{ //only add kfd and dri paths if you actually going to run somthing on GPUs - dockerOpts = "--network=host --device=/dev/kfd --device=/dev/dri --group-add video --group-add render --cap-add=SYS_PTRACE --security-opt seccomp=unconfined" - } - if (params.COMPILER_VERSION == "develop" || params.COMPILER_VERSION == "amd-staging" || params.COMPILER_VERSION == "therock" || params.COMPILER_COMMIT != ""){ - // the --env COMPRESSED_BUNDLE_FORMAT_VERSION=2 env variable is required when building code with offload-compress flag with - // newer clang22 compilers and running with older hip runtima libraries - dockerOpts = dockerOpts + " --env HIP_CLANG_PATH='/llvm-project/build/bin' --env COMPRESSED_BUNDLE_FORMAT_VERSION=2 --env HIP_PLATFORM=amd " - } - // on some machines the group ids for video and render groups may not be the same as in the docker image! - def video_id = sh(returnStdout: true, script: 'getent group video | cut -d: -f3') - def render_id = sh(returnStdout: true, script: 'getent group render | cut -d: -f3') - dockerOpts = dockerOpts + " --group-add=${video_id} --group-add=${render_id} -v /var/jenkins/ref-repo/:/var/jenkins/ref-repo/ " - echo "Docker flags: ${dockerOpts}" - return dockerOpts -} - -def build_client_examples(String arch){ - def cmd = """ cd ../client_example && rm -rf build && mkdir build && cd build && \ - cmake -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/projects/composablekernel/install;/opt/rocm" \ - -DGPU_TARGETS="${arch}" \ - -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_CXX_FLAGS=" -O3 " .. && make -j """ - return cmd -} - -def build_client_examples_and_codegen_tests(String arch){ - def cmd = """ cd ../codegen && rm -rf build && mkdir build && cd build && \ - cmake -DCMAKE_PREFIX_PATH=/opt/rocm -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" .. && \ - make -j64 check && \ - cd ../../client_example && rm -rf build && mkdir build && cd build && \ - cmake -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/projects/composablekernel/install;/opt/rocm" \ - -DGPU_TARGETS="${arch}" \ - -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_CXX_FLAGS=" -O3 " .. && make -j """ - return cmd -} - -def build_and_run_fmha(String arch){ - def cmd = """ cmake -G Ninja -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/projects/composablekernel/install;/opt/rocm" \ - -DGPU_TARGETS="${arch}" \ - -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" .. && \ - ninja -j128 tile_example_fmha_fwd tile_example_fmha_bwd && \ - cd ../ && - example/ck_tile/01_fmha/script/run_full_test.sh "CI_${params.COMPILER_VERSION}" "${env.BRANCH_NAME}" "${NODE_NAME}" "${arch}" """ - return cmd -} - -def cmake_build(Map conf=[:]){ - - def config_targets = conf.get("config_targets","check") - def build_envs = "CTEST_PARALLEL_LEVEL=4 " + conf.get("build_env","") - def prefixpath = conf.get("prefixpath","/opt/rocm") - def setup_args = conf.get("setup_args","") - // make sure all unit tests always run on develop branch - def runAllUnitTests = (env.BRANCH_NAME == "develop") ? true : params.RUN_ALL_UNIT_TESTS - - if (prefixpath != "/usr/local"){ - setup_args = setup_args + " -DCMAKE_PREFIX_PATH=${prefixpath} " - } - - //cmake_env can overwrite default CXX variables. - def cmake_envs - if(!setup_args.contains("gfx1250")){ - cmake_envs = "CXX=${params.BUILD_COMPILER} CXXFLAGS='-Werror' " + conf.get("cmake_ex_env","") - } - else{ //use default compiler for gfx1250 - cmake_envs = "CXX=/opt/rocm/llvm/bin/clang++ CXXFLAGS='-Werror' " + conf.get("cmake_ex_env","") - } - - if(conf.get("build_install","") == "true") - { - config_targets = 'install ' + config_targets - setup_args = ' -DBUILD_DEV=On -DCMAKE_INSTALL_PREFIX=../install' + setup_args - } else{ - setup_args = ' -DBUILD_DEV=On' + setup_args - } - if (params.DISABLE_DL_KERNELS){ - setup_args = setup_args + " -DDISABLE_DL_KERNELS=ON " - } - - setup_args = " -DCMAKE_BUILD_TYPE=release " + setup_args - - def pre_setup_cmd = """ - #!/bin/bash - cd projects/composablekernel - ulimit -c unlimited - rm -rf build - mkdir build - rm -rf install - mkdir install - cd build - """ - def invocation_tag="" - if (setup_args.contains("gfx12")){ - invocation_tag="gfx12" - } - if (setup_args.contains("gfx11")){ - invocation_tag="gfx11" - } - if (setup_args.contains("gfx101")){ - invocation_tag="gfx101" - } - if (setup_args.contains("gfx103")){ - invocation_tag="gfx103" - } - if (setup_args.contains("gfx908")){ - invocation_tag="gfx908" - } - if (setup_args.contains("gfx90a")){ - invocation_tag="gfx90a" - } - if (setup_args.contains("gfx94")){ - invocation_tag="gfx94" - } - if (setup_args.contains("gfx95")){ - invocation_tag="gfx95" - } - echo "invocation tag: ${invocation_tag}" - def redis_pre_setup_cmd = pre_setup_cmd - if(check_host() && params.USE_SCCACHE && "${env.CK_SCCACHE}" != "null" && "${invocation_tag}" != "") { - redis_pre_setup_cmd = pre_setup_cmd + """ - #!/bin/bash - export ROCM_PATH=/opt/rocm - export SCCACHE_ENABLED=true - export SCCACHE_LOG_LEVEL=debug - export SCCACHE_IDLE_TIMEOUT=14400 - export COMPILERS_HASH_DIR=/tmp/.sccache - export SCCACHE_BIN=/usr/local/.cargo/bin/sccache - export SCCACHE_EXTRAFILES=/tmp/.sccache/rocm_compilers_hash_file - export SCCACHE_REDIS="redis://${env.CK_SCCACHE}" - echo "connect = ${env.CK_SCCACHE}" >> ../script/redis-cli.conf - export SCCACHE_C_CUSTOM_CACHE_BUSTER="${invocation_tag}" - echo \$SCCACHE_C_CUSTOM_CACHE_BUSTER - stunnel ../script/redis-cli.conf - ../script/sccache_wrapper.sh --enforce_redis - """ - try { - def cmd1 = conf.get("cmd1", """ - ${redis_pre_setup_cmd} - """) - sh cmd1 - setup_args = " -DCMAKE_HIP_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache -DCMAKE_C_COMPILER_LAUNCHER=sccache " + setup_args - } - catch(Exception err){ - echo "could not connect to redis server: ${err.getMessage()}. will not use sccache." - def cmd2 = conf.get("cmd2", """ - ${pre_setup_cmd} - """) - sh cmd2 - } - } - else{ - def cmd3 = conf.get("cmd3", """ - ${pre_setup_cmd} - """) - sh cmd3 - } - - // reduce parallelism when compiling, clang uses too much memory - def nt = nthreads() - def cmd - def setup_cmd - def build_cmd - def execute_cmd = conf.get("execute_cmd", "") - //check the node gpu architecture - def arch_name = check_arch_name() - if(!setup_args.contains("NO_CK_BUILD")){ - if (params.NINJA_BUILD_TRACE) { - echo "running ninja build trace" - } - if (params.RUN_BUILDER_TESTS && !setup_args.contains("-DCK_CXX_STANDARD=") && !setup_args.contains("gfx10") && !setup_args.contains("gfx11")) { - setup_args = " -D CK_EXPERIMENTAL_BUILDER=ON " + setup_args - } - if (params.RUN_ROCM_CK_TESTS) { - setup_args = " -D CK_ENABLE_ROCM_CK=ON " + setup_args - } - setup_cmd = conf.get( - "setup_cmd", - """${cmake_envs} cmake -G Ninja ${setup_args} -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DCMAKE_CXX_FLAGS=" -O3 " .. """ - ) - - // Smart-build: Only build if running all tests or forced - // Otherwise, smart-build will determine what to build after cmake configure - if (runAllUnitTests) { - build_cmd = conf.get( - "build_cmd", - "${build_envs} ninja -j${nt} ${config_targets}" - ) - } else { - // Smart-build enabled: skip full build and execute_cmd (client examples) - build_cmd = "" - execute_cmd = "" - } - - cmd = conf.get("cmd", """ - ${setup_cmd} - ${build_cmd} - ${execute_cmd} - """) - } - else{ - cmd = conf.get("cmd", """ - ${execute_cmd} - """) - } - - echo cmd - - dir("projects/composablekernel/build"){ - // Start sccache monitoring - if(check_host() && params.USE_SCCACHE && "${env.CK_SCCACHE}" != "null" && "${invocation_tag}" != "") { - sh """ - chmod +x ../script/monitor_sccache_during_build.sh - mkdir -p logs - export SCCACHE_C_CUSTOM_CACHE_BUSTER="${invocation_tag}" - ../script/monitor_sccache_during_build.sh build_monitor & - MONITOR_PID=\$! - echo "Monitor PID: \$MONITOR_PID" - echo \$MONITOR_PID > monitor.pid - """ - } - try { - //build CK - sh cmd - if (runAllUnitTests){ - // Archive artifacts if they were generated - if (fileExists("ck_build_trace_${arch_name}.json")) { - archiveArtifacts "ck_build_trace_${arch_name}.json" - } - if (fileExists("clang_build_analysis_${arch_name}.log")) { - archiveArtifacts "clang_build_analysis_${arch_name}.log" - } - // Process ninja build trace after full build - if(fileExists(".ninja_log")) { - sh "python3 ../script/ninja_json_converter.py .ninja_log --legacy-format --output ck_build_trace_${arch_name}.json" - archiveArtifacts "ck_build_trace_${arch_name}.json" - sh "python3 ../script/parse_ninja_trace.py ck_build_trace_${arch_name}.json" - } - - if (params.NINJA_FTIME_TRACE) { - echo "running ClangBuildAnalyzer" - sh "/ClangBuildAnalyzer/build/ClangBuildAnalyzer --all . clang_build.log" - sh "/ClangBuildAnalyzer/build/ClangBuildAnalyzer --analyze clang_build.log > clang_build_analysis_${arch_name}.log" - archiveArtifacts "clang_build_analysis_${arch_name}.log" - } - } - } catch (Exception buildError) { - echo "Build failed: ${buildError.getMessage()}" - throw buildError - } finally { - // Stop sccache monitoring - if(check_host() && params.USE_SCCACHE && "${env.CK_SCCACHE}" != "null" && "${invocation_tag}" != "") { - sh """ - # Stop monitoring - if [ -f monitor.pid ]; then - MONITOR_PID=\$(cat monitor.pid) - kill \$MONITOR_PID 2>/dev/null || echo "Monitor already stopped" - rm -f monitor.pid - fi - """ - - // Archive the monitoring logs - try { - archiveArtifacts artifacts: "logs/*monitor*.log", allowEmptyArchive: true - } catch (Exception e) { - echo "Could not archive sccache monitoring logs: ${e.getMessage()}" - } - } - } - - //run tests except when NO_CK_BUILD is set and except on gfx1250 - if(!setup_args.contains("NO_CK_BUILD")){ - // run unit tests unless building library for all targets - // Note: This else block is when NINJA_BUILD_TRACE=false and BUILD_INSTANCES_ONLY=false - // So no ninja trace processing needed here - if (!params.BUILD_INSTANCES_ONLY){ - if (!runAllUnitTests && !setup_args.contains("gfx1250") ){ - // Smart Build: Run smart_build_and_test.sh - sh """ - export WORKSPACE_ROOT=${env.WORKSPACE} - export PARALLEL=32 - export NINJA_JOBS=${nt} - export ARCH_NAME=${arch_name} - export PROCESS_NINJA_TRACE=false - export NINJA_FTIME_TRACE=false - bash ../script/dependency-parser/smart_build_and_test.sh - """ - } - else{ //run all tests - if(!setup_args.contains("gfx1250")){ - echo "Full test suite requested (RUN_ALL_UNIT_TESTS=true or develop branch)" - sh "ninja -j${nt} check" - } - else{ //do not run tests on gfx1250, just build everything - echo "Building for gfx1250" - sh "ninja -j${nt}" - } - if (params.RUN_ROCM_CK_TESTS) { - sh 'ninja check-rocm-ck' - } - if(params.BUILD_PACKAGES || params.BUILD_INSTANCES_ONLY){ - echo "Build ckProfiler packages" - sh 'ninja -j64 package' - sh "mv composablekernel-ckprofiler_*.deb composablekernel-ckprofiler_1.2.0_amd64_${arch_name}.deb" - stash includes: "composablekernel-ckprofiler**.deb", name: "profiler_package_${arch_name}" - } - } - if (params.RUN_BUILDER_TESTS && !setup_args.contains("-DCK_CXX_STANDARD=") && !setup_args.contains("gfx10") && !setup_args.contains("gfx11")) { - sh 'ninja check-builder' - } - } - } - } - - if (params.RUN_CK_TILE_FMHA_TESTS){ - try{ - dir("projects/composablekernel"){ - archiveArtifacts "perf_fmha_*.log" - stash includes: "perf_fmha_**.log", name: "perf_fmha_log_${arch_name}" - } - } - catch(Exception err){ - echo "could not locate the requested artifacts: ${err.getMessage()}. will skip the stashing." - } - } -} - -def buildHipClangJob(Map conf=[:]){ - show_node_info() - checkoutComposableKernel() - def prefixpath = conf.get("prefixpath", "/opt/rocm") - def dockerOpts = get_docker_options() - def image - def retimage - (retimage, image) = getDockerImage(conf) - - setGithubStatus("${env.STAGE_NAME}", 'pending', "Starting ${env.STAGE_NAME}") - try { - withDockerContainer(image: image, args: dockerOpts) { - timeout(time: 20, unit: 'HOURS') - { - cmake_build(conf) - } - } - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - return retimage -} - -def buildHipClangJobAndReboot(Map conf=[:]){ - try{ - buildHipClangJob(conf) - } - catch(e){ - echo "throwing error exception for the stage" - echo 'Exception occurred: ' + e.toString() - throw e - } -} - -def Build_CK(Map conf=[:]){ - show_node_info() - checkoutComposableKernel() - def prefixpath = conf.get("prefixpath", "/opt/rocm") - def dockerOpts=get_docker_options() - def image - def retimage - - setGithubStatus("${env.STAGE_NAME}", 'pending', "Starting ${env.STAGE_NAME}") - try { - try { - (retimage, image) = getDockerImage(conf) - withDockerContainer(image: image, args: dockerOpts) { - timeout(time: 2, unit: 'MINUTES'){ - sh 'rocminfo | tee rocminfo.log' - if ( !runShell('grep -n "gfx" rocminfo.log') ){ - throw new Exception ("GPU not found") - } - else{ - echo "GPU is OK" - } - } - } - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - echo "The job was cancelled or aborted" - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - withDockerContainer(image: image, args: dockerOpts) { - timeout(time: 20, unit: 'HOURS') - { - //check whether to run performance tests on this node - def arch = check_arch_name() - cmake_build(conf) - if ( params.RUN_INDUCTOR_TESTS && arch == "gfx90a" ){ - echo "Run inductor codegen tests" - sh "projects/composablekernel/script/run_inductor_tests.sh" - } - // run performance tests, stash the logs, results will be processed on the master node - dir("projects/composablekernel/script"){ - if (params.RUN_PERFORMANCE_TESTS){ - if (params.RUN_FULL_QA && (arch == "gfx90a" || arch == "gfx942")){ - // run full tests on gfx90a or gfx942 - echo "Run full performance tests" - sh "./run_full_performance_tests.sh 0 QA_${params.COMPILER_VERSION} ${env.BRANCH_NAME} ${NODE_NAME} ${arch}" - archiveArtifacts "perf_*.log" - stash includes: "perf_**.log", name: "perf_log_${arch}" - } - else if (!params.RUN_FULL_QA && (arch == "gfx90a" || arch == "gfx942")){ - // run standard tests on gfx90a or gfx942 - echo "Run performance tests" - sh "./run_performance_tests.sh 0 CI_${params.COMPILER_VERSION} ${env.BRANCH_NAME} ${NODE_NAME} ${arch}" - archiveArtifacts "perf_*.log" - stash includes: "perf_**.log", name: "perf_log_${arch}" - } - else if ( arch != "gfx10"){ - // run basic tests on gfx11/gfx12/gfx908/gfx950, but not on gfx10, it takes too long - echo "Run gemm performance tests" - sh "./run_gemm_performance_tests.sh 0 CI_${params.COMPILER_VERSION} ${env.BRANCH_NAME} ${NODE_NAME} ${arch}" - archiveArtifacts "perf_onnx_gemm_*.log" - stash includes: "perf_onnx_gemm_**.log", name: "perf_log_${arch}" - } - } - } - if (params.hipTensor_test && arch == "gfx90a" ){ - // build and test hipTensor on gfx90a node - sh """#!/bin/bash - rm -rf rocm-libraries - git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-libraries.git - cd rocm-libraries - git sparse-checkout init --cone - git sparse-checkout set projects/hiptensor - git checkout "${params.hipTensor_branch}" - """ - dir("rocm-libraries/projects/hiptensor"){ - sh """#!/bin/bash - mkdir -p build - ls -ltr - CC=hipcc CXX=hipcc cmake -Bbuild . -D CMAKE_PREFIX_PATH="${env.WORKSPACE}/install" - cmake --build build -- -j - ctest --test-dir build - """ - } - } - } - } - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") +def loadCk() { + if (ck == null) { + checkout scm + dir("projects/composablekernel") { + ck = load "vars/ck.groovy" } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - return retimage -} - -def Build_CK_and_Reboot(Map conf=[:]){ - try{ - Build_CK(conf) - } - catch(e){ - echo "throwing error exception while building CK" - echo 'Exception occurred: ' + e.toString() - throw e } } -def process_results(Map conf=[:]){ - checkoutComposableKernel() - //use older image that has user jenkins - def image = "${env.CK_DOCKERHUB}:ck_ub22.04_rocm6.3" - - setGithubStatus("${env.STAGE_NAME}", 'pending', 'Processing results...') - try { - try - { - echo "Pulling image: ${image}" - def retimage = docker.image("${image}") - withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { - retimage.pull() - } - } - catch(Exception ex) - { - error "Unable to locate image: ${image}" - } - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - - withDockerContainer(image: image, args: '--cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v=/var/jenkins/:/var/jenkins') { - timeout(time: 15, unit: 'MINUTES'){ - try{ - dir("projects/composablekernel/script"){ - if (params.RUN_CK_TILE_FMHA_TESTS){ - try{ - unstash "perf_fmha_log_gfx942" - } - catch(Exception err){ - echo "could not locate the FMHA performance logs for gfx942: ${err.getMessage()}." - } - try{ - unstash "perf_fmha_log_gfx90a" - } - catch(Exception err){ - echo "could not locate the FMHA performance logs for gfx90a: ${err.getMessage()}." - } - try{ - unstash "perf_fmha_log_gfx950" - } - catch(Exception err){ - echo "could not locate the FMHA performance logs for gfx950: ${err.getMessage()}." - } - - } - if (params.BUILD_INSTANCES_ONLY){ - // unstash deb packages - try{ - unstash "lib_package" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate lib_package." - } - } - if (params.BUILD_PACKAGES){ - // unstash deb packages - try{ - unstash "profiler_package_gfx90a" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate profiler_package_gfx90a." - } - try{ - unstash "profiler_package_gfx942" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate profiler_package_gfx942." - } - try{ - unstash "profiler_package_gfx950" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate profiler_package_gfx950." - } - try{ - unstash "profiler_package_gfx12" - sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" - } - catch(Exception err){ - echo "could not locate profiler_package_gfx12." - } - } - else{ - // unstash perf files to master - try{ - unstash "perf_log_gfx90a" - } - catch(Exception err){ - echo "could not locate the gfx90a performance logs: ${err.getMessage()}." - } - try{ - unstash "perf_log_gfx942" - } - catch(Exception err){ - echo "could not locate the gfx942 performance logs: ${err.getMessage()}." - } - try{ - unstash "perf_log_gfx950" - } - catch(Exception err){ - echo "could not locate the gfx950 performance logs: ${err.getMessage()}." - } - try{ - unstash "perf_log_gfx908" - } - catch(Exception err){ - echo "could not locate the gfx908 performance logs: ${err.getMessage()}." - } - try{ - unstash "perf_log_gfx11" - } - catch(Exception err){ - echo "could not locate the gfx11 performance logs: ${err.getMessage()}." - } - try{ - - unstash "perf_log_gfx12" - } - catch(Exception err){ - echo "could not locate the gfx12 performance logs: ${err.getMessage()}." - } - } - // process the logs - sh "./process_perf_data.sh" - } - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - finally{ - echo "Finished processing performance test results" - } - } - } -} - -def run_downstream_tests(Map conf=[:]){ - show_node_info() - checkoutComposableKernel() - def dockerOpts = get_docker_options() + ' --group-add irc ' - - setGithubStatus("${env.STAGE_NAME}", 'pending', "Starting ${env.STAGE_NAME}") - try { - try - { - echo "Pulling image: ${conf.image}" - retimage = docker.image("${conf.image}") - withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { - retimage.pull() - } - } - catch(Exception ex) - { - error "Unable to locate image: ${conf.image}" - } - } - catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ - setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") - throw e - } - - withDockerContainer(image: conf.image, args: dockerOpts) { - timeout(time: conf.get("timeoutHours", 2), unit: 'HOURS'){ - try{ - sh "rocminfo" - sh "python3 --version" - for (cmd in conf.execute_cmds) { - sh "${cmd}" - } - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") - } - catch(e){ - echo "Throwing error exception while running ${env.STAGE_NAME}" - echo 'Exception occurred: ' + e.toString() - setGithubStatus("${env.STAGE_NAME}", 'error', "Stage ${env.STAGE_NAME} failed") - throw e - } - finally{ - echo "Finished running ${env.STAGE_NAME}" - } - } - } -} - -def getPytorchTestsCmds() { - return [ - "mkdir pytorch", - "cp -r /var/jenkins/workspace/pytorch/* pytorch/", - "ls -ltr pytorch", - "python3 pytorch/tools/amd_build/build_amd.py", - "cd pytorch && USE_ROCM_CK_SDPA=1 PYTORCH_ROCM_ARCH=gfx942 python3 setup.py develop" - ] -} -def getAiterTestsCmds() { - return [ - "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8_blockscale.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_mha.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_mha_varlen.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_batch_prefill.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_moe.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_2stage.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_blockscale.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_ep.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_sorting.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_sorting_mxfp4.py", - "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_tkw1.py" - ] -} -def getFaTestsCmds() { - return [ - "python3 -u -m pytest /home/jenkins/workspace/flash-attention/tests/test_flash_attn_ck.py" - ] -} - //launch develop branch daily jobs CRON_SETTINGS = BRANCH_NAME == "develop" ? '''0 23 * * * % RUN_FULL_QA=true;RUN_CK_TILE_FMHA_TESTS=true;RUN_PERFORMANCE_TESTS=true;FORCE_CI=true 0 22 * * * % RUN_FULL_QA=true;DISABLE_DL_KERNELS=true;RUN_TILE_ENGINE_BASIC_TESTS=true;RUN_TILE_ENGINE_GEMM_TESTS=true;RUN_PERFORMANCE_TESTS=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true @@ -1482,8 +273,9 @@ pipeline { agent{ label rocmnode("nogpu") } steps { script { - checkoutComposableKernel() - env.SHOULD_RUN_CI = String.valueOf(params.FORCE_CI.toBoolean() || shouldRunCICheck()) + loadCk() + ck.checkoutComposableKernel() + env.SHOULD_RUN_CI = String.valueOf(params.FORCE_CI.toBoolean() || ck.shouldRunCICheck()) echo "SHOULD_RUN_CI: ${env.SHOULD_RUN_CI}" } } @@ -1498,7 +290,10 @@ pipeline { agent{ label rocmnode("nogpu") } steps{ deleteDir() - buildDocker('/opt/rocm') + script { + loadCk() + ck.buildDocker('/opt/rocm') + } cleanWs() } } @@ -1516,21 +311,9 @@ pipeline { expression { params.RUN_CPPCHECK.toBoolean() } } agent{ label rocmnode("nogpu") } - environment{ - setup_args = "NO_CK_BUILD" - execute_cmd = """cd .. && \ - find . -type f \\( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.cl' \\) \ - -not -path '*/build/*' -not -path '*/include/rapidjson/*' | \ - xargs -P 8 -I{} sh -c 'clang-format-18 -style=file {} | diff -u - {} || (echo "ERROR: {} needs formatting" && exit 1)' && \ - /cppcheck/build/bin/cppcheck ../* -v -j \$(nproc) -I ../include -I ../profiler/include -I ../library/include \ - -D CK_ENABLE_FP64 -D CK_ENABLE_FP32 -D CK_ENABLE_FP16 -D CK_ENABLE_FP8 -D CK_ENABLE_BF16 -D CK_ENABLE_BF8 -D CK_ENABLE_INT8 \ - -D __gfx908__ -D __gfx90a__ -D __gfx942__ -D __gfx1030__ -D __gfx1100__ -D __gfx1101__ -D __gfx1102__ \ - -U __gfx803__ -U __gfx900__ -U __gfx906__ -U CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4 \ - --file-filter=*.cpp --force --enable=all --output-file=ck_cppcheck.log""" - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, setup_cmd: "", build_cmd: "", execute_cmd: execute_cmd) + script { loadCk(); ck.runClangFormatAndCppcheck() } archiveArtifacts "build/ck_cppcheck.log" cleanWs() } @@ -1541,16 +324,9 @@ pipeline { expression { !params.RUN_CPPCHECK.toBoolean() } } agent{ label rocmnode("nogpu") } - environment{ - setup_args = "NO_CK_BUILD" - execute_cmd = """cd .. && \ - find . -type f \\( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.cl' \\) \ - -not -path '*/build/*' -not -path '*/include/rapidjson/*' | \ - xargs -P 8 -I{} sh -c 'clang-format-18 -style=file {} | diff -u - {} || (echo "ERROR: {} needs formatting" && exit 1)'""" - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, setup_cmd: "", build_cmd: "", execute_cmd: execute_cmd) + script { loadCk(); ck.runClangFormat() } cleanWs() } } @@ -1572,7 +348,10 @@ pipeline { } agent{ label rocmnode("gfx942")} steps{ - run_downstream_tests(image: "${env.CK_PYTORCH_IMAGE}", timeoutHours: 2, execute_cmds: getPytorchTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_PYTORCH_IMAGE}", timeoutHours: 2, execute_cmds: ck.getPytorchTestsCmds()) + } cleanWs() } } @@ -1584,7 +363,10 @@ pipeline { } agent{ label rocmnode("gfx942")} steps{ - run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: getAiterTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: ck.getAiterTestsCmds()) + } cleanWs() } } @@ -1596,7 +378,10 @@ pipeline { } agent{ label rocmnode("gfx950")} steps{ - run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: getAiterTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: ck.getAiterTestsCmds()) + } cleanWs() } } @@ -1608,7 +393,10 @@ pipeline { } agent{ label rocmnode("gfx942")} steps{ - run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: getFaTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: ck.getFaTestsCmds()) + } cleanWs() } } @@ -1620,7 +408,10 @@ pipeline { } agent{ label rocmnode("gfx950")} steps{ - run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: getFaTestsCmds()) + script { + loadCk() + ck.run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: ck.getFaTestsCmds()) + } cleanWs() } } @@ -1641,17 +432,9 @@ pipeline { expression { params.RUN_FULL_CONV_TILE_TESTS.toBoolean() } } agent{ label rocmnode("gfx90a")} - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ python3 ../experimental/grouped_convolution_tile_instances/generate_instances.py --mode=profiler && \ - cmake .. --preset dev-gfx90a -D CK_EXPERIMENTAL_BUILDER=ON && \ - make -j64 test_grouped_convnd_fwd_tile test_grouped_convnd_bwd_weight_tile && \ - ./bin/test_grouped_convnd_bwd_weight_tile && \ - ./bin/test_grouped_convnd_fwd_tile""" - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runFullGroupedConvTileTests() } cleanWs() } } @@ -1672,15 +455,9 @@ pipeline { expression { params.RUN_GROUPED_CONV_LARGE_CASES_TESTS.toBoolean() } } agent{ label rocmnode("gfx90a")} - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake .. --preset dev-gfx90a && \ - make -j64 test_grouped_convnd_fwd_large_cases test_grouped_convnd_bwd_data_large_cases test_grouped_convnd_fwd_bias_clamp_large_cases && \ - ./bin/test_grouped_convnd_fwd_large_cases && ./bin/test_grouped_convnd_bwd_data_large_cases && ./bin/test_grouped_convnd_fwd_bias_clamp_large_cases""" - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runGroupedConvLargeCaseTests() } cleanWs() } } @@ -1701,27 +478,9 @@ pipeline { expression { params.RUN_CONV_COMPREHENSIVE_DATASET.toBoolean() } } agent{ label rocmnode("gfx90a")} - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cd ../build && \ - cmake .. --preset dev-gfx90a && \ - make -j64 test_grouped_convnd_fwd_dataset_xdl && \ - test_grouped_convnd_bwd_data_dataset_xdl \ - test_grouped_convnd_bwd_weight_dataset_xdl && \ - cd ../test_data && \ - # Dataset generation modes: - # - small: ~60 test cases (minimal, quick testing - 3 models, 2 batch sizes, 2 image sizes) - # - half: ~300 test cases (moderate coverage - 16 models, 3 batch sizes, 5 image sizes), ~ 17 hours testing time - # - full: ~600 test cases (comprehensive - 16 models, 5 batch sizes, 9 image sizes), ~ 40 hours testing time - ./generate_test_dataset.sh small && \ - cd ../build && \ - ./bin/test_grouped_convnd_fwd_dataset_xdl && \ - ./bin/test_grouped_convnd_bwd_data_dataset_xdl && \ - ./bin/test_grouped_convnd_bwd_weight_dataset_xdl""" - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runComprehensiveConvDatasetTests() } cleanWs() } } @@ -1744,11 +503,14 @@ pipeline { agent{ label rocmnode("gfx90a") } environment{ setup_args = "NO_CK_BUILD" - execute_args = build_and_run_fmha("gfx90a") + execute_args = ck.build_and_run_fmha("gfx90a") } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { + loadCk() + ck.buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + } cleanWs() } } @@ -1761,11 +523,14 @@ pipeline { agent{ label rocmnode("gfx942") } environment{ setup_args = "NO_CK_BUILD" - execute_args = build_and_run_fmha("gfx942") + execute_args = ck.build_and_run_fmha("gfx942") } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { + loadCk() + ck.buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + } cleanWs() } } @@ -1778,11 +543,14 @@ pipeline { agent{ label rocmnode("gfx950") } environment{ setup_args = "NO_CK_BUILD" - execute_args = build_and_run_fmha("gfx950") + execute_args = ck.build_and_run_fmha("gfx950") } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { + loadCk() + ck.buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + } cleanWs() } } @@ -1795,11 +563,14 @@ pipeline { agent{ label rocmnode("gfx1201") } environment{ setup_args = "NO_CK_BUILD" - execute_args = build_and_run_fmha("gfx1201") + execute_args = ck.build_and_run_fmha("gfx1201") } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { + loadCk() + ck.buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + } cleanWs() } } @@ -1820,30 +591,9 @@ pipeline { expression { params.RUN_TILE_ENGINE_BASIC_TESTS.toBoolean() } } agent{ label rocmnode("gfx942") } - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx942" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_UNIVERSAL_CONFIG_FILE="default_ci_config.json" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_MULTI_D_CONFIG_FILE="default_ci_config.json" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ - -D GEMM_PRESHUFFLE_CONFIG_FILE="default_ci_config.json" .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json """ - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runTileEngineBasicTests(params.BUILD_COMPILER) } cleanWs() } } @@ -1864,33 +614,9 @@ pipeline { expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } } agent{ label rocmnode("gfx942") } - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx942" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16;bf8;bf16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_STREAMK_DATATYPE="fp8;fp16" \ - -D GEMM_STREAMK_LAYOUT="rcr" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ - -D GROUPED_GEMM_DATATYPE="fp8;fp16" \ - -D GROUPED_GEMM_LAYOUT="rcr;rrr;crr;ccr" \ - -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all benchmark_gemm_streamk_all benchmark_grouped_gemm_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_universal_results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/grouped_gemm/grouped_gemm_benchmark.py . --problem-sizes "1024,1024,1024" --group-counts 8 --warmup 5 --repeat 5 --verbose --json grouped_gemm_results.json """ - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runTileEngineGemmTests("gfx942", params.BUILD_COMPILER) } cleanWs() } } @@ -1901,28 +627,9 @@ pipeline { expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } } agent{ label rocmnode("gfx950") } - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx950" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ - -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json """ - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runTileEngineGemmTests("gfx950", params.BUILD_COMPILER) } cleanWs() } } @@ -1933,22 +640,9 @@ pipeline { expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } } agent{ label rocmnode("gfx1201") } - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx1201" \ - -D GEMM_UNIVERSAL_DATATYPE="fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json """ - } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args:setup_args, build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runTileEngineGemmTests("gfx1201", params.BUILD_COMPILER) } cleanWs() } } @@ -1970,13 +664,9 @@ pipeline { expression { (params.BUILD_GFX942.toBoolean() || params.RUN_FULL_QA.toBoolean()) && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx942") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx942" """ - execute_args = build_client_examples("gfx942") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx942") } cleanWs() } } @@ -1987,13 +677,9 @@ pipeline { expression { params.BUILD_GFX950.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx950") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx950" """ - execute_args = build_client_examples("gfx950") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx950") } cleanWs() } } @@ -2005,13 +691,9 @@ pipeline { expression { params.BUILD_GFX908.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx908") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx908" """ - execute_args = build_client_examples("gfx908") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx908") } cleanWs() } } @@ -2023,13 +705,9 @@ pipeline { expression { params.BUILD_GFX90A.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx90a") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx90a" -DCK_CXX_STANDARD="17" """ - execute_args = build_client_examples_and_codegen_tests("gfx90a") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx90a") } cleanWs() } } @@ -2050,7 +728,7 @@ pipeline { } steps{ deleteDir() - buildHipClangJobAndReboot(setup_args: setup_args, build_cmd: "", build_type: 'Release', execute_cmd: execute_args) + script { loadCk(); ck.runBuildInstancesOnly(params.BUILD_COMPILER) } cleanWs() } } @@ -2062,13 +740,9 @@ pipeline { expression { params.BUILD_GFX101.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx1010") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx10-1-generic" """ - execute_args = build_client_examples("gfx10-1-generic") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx10-1-generic") } cleanWs() } } @@ -2080,13 +754,9 @@ pipeline { expression { params.BUILD_GFX103.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx1030") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx10-3-generic" """ - execute_args = build_client_examples("gfx10-3-generic") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx10-3-generic") } cleanWs() } } @@ -2097,13 +767,9 @@ pipeline { expression { params.BUILD_GFX11.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label 'miopen && (gfx1101 || gfx1100)' } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx11-generic" """ - execute_args = build_client_examples("gfx11-generic") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx11-generic") } cleanWs() } } @@ -2114,13 +780,9 @@ pipeline { expression { params.BUILD_GFX12.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx1201") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx12-generic" """ - execute_args = build_client_examples("gfx12-generic") - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx12-generic") } cleanWs() } } @@ -2131,12 +793,9 @@ pipeline { expression { params.BUILD_GFX1250.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } } agent{ label rocmnode("gfx90a") } - environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx1250" -DDISABLE_DL_KERNELS="ON" """ - } steps{ deleteDir() - Build_CK_and_Reboot(setup_args: setup_args, docker_name: "${env.CK_DOCKERHUB_PRIVATE}:npi-mi450-latest", config_targets: "install", no_reboot:true, build_type: 'Release', prefixpath: '/usr/local') + script { loadCk(); ck.runBuildCKAndTests("gfx1250") } cleanWs() } } @@ -2145,12 +804,13 @@ pipeline { always { node(rocmnode("nogpu")) { script { + loadCk() // Simulate capture - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx11.json") - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx12.json") - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx90a.json") - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx942.json") - generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx950.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx11.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx12.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx90a.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx942.json") + ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx950.json") } cleanWs() } @@ -2158,8 +818,9 @@ pipeline { success { script { node(rocmnode("nogpu")) { + loadCk() // Report the parent stage build ck and run tests status - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") + ck.setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") echo "Reporting success status for build ck and run tests" } } @@ -2178,7 +839,10 @@ pipeline { agent { label 'mici' } steps{ deleteDir() - process_results() + script { + loadCk() + ck.process_results() + } cleanWs() } } @@ -2187,8 +851,9 @@ pipeline { success { script { node(rocmnode("nogpu")) { + loadCk() // Report the skipped parent's stage status - setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") + ck.setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") echo "Process Performance Test Results stage skipped." } } @@ -2200,17 +865,17 @@ pipeline { success { script { node(rocmnode("nogpu")) { - setGithubStatus('Math CI Summary', 'success', "Math CI passed") + loadCk() + ck.setGithubStatus('Math CI Summary', 'success', "Math CI passed") } } } failure { script { node(rocmnode("nogpu")) { - setGithubStatus('Math CI Summary', 'failure', "Math CI failed") - script { - checkoutComposableKernel() - } + loadCk() + ck.setGithubStatus('Math CI Summary', 'failure', "Math CI failed") + ck.checkoutComposableKernel() withCredentials([string(credentialsId: 'ck_ci_errors_webhook_url', variable: 'WEBHOOK_URL')]) { sh 'bash projects/composablekernel/script/infra_helper/send_failure_notifications.sh' } diff --git a/vars/ck.groovy b/vars/ck.groovy new file mode 100644 index 0000000000..e2caf8007f --- /dev/null +++ b/vars/ck.groovy @@ -0,0 +1,1414 @@ +@NonCPS +String getGitHubCommitHash(def build) +{ + def scmAction = build?.actions.find { action -> + action instanceof jenkins.scm.api.SCMRevisionAction + } + if (scmAction?.revision instanceof org.jenkinsci.plugins.github_branch_source.PullRequestSCMRevision) + { + return scmAction.revision.pullHash + } + else if (scmAction?.revision instanceof jenkins.plugins.git.AbstractGitSCMSource$SCMRevisionImpl) + { + return scmAction.revision.hash + } + return null +} + +def show_node_info() { + sh """ + echo "NODE_NAME = \$NODE_NAME" + hostname + lsb_release -sd + uname -r + cat /sys/module/amdgpu/version + ls /opt/ -la + """ +} + +def setGithubStatus(String context, String state, String description) { + def sha = env.GIT_COMMIT + def targetUrl = env.RUN_DISPLAY_URL ?: env.BUILD_URL + def statusUrl = "https://api.github.com/repos/ROCm/rocm-libraries/statuses/${sha}" + withCredentials([usernamePassword(credentialsId: 'github-app-miopen', usernameVariable: 'GITHUB_APP', passwordVariable: 'GITHUB_TOKEN')]) { + def code = '0' + try { + retry(3) { + code = sh(returnStdout: true, script: """ + curl -s -w "%{http_code}" -o /dev/null -X POST '${statusUrl}' \\ + -H "Authorization: token \$GITHUB_TOKEN" \\ + -H 'Content-Type: application/json' \\ + -d '{"state":"${state}","context":"${context}","description":"${description}","target_url":"${targetUrl}"}' + """).trim() + if (!code.startsWith('2')) { + error("GitHub status POST returned ${code}") + } + } + } catch (Exception e) { + echo "WARNING: GitHub status POST failed after retries (context=${context}, state=${state}, code=${code})" + } + } +} + +def cloneUpdateRefRepo() { + def refRepoPath = "/var/jenkins/ref-repo/rocm-libraries" + def lockLabel = "git ref repo lock - ${env.NODE_NAME}" + def folderExists = sh( + script: "test -d ${refRepoPath}/refs", + returnStatus: true + ) == 0 + + if (!folderExists) { + echo "rocm-libraries repo does not exist at ${refRepoPath}, creating mirror clone..." + echo "locking on label: ${lockLabel}" + lock(lockLabel) { + def cloneCommand = """ + set -ex + rm -rf ${refRepoPath} && mkdir -p ${refRepoPath} + git clone --mirror https://github.com/ROCm/rocm-libraries.git ${refRepoPath} + """ + sh(script: cloneCommand, label: "clone ref repo") + } + echo "Completed git clone, lock released" + } + echo "rocm-libraries repo exists at ${refRepoPath}, performing git remote update..." + echo "locking on label: ${lockLabel}" + lock(lockLabel) { + def fetchCommand = """ + set -ex + cd ${refRepoPath} + git remote prune origin + git remote update + """ + sh(script: fetchCommand, label: "update ref repo") + } + echo "Completed git ref repo fetch, lock released" +} + +def checkoutComposableKernel() +{ + //update ref repo + cloneUpdateRefRepo() + // checkout project + def scmVars = checkout scm + // getGitHubCommitHash reads SCMRevisionAction recorded before any local merge, + // giving the true PR branch tip (pullHash) or branch HEAD (hash). + // Falls back to ORIG_HEAD (pre-merge HEAD set by git merge) when SCMRevisionAction + // is unavailable, then to HEAD for branch builds where no merge occurred. + env.GIT_COMMIT = getGitHubCommitHash(currentBuild.rawBuild) ?: sh(returnStdout: true, script: ''' + git rev-parse ORIG_HEAD 2>/dev/null || git rev-parse HEAD + ''').trim() +} + +def generateAndArchiveBuildTraceVisualization(String buildTraceFileName) { + try { + checkoutComposableKernel() + + // Retrieve the build trace artifact + def traceFileExists = false + try { + copyArtifacts( + projectName: env.JOB_NAME, + selector: specific(env.BUILD_NUMBER), + filter: buildTraceFileName + ) + traceFileExists = fileExists(buildTraceFileName) + } catch (Exception e) { + echo "Could not copy build trace artifact: ${e.getMessage()}" + traceFileExists = false + return + } + + sh """ + echo "post artifact download:" + ls -la + """ + + // Pull image + def image = "ghcr.io/puppeteer/puppeteer:24.30.0" + echo "Pulling image: ${image}" + def retimage = docker.image("${image}") + retimage.pull() + + // Create a temporary workspace + sh """#!/bin/bash + ls -la + mkdir -p workspace + cp ./projects/composablekernel/script/infra_helper/capture_build_trace.js ./workspace + cp ${buildTraceFileName} ./workspace/${buildTraceFileName} + chmod 777 ./workspace + ls -la ./workspace + """ + + // Run container to get snapshot + def dockerOpts = "--cap-add=SYS_ADMIN -v \"\$(pwd)/workspace:/workspace\" -e NODE_PATH=/home/pptruser/node_modules -e BUILD_TRACE_FILE=${buildTraceFileName}" + // Create unique image name by sanitizing job name + def sanitizedJobName = env.JOB_NAME.replaceAll(/[\/\\:*?"<>| ]/, '_').replaceAll('%2F', '_') + def architectureName = (buildTraceFileName =~ /(gfx[0-9a-zA-Z]+)/)[0][1] + def imageName = "perfetto_snapshot_${sanitizedJobName}_build_${env.BUILD_NUMBER}_${architectureName}.png" + sh """ + docker run --rm ${dockerOpts} ${image} node /workspace/capture_build_trace.js + mv ./workspace/perfetto_snapshot_build.png ./workspace/${imageName} + """ + + // Archive the snapshot + sh """ + mv ./workspace/${imageName} ${imageName} + """ + archiveArtifacts "${imageName}" + + // Notify the channel + withCredentials([string(credentialsId: 'ck_ci_build_perf_webhook_url', variable: 'WEBHOOK_URL')]) { + sh ''' + # Create build trace filename with build number based on the original filename + BUILD_TRACE_WITH_NUMBER=$(echo "''' + buildTraceFileName + '''" | sed 's/.json/_''' + sanitizedJobName + '''_''' + env.BUILD_NUMBER + '''_''' + architectureName + '''.json/') + + # Convert image to base64 + echo "Converting image to base64..." + IMAGE_BASE64=$(base64 -w 0 ''' + imageName + ''') + echo "Image base64 length: ${#IMAGE_BASE64}" + + # Convert build trace to base64 + echo "Converting build trace to base64..." + BUILD_TRACE_BASE64=$(base64 -w 0 ''' + buildTraceFileName + ''') + echo "Build trace base64 length: ${#BUILD_TRACE_BASE64}" + + # Create JSON payload with base64 data + echo "Creating JSON payload..." + { + printf '{\n' + printf ' "jobName": "%s",\n' "''' + env.JOB_NAME + '''" + printf ' "buildNumber": "%s",\n' "''' + env.BUILD_NUMBER + '''" + printf ' "jobUrl": "%s",\n' "''' + env.RUN_DISPLAY_URL + '''" + printf ' "imageName": "%s",\n' "''' + imageName + '''" + printf ' "architecture": "%s",\n' "''' + architectureName + '''" + printf ' "imageData": "%s",\n' "$IMAGE_BASE64" + printf ' "buildTraceName": "%s",\n' "$BUILD_TRACE_WITH_NUMBER" + printf ' "buildTraceData": "%s"\n' "$BUILD_TRACE_BASE64" + printf '}\n' + } > webhook_payload.json + + echo "JSON payload created, size: $(wc -c < webhook_payload.json) bytes" + + curl -X POST "${WEBHOOK_URL}" \ + -H "Content-Type: application/json" \ + -d @webhook_payload.json + + # Clean up temporary file + rm -f webhook_payload.json + ''' + } + } catch (Exception e) { + echo "Throwing error exception while generating build trace visualization" + echo 'Exception occurred: ' + e.toString() + } +} + +def nthreads() { + def nproc = sh(returnStdout: true, script: 'nproc') + echo "Number of cores: ${nproc}" + def n = nproc.toInteger() + if (n > 64){ + n = 64 + } + echo "Number of threads used for building: ${n}" + return n +} + +def runShell(String command){ + def responseCode = sh returnStatus: true, script: "${command} > tmp.txt" + def output = readFile(file: "tmp.txt") + return (output != "") +} + +def shouldRunCICheck() { + // File patterns that should not trigger CI + def skipFilePatterns = [ + /^projects\/composablekernel\/\.github\/.*/, // GitHub workflow files + /^projects\/composablekernel\/docs\/.*/, // Documentation files + /^projects\/composablekernel\/LICENSE$/, // License file + /^projects\/composablekernel\/.*\.gitignore$/, // Git ignore files + /^projects\/composablekernel\/.*\.md$/ // Markdown files + ] + + try { + // Always run if this is a base branch build + def baseBranch = "develop" + def isBaseBranchBuild = (env.CHANGE_ID == null && env.BRANCH_NAME == baseBranch) + + if (isBaseBranchBuild) { + echo "Base branch (${baseBranch}) build detected - always running CI for safety" + return true + } + + // Get the list of changed files (all files touched in any commit, even if reverted) + def changedFiles = sh( + returnStdout: true, + script: ''' + BASE_BRANCH="develop" + + if [ "$CHANGE_ID" != "" ]; then + # For PR builds, get all files touched in any commit + echo "PR build detected, checking all touched files against origin/$CHANGE_TARGET" >&2 + git log --name-only --pretty=format: origin/$CHANGE_TARGET..HEAD -- projects/composablekernel/ | sort -u | grep -v '^$' || true + else + # For feature branch builds, compare against merge-base with base branch + MERGE_BASE=$(git merge-base HEAD origin/$BASE_BRANCH 2>/dev/null || echo "HEAD~1") + echo "Branch build detected, checking all touched files since merge-base: $MERGE_BASE" >&2 + git log --name-only --pretty=format: $MERGE_BASE..HEAD -- projects/composablekernel/ | sort -u | grep -v '^$' || true + fi + ''' + ).trim().split('\n') + + if (changedFiles.size() == 1 && changedFiles[0] == '') { + echo "No changed files detected - this might be a manual trigger or merge commit, running CI for safety" + return true + } + + echo "Changed files: ${changedFiles.join(', ')}" + + // Separate files into those requiring CI and those that can be skipped + def filesRequiringCI = [] + def skippedFiles = [] + + changedFiles.each { file -> + def shouldSkip = skipFilePatterns.any { pattern -> + file ==~ pattern + } + + if (shouldSkip) { + skippedFiles.add(file) + } else { + filesRequiringCI.add(file) + } + } + + // Debug output + if (skippedFiles.size() > 0) { + echo "Files that don't require CI (${skippedFiles.size()}):" + skippedFiles.each { echo " - ${it}" } + } + + if (filesRequiringCI.size() > 0) { + echo "Files that require CI (${filesRequiringCI.size()}):" + filesRequiringCI.each { echo " - ${it}" } + return true + } else { + echo "Only non-relevant files changed, skipping CI" + return false + } + } catch (Exception e) { + echo "Error checking changed files: ${e.getMessage()}, running CI by default" + return true + } +} + +def getBaseDockerImageName(){ + def img + if (params.USE_CUSTOM_DOCKER != ""){ + img = "${params.USE_CUSTOM_DOCKER}" + } + else{ + img = "${env.CK_DOCKERHUB}:ck_ub24.04_rocm${params.ROCMVERSION}" + } + return img +} + +def getDockerImageName(){ + def img + def base_name = getBaseDockerImageName() + if (params.USE_CUSTOM_DOCKER != ""){ + img = "${params.USE_CUSTOM_DOCKER}" + } + else{ + if (params.COMPILER_VERSION == "") { + img = "${base_name}" + } + else{ + if (params.COMPILER_COMMIT == ""){ + img = "${base_name}_${params.COMPILER_VERSION}" + } + else{ + def commit = "${params.COMPILER_COMMIT}"[0..6] + img = "${base_name}_${params.COMPILER_VERSION}_${commit}" + } + } + } + return img +} + +def check_host() { + if ("${env.CK_SCCACHE}" != "null"){ + def SCCACHE_SERVER="${env.CK_SCCACHE.split(':')[0]}" + echo "sccache server: ${SCCACHE_SERVER}" + sh "chmod +w -R ${env.WORKSPACE}" + sh '''ping -c 1 -p 6379 "${SCCACHE_SERVER}" | echo $? > tmp.txt''' + def output = readFile(file: "tmp.txt") + echo "tmp.txt contents: \$output" + return (output != "0") + } + else{ + return 1 + } +} + +def check_arch_name(){ + sh 'rocminfo | tee rocminfo.log' + if ( runShell('grep -n "gfx90a" rocminfo.log') ){ + return "gfx90a" + } + else if ( runShell('grep -n "gfx942" rocminfo.log') ) { + return "gfx942" + } + else if ( runShell('grep -n "gfx101" rocminfo.log') ) { + return "gfx101" + } + else if ( runShell('grep -n "gfx103" rocminfo.log') ) { + return "gfx103" + } + else if ( runShell('grep -n "gfx11" rocminfo.log') ) { + return "gfx11" + } + else if ( runShell('grep -n "gfx120" rocminfo.log') ) { + return "gfx12" + } + else if ( runShell('grep -n "gfx908" rocminfo.log') ) { + return "gfx908" + } + else if ( runShell('grep -n "gfx950" rocminfo.log') ) { + return "gfx950" + } + else { + return "" + } +} + +def getDockerImage(Map conf=[:]){ + def image + if ( conf.get("docker_name", "") != "" ){ + image = conf.get("docker_name", "") + echo "Using special docker: ${image}" + } + else{ + image = getDockerImageName() + echo "Using default docker: ${image}" + } + //Check if image exists + def retimage + try + { + echo "Pulling image: ${image}" + retimage = docker.image("${image}") + withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { + retimage.pull() + } + } + catch(Exception ex) + { + error "Unable to locate image: ${image}" + } + return [retimage, image] +} + +// Build and push a docker image, capturing its digest into the specified env var. +// If forceBuild is false, will skip building if the image already exists in the registry. +def buildAndPushDockerImage(String install_prefix, String image_name, String dockerExtraArgs, boolean forceBuild){ + show_node_info() + env.DOCKER_BUILDKIT=1 + checkoutComposableKernel() + def dockerArgs = "--build-arg PREFIX=${install_prefix} --build-arg compiler_version='${params.COMPILER_VERSION}' --build-arg compiler_commit='${params.COMPILER_COMMIT}' --build-arg ROCMVERSION='${params.ROCMVERSION}' " + dockerArgs += " " + dockerExtraArgs + + if(!forceBuild){ + try{ + echo "Checking for image: ${image_name}" + withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { + sh "docker manifest inspect --insecure ${image_name}" + } + echo "Image: ${image_name} found! Skipping building image" + return image_name + } + catch(Exception ex){ + echo "Unable to locate image: ${image_name}. Will attempt to build image now." + } + } + + echo "Building image: ${image_name} with args: ${dockerArgs}" + def retimage = docker.build("${image_name}", dockerArgs) + withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { + retimage.push() + } + def digest = sh(returnStdout: true, script: "docker inspect --format='{{index .RepoDigests 0}}' ${image_name}").trim() + echo "Built image digest: ${digest}" + echo "Pruning dangling Docker images to free disk space on CI agent" + sh "docker image prune -f --filter 'dangling=true' || true" + return digest +} + +def buildDockerBase(install_prefix){ + def image_name = getDockerImageName() + def base_image_name = getBaseDockerImageName() + echo "Building Docker for ${image_name}" + def dockerExtraArgs = " -f projects/composablekernel/Dockerfile . " + if(params.COMPILER_VERSION == "develop" || params.COMPILER_VERSION == "amd-staging" || params.COMPILER_COMMIT != ""){ + dockerExtraArgs = " --no-cache --build-arg BASE_DOCKER='${base_image_name}' -f projects/composablekernel/Dockerfile.compiler . " + } + else if(params.COMPILER_VERSION == "therock"){ + dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile . " + } + env.CK_BASE_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, params.BUILD_DOCKER.toBoolean()) +} + +def buildDockerPytorch(install_prefix){ + def image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_pytorch" + def dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile.pytorch --build-arg CK_PYTORCH_BRANCH='${params.ck_pytorch_branch}' . " + env.CK_PYTORCH_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, true) +} + +def buildDockerAiter(install_prefix){ + def image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_aiter" + def dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile.aiter --build-arg AITER_BRANCH='${params.aiter_branch}' --build-arg CK_AITER_BRANCH='${params.ck_aiter_branch}' . " + env.CK_AITER_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, true) +} + +def buildDockerFa(install_prefix){ + def image_name = "${env.CK_DOCKERHUB_PRIVATE}:ck_fa" + def dockerExtraArgs = " --no-cache -f projects/composablekernel/Dockerfile.fa" + dockerExtraArgs += " --build-arg BASE_DOCKER='${params.fa_base_docker}'" + dockerExtraArgs += " --build-arg FA_BRANCH='${params.fa_branch}'" + dockerExtraArgs += " --build-arg CK_FA_BRANCH='${params.ck_fa_branch}'" + dockerExtraArgs += " --build-arg GPU_ARCHS='gfx942;gfx950'" + dockerExtraArgs += " . " + env.CK_FA_IMAGE = buildAndPushDockerImage(install_prefix, image_name, dockerExtraArgs, true) +} + +def buildDocker(install_prefix){ + buildDockerBase(install_prefix) + if (params.RUN_PYTORCH_TESTS.toBoolean()) { + buildDockerPytorch(install_prefix) + } + if (params.RUN_AITER_TESTS.toBoolean()) { + buildDockerAiter(install_prefix) + } + if (params.RUN_FA_TESTS.toBoolean()) { + buildDockerFa(install_prefix) + } +} + +def get_docker_options(){ + def dockerOpts + if ( params.BUILD_INSTANCES_ONLY ){ + dockerOpts = "--network=host --group-add video --group-add render --cap-add=SYS_PTRACE --security-opt seccomp=unconfined" + } + else{ //only add kfd and dri paths if you actually going to run somthing on GPUs + dockerOpts = "--network=host --device=/dev/kfd --device=/dev/dri --group-add video --group-add render --cap-add=SYS_PTRACE --security-opt seccomp=unconfined" + } + if (params.COMPILER_VERSION == "develop" || params.COMPILER_VERSION == "amd-staging" || params.COMPILER_VERSION == "therock" || params.COMPILER_COMMIT != ""){ + // the --env COMPRESSED_BUNDLE_FORMAT_VERSION=2 env variable is required when building code with offload-compress flag with + // newer clang22 compilers and running with older hip runtima libraries + dockerOpts = dockerOpts + " --env HIP_CLANG_PATH='/llvm-project/build/bin' --env COMPRESSED_BUNDLE_FORMAT_VERSION=2 --env HIP_PLATFORM=amd " + } + // on some machines the group ids for video and render groups may not be the same as in the docker image! + def video_id = sh(returnStdout: true, script: 'getent group video | cut -d: -f3') + def render_id = sh(returnStdout: true, script: 'getent group render | cut -d: -f3') + dockerOpts = dockerOpts + " --group-add=${video_id} --group-add=${render_id} -v /var/jenkins/ref-repo/:/var/jenkins/ref-repo/ " + echo "Docker flags: ${dockerOpts}" + return dockerOpts +} + +def build_client_examples(String arch){ + def cmd = """ cd ../client_example && rm -rf build && mkdir build && cd build && \ + cmake -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/projects/composablekernel/install;/opt/rocm" \ + -DGPU_TARGETS="${arch}" \ + -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ + -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" \ + -DCMAKE_CXX_FLAGS=" -O3 " .. && make -j """ + return cmd +} + +def build_client_examples_and_codegen_tests(String arch){ + def cmd = """ cd ../codegen && rm -rf build && mkdir build && cd build && \ + cmake -DCMAKE_PREFIX_PATH=/opt/rocm -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" .. && \ + make -j64 check && \ + cd ../../client_example && rm -rf build && mkdir build && cd build && \ + cmake -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/projects/composablekernel/install;/opt/rocm" \ + -DGPU_TARGETS="${arch}" \ + -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ + -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" \ + -DCMAKE_CXX_FLAGS=" -O3 " .. && make -j """ + return cmd +} + +def build_and_run_fmha(String arch){ + def cmd = """ cmake -G Ninja -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/projects/composablekernel/install;/opt/rocm" \ + -DGPU_TARGETS="${arch}" \ + -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ + -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" .. && \ + ninja -j128 tile_example_fmha_fwd tile_example_fmha_bwd && \ + cd ../ && + example/ck_tile/01_fmha/script/run_full_test.sh "CI_${params.COMPILER_VERSION}" "${env.BRANCH_NAME}" "${NODE_NAME}" "${arch}" """ + return cmd +} + +def cmake_build(Map conf=[:]){ + + def config_targets = conf.get("config_targets","check") + def build_envs = "CTEST_PARALLEL_LEVEL=4 " + conf.get("build_env","") + def prefixpath = conf.get("prefixpath","/opt/rocm") + def setup_args = conf.get("setup_args","") + // make sure all unit tests always run on develop branch + def runAllUnitTests = (env.BRANCH_NAME == "develop") ? true : params.RUN_ALL_UNIT_TESTS + + if (prefixpath != "/usr/local"){ + setup_args = setup_args + " -DCMAKE_PREFIX_PATH=${prefixpath} " + } + + //cmake_env can overwrite default CXX variables. + def cmake_envs + if(!setup_args.contains("gfx1250")){ + cmake_envs = "CXX=${params.BUILD_COMPILER} CXXFLAGS='-Werror' " + conf.get("cmake_ex_env","") + } + else{ //use default compiler for gfx1250 + cmake_envs = "CXX=/opt/rocm/llvm/bin/clang++ CXXFLAGS='-Werror' " + conf.get("cmake_ex_env","") + } + + if(conf.get("build_install","") == "true") + { + config_targets = 'install ' + config_targets + setup_args = ' -DBUILD_DEV=On -DCMAKE_INSTALL_PREFIX=../install' + setup_args + } else{ + setup_args = ' -DBUILD_DEV=On' + setup_args + } + if (params.DISABLE_DL_KERNELS){ + setup_args = setup_args + " -DDISABLE_DL_KERNELS=ON " + } + + setup_args = " -DCMAKE_BUILD_TYPE=release " + setup_args + + def pre_setup_cmd = """ + #!/bin/bash + cd projects/composablekernel + ulimit -c unlimited + rm -rf build + mkdir build + rm -rf install + mkdir install + cd build + """ + def invocation_tag="" + if (setup_args.contains("gfx12")){ + invocation_tag="gfx12" + } + if (setup_args.contains("gfx11")){ + invocation_tag="gfx11" + } + if (setup_args.contains("gfx101")){ + invocation_tag="gfx101" + } + if (setup_args.contains("gfx103")){ + invocation_tag="gfx103" + } + if (setup_args.contains("gfx908")){ + invocation_tag="gfx908" + } + if (setup_args.contains("gfx90a")){ + invocation_tag="gfx90a" + } + if (setup_args.contains("gfx94")){ + invocation_tag="gfx94" + } + if (setup_args.contains("gfx95")){ + invocation_tag="gfx95" + } + echo "invocation tag: ${invocation_tag}" + def redis_pre_setup_cmd = pre_setup_cmd + if(check_host() && params.USE_SCCACHE && "${env.CK_SCCACHE}" != "null" && "${invocation_tag}" != "") { + redis_pre_setup_cmd = pre_setup_cmd + """ + #!/bin/bash + export ROCM_PATH=/opt/rocm + export SCCACHE_ENABLED=true + export SCCACHE_LOG_LEVEL=debug + export SCCACHE_IDLE_TIMEOUT=14400 + export COMPILERS_HASH_DIR=/tmp/.sccache + export SCCACHE_BIN=/usr/local/.cargo/bin/sccache + export SCCACHE_EXTRAFILES=/tmp/.sccache/rocm_compilers_hash_file + export SCCACHE_REDIS="redis://${env.CK_SCCACHE}" + echo "connect = ${env.CK_SCCACHE}" >> ../script/redis-cli.conf + export SCCACHE_C_CUSTOM_CACHE_BUSTER="${invocation_tag}" + echo \$SCCACHE_C_CUSTOM_CACHE_BUSTER + stunnel ../script/redis-cli.conf + ../script/sccache_wrapper.sh --enforce_redis + """ + try { + def cmd1 = conf.get("cmd1", """ + ${redis_pre_setup_cmd} + """) + sh cmd1 + setup_args = " -DCMAKE_HIP_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache -DCMAKE_C_COMPILER_LAUNCHER=sccache " + setup_args + } + catch(Exception err){ + echo "could not connect to redis server: ${err.getMessage()}. will not use sccache." + def cmd2 = conf.get("cmd2", """ + ${pre_setup_cmd} + """) + sh cmd2 + } + } + else{ + def cmd3 = conf.get("cmd3", """ + ${pre_setup_cmd} + """) + sh cmd3 + } + + // reduce parallelism when compiling, clang uses too much memory + def nt = nthreads() + def cmd + def setup_cmd + def build_cmd + def execute_cmd = conf.get("execute_cmd", "") + //check the node gpu architecture + def arch_name = check_arch_name() + if(!setup_args.contains("NO_CK_BUILD")){ + if (params.NINJA_BUILD_TRACE) { + echo "running ninja build trace" + } + if (params.RUN_BUILDER_TESTS && !setup_args.contains("-DCK_CXX_STANDARD=") && !setup_args.contains("gfx10") && !setup_args.contains("gfx11")) { + setup_args = " -D CK_EXPERIMENTAL_BUILDER=ON " + setup_args + } + if (params.RUN_ROCM_CK_TESTS) { + setup_args = " -D CK_ENABLE_ROCM_CK=ON " + setup_args + } + setup_cmd = conf.get( + "setup_cmd", + """${cmake_envs} cmake -G Ninja ${setup_args} -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DCMAKE_CXX_FLAGS=" -O3 " .. """ + ) + + // Smart-build: Only build if running all tests or forced + // Otherwise, smart-build will determine what to build after cmake configure + if (runAllUnitTests) { + build_cmd = conf.get( + "build_cmd", + "${build_envs} ninja -j${nt} ${config_targets}" + ) + } else { + // Smart-build enabled: skip full build and execute_cmd (client examples) + build_cmd = "" + execute_cmd = "" + } + + cmd = conf.get("cmd", """ + ${setup_cmd} + ${build_cmd} + ${execute_cmd} + """) + } + else{ + cmd = conf.get("cmd", """ + ${execute_cmd} + """) + } + + echo cmd + + dir("projects/composablekernel/build"){ + // Start sccache monitoring + if(check_host() && params.USE_SCCACHE && "${env.CK_SCCACHE}" != "null" && "${invocation_tag}" != "") { + sh """ + chmod +x ../script/monitor_sccache_during_build.sh + mkdir -p logs + export SCCACHE_C_CUSTOM_CACHE_BUSTER="${invocation_tag}" + ../script/monitor_sccache_during_build.sh build_monitor & + MONITOR_PID=\$! + echo "Monitor PID: \$MONITOR_PID" + echo \$MONITOR_PID > monitor.pid + """ + } + try { + //build CK + sh cmd + if (runAllUnitTests){ + // Archive artifacts if they were generated + if (fileExists("ck_build_trace_${arch_name}.json")) { + archiveArtifacts "ck_build_trace_${arch_name}.json" + } + if (fileExists("clang_build_analysis_${arch_name}.log")) { + archiveArtifacts "clang_build_analysis_${arch_name}.log" + } + // Process ninja build trace after full build + if(fileExists(".ninja_log")) { + sh "python3 ../script/ninja_json_converter.py .ninja_log --legacy-format --output ck_build_trace_${arch_name}.json" + archiveArtifacts "ck_build_trace_${arch_name}.json" + sh "python3 ../script/parse_ninja_trace.py ck_build_trace_${arch_name}.json" + } + + if (params.NINJA_FTIME_TRACE) { + echo "running ClangBuildAnalyzer" + sh "/ClangBuildAnalyzer/build/ClangBuildAnalyzer --all . clang_build.log" + sh "/ClangBuildAnalyzer/build/ClangBuildAnalyzer --analyze clang_build.log > clang_build_analysis_${arch_name}.log" + archiveArtifacts "clang_build_analysis_${arch_name}.log" + } + } + } catch (Exception buildError) { + echo "Build failed: ${buildError.getMessage()}" + throw buildError + } finally { + // Stop sccache monitoring + if(check_host() && params.USE_SCCACHE && "${env.CK_SCCACHE}" != "null" && "${invocation_tag}" != "") { + sh """ + # Stop monitoring + if [ -f monitor.pid ]; then + MONITOR_PID=\$(cat monitor.pid) + kill \$MONITOR_PID 2>/dev/null || echo "Monitor already stopped" + rm -f monitor.pid + fi + """ + + // Archive the monitoring logs + try { + archiveArtifacts artifacts: "logs/*monitor*.log", allowEmptyArchive: true + } catch (Exception e) { + echo "Could not archive sccache monitoring logs: ${e.getMessage()}" + } + } + } + + //run tests except when NO_CK_BUILD is set and except on gfx1250 + if(!setup_args.contains("NO_CK_BUILD")){ + // run unit tests unless building library for all targets + // Note: This else block is when NINJA_BUILD_TRACE=false and BUILD_INSTANCES_ONLY=false + // So no ninja trace processing needed here + if (!params.BUILD_INSTANCES_ONLY){ + if (!runAllUnitTests && !setup_args.contains("gfx1250") ){ + // Smart Build: Run smart_build_and_test.sh + sh """ + export WORKSPACE_ROOT=${env.WORKSPACE} + export PARALLEL=32 + export NINJA_JOBS=${nt} + export ARCH_NAME=${arch_name} + export PROCESS_NINJA_TRACE=false + export NINJA_FTIME_TRACE=false + bash ../script/dependency-parser/smart_build_and_test.sh + """ + } + else{ //run all tests + if(!setup_args.contains("gfx1250")){ + echo "Full test suite requested (RUN_ALL_UNIT_TESTS=true or develop branch)" + sh "ninja -j${nt} check" + } + else{ //do not run tests on gfx1250, just build everything + echo "Building for gfx1250" + sh "ninja -j${nt}" + } + if (params.RUN_ROCM_CK_TESTS) { + sh 'ninja check-rocm-ck' + } + if(params.BUILD_PACKAGES || params.BUILD_INSTANCES_ONLY){ + echo "Build ckProfiler packages" + sh 'ninja -j64 package' + sh "mv composablekernel-ckprofiler_*.deb composablekernel-ckprofiler_1.2.0_amd64_${arch_name}.deb" + stash includes: "composablekernel-ckprofiler**.deb", name: "profiler_package_${arch_name}" + } + } + if (params.RUN_BUILDER_TESTS && !setup_args.contains("-DCK_CXX_STANDARD=") && !setup_args.contains("gfx10") && !setup_args.contains("gfx11")) { + sh 'ninja check-builder' + } + } + } + } + + if (params.RUN_CK_TILE_FMHA_TESTS){ + try{ + dir("projects/composablekernel"){ + archiveArtifacts "perf_fmha_*.log" + stash includes: "perf_fmha_**.log", name: "perf_fmha_log_${arch_name}" + } + } + catch(Exception err){ + echo "could not locate the requested artifacts: ${err.getMessage()}. will skip the stashing." + } + } +} + +def buildHipClangJob(Map conf=[:]){ + show_node_info() + checkoutComposableKernel() + def prefixpath = conf.get("prefixpath", "/opt/rocm") + def dockerOpts = get_docker_options() + def image + def retimage + (retimage, image) = getDockerImage(conf) + + setGithubStatus("${env.STAGE_NAME}", 'pending', "Starting ${env.STAGE_NAME}") + try { + withDockerContainer(image: image, args: dockerOpts) { + timeout(time: 20, unit: 'HOURS') + { + cmake_build(conf) + } + } + setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") + } + catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ + setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") + throw e + } + return retimage +} + +def buildHipClangJobAndReboot(Map conf=[:]){ + try{ + buildHipClangJob(conf) + } + catch(e){ + echo "throwing error exception for the stage" + echo 'Exception occurred: ' + e.toString() + throw e + } +} + +def Build_CK(Map conf=[:]){ + show_node_info() + checkoutComposableKernel() + def prefixpath = conf.get("prefixpath", "/opt/rocm") + def dockerOpts=get_docker_options() + def image + def retimage + + setGithubStatus("${env.STAGE_NAME}", 'pending', "Starting ${env.STAGE_NAME}") + try { + try { + (retimage, image) = getDockerImage(conf) + withDockerContainer(image: image, args: dockerOpts) { + timeout(time: 2, unit: 'MINUTES'){ + sh 'rocminfo | tee rocminfo.log' + if ( !runShell('grep -n "gfx" rocminfo.log') ){ + throw new Exception ("GPU not found") + } + else{ + echo "GPU is OK" + } + } + } + } + catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ + echo "The job was cancelled or aborted" + setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") + throw e + } + withDockerContainer(image: image, args: dockerOpts) { + timeout(time: 20, unit: 'HOURS') + { + //check whether to run performance tests on this node + def arch = check_arch_name() + cmake_build(conf) + if ( params.RUN_INDUCTOR_TESTS && arch == "gfx90a" ){ + echo "Run inductor codegen tests" + sh "projects/composablekernel/script/run_inductor_tests.sh" + } + // run performance tests, stash the logs, results will be processed on the master node + dir("projects/composablekernel/script"){ + if (params.RUN_PERFORMANCE_TESTS){ + if (params.RUN_FULL_QA && (arch == "gfx90a" || arch == "gfx942")){ + // run full tests on gfx90a or gfx942 + echo "Run full performance tests" + sh "./run_full_performance_tests.sh 0 QA_${params.COMPILER_VERSION} ${env.BRANCH_NAME} ${NODE_NAME} ${arch}" + archiveArtifacts "perf_*.log" + stash includes: "perf_**.log", name: "perf_log_${arch}" + } + else if (!params.RUN_FULL_QA && (arch == "gfx90a" || arch == "gfx942")){ + // run standard tests on gfx90a or gfx942 + echo "Run performance tests" + sh "./run_performance_tests.sh 0 CI_${params.COMPILER_VERSION} ${env.BRANCH_NAME} ${NODE_NAME} ${arch}" + archiveArtifacts "perf_*.log" + stash includes: "perf_**.log", name: "perf_log_${arch}" + } + else if ( arch != "gfx10"){ + // run basic tests on gfx11/gfx12/gfx908/gfx950, but not on gfx10, it takes too long + echo "Run gemm performance tests" + sh "./run_gemm_performance_tests.sh 0 CI_${params.COMPILER_VERSION} ${env.BRANCH_NAME} ${NODE_NAME} ${arch}" + archiveArtifacts "perf_onnx_gemm_*.log" + stash includes: "perf_onnx_gemm_**.log", name: "perf_log_${arch}" + } + } + } + if (params.hipTensor_test && arch == "gfx90a" ){ + // build and test hipTensor on gfx90a node + sh """#!/bin/bash + rm -rf rocm-libraries + git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-libraries.git + cd rocm-libraries + git sparse-checkout init --cone + git sparse-checkout set projects/hiptensor + git checkout "${params.hipTensor_branch}" + """ + dir("rocm-libraries/projects/hiptensor"){ + sh """#!/bin/bash + mkdir -p build + ls -ltr + CC=hipcc CXX=hipcc cmake -Bbuild . -D CMAKE_PREFIX_PATH="${env.WORKSPACE}/install" + cmake --build build -- -j + ctest --test-dir build + """ + } + } + } + } + setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") + } + catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ + setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") + throw e + } + return retimage +} + +def Build_CK_and_Reboot(Map conf=[:]){ + try{ + Build_CK(conf) + } + catch(e){ + echo "throwing error exception while building CK" + echo 'Exception occurred: ' + e.toString() + throw e + } +} + +def process_results(Map conf=[:]){ + checkoutComposableKernel() + //use older image that has user jenkins + def image = "${env.CK_DOCKERHUB}:ck_ub22.04_rocm6.3" + + setGithubStatus("${env.STAGE_NAME}", 'pending', 'Processing results...') + try { + try + { + echo "Pulling image: ${image}" + def retimage = docker.image("${image}") + withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { + retimage.pull() + } + } + catch(Exception ex) + { + error "Unable to locate image: ${image}" + } + } + catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ + setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") + throw e + } + + withDockerContainer(image: image, args: '--cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v=/var/jenkins/:/var/jenkins') { + timeout(time: 15, unit: 'MINUTES'){ + try{ + dir("projects/composablekernel/script"){ + if (params.RUN_CK_TILE_FMHA_TESTS){ + try{ + unstash "perf_fmha_log_gfx942" + } + catch(Exception err){ + echo "could not locate the FMHA performance logs for gfx942: ${err.getMessage()}." + } + try{ + unstash "perf_fmha_log_gfx90a" + } + catch(Exception err){ + echo "could not locate the FMHA performance logs for gfx90a: ${err.getMessage()}." + } + try{ + unstash "perf_fmha_log_gfx950" + } + catch(Exception err){ + echo "could not locate the FMHA performance logs for gfx950: ${err.getMessage()}." + } + + } + if (params.BUILD_INSTANCES_ONLY){ + // unstash deb packages + try{ + unstash "lib_package" + sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" + } + catch(Exception err){ + echo "could not locate lib_package." + } + } + if (params.BUILD_PACKAGES){ + // unstash deb packages + try{ + unstash "profiler_package_gfx90a" + sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" + } + catch(Exception err){ + echo "could not locate profiler_package_gfx90a." + } + try{ + unstash "profiler_package_gfx942" + sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" + } + catch(Exception err){ + echo "could not locate profiler_package_gfx942." + } + try{ + unstash "profiler_package_gfx950" + sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" + } + catch(Exception err){ + echo "could not locate profiler_package_gfx950." + } + try{ + unstash "profiler_package_gfx12" + sh "sshpass -p ${env.ck_deb_pw} scp -o StrictHostKeyChecking=no composablekernel-ckprofiler*.deb ${env.ck_deb_user}@${env.ck_deb_ip}:/var/www/html/composable_kernel/" + } + catch(Exception err){ + echo "could not locate profiler_package_gfx12." + } + } + else{ + // unstash perf files to master + try{ + unstash "perf_log_gfx90a" + } + catch(Exception err){ + echo "could not locate the gfx90a performance logs: ${err.getMessage()}." + } + try{ + unstash "perf_log_gfx942" + } + catch(Exception err){ + echo "could not locate the gfx942 performance logs: ${err.getMessage()}." + } + try{ + unstash "perf_log_gfx950" + } + catch(Exception err){ + echo "could not locate the gfx950 performance logs: ${err.getMessage()}." + } + try{ + unstash "perf_log_gfx908" + } + catch(Exception err){ + echo "could not locate the gfx908 performance logs: ${err.getMessage()}." + } + try{ + unstash "perf_log_gfx11" + } + catch(Exception err){ + echo "could not locate the gfx11 performance logs: ${err.getMessage()}." + } + try{ + + unstash "perf_log_gfx12" + } + catch(Exception err){ + echo "could not locate the gfx12 performance logs: ${err.getMessage()}." + } + } + // process the logs + sh "./process_perf_data.sh" + } + setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") + } + catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ + setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") + throw e + } + finally{ + echo "Finished processing performance test results" + } + } + } +} + +def run_downstream_tests(Map conf=[:]){ + show_node_info() + checkoutComposableKernel() + def dockerOpts = get_docker_options() + ' --group-add irc ' + + setGithubStatus("${env.STAGE_NAME}", 'pending', "Starting ${env.STAGE_NAME}") + try { + try + { + echo "Pulling image: ${conf.image}" + retimage = docker.image("${conf.image}") + withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) { + retimage.pull() + } + } + catch(Exception ex) + { + error "Unable to locate image: ${conf.image}" + } + } + catch (org.jenkinsci.plugins.workflow.steps.FlowInterruptedException e){ + setGithubStatus("${env.STAGE_NAME}", 'failure', "Stage ${env.STAGE_NAME} failed") + throw e + } + + withDockerContainer(image: conf.image, args: dockerOpts) { + timeout(time: conf.get("timeoutHours", 2), unit: 'HOURS'){ + try{ + sh "rocminfo" + sh "python3 --version" + for (cmd in conf.execute_cmds) { + sh "${cmd}" + } + setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") + } + catch(e){ + echo "Throwing error exception while running ${env.STAGE_NAME}" + echo 'Exception occurred: ' + e.toString() + setGithubStatus("${env.STAGE_NAME}", 'error', "Stage ${env.STAGE_NAME} failed") + throw e + } + finally{ + echo "Finished running ${env.STAGE_NAME}" + } + } + } +} + +def getPytorchTestsCmds() { + return [ + "mkdir pytorch", + "cp -r /var/jenkins/workspace/pytorch/* pytorch/", + "ls -ltr pytorch", + "python3 pytorch/tools/amd_build/build_amd.py", + "cd pytorch && USE_ROCM_CK_SDPA=1 PYTORCH_ROCM_ARCH=gfx942 python3 setup.py develop" + ] +} +def getAiterTestsCmds() { + return [ + "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8_blockscale.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_mha.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_mha_varlen.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_batch_prefill.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_moe.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_2stage.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_blockscale.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_ep.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_sorting.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_sorting_mxfp4.py", + "python3 /home/jenkins/workspace/aiter/op_tests/test_moe_tkw1.py" + ] +} +def getFaTestsCmds() { + return [ + "python3 -u -m pytest /home/jenkins/workspace/flash-attention/tests/test_flash_attn_ck.py" + ] +} + +def runClangFormat() { + buildHipClangJobAndReboot( + setup_args: "NO_CK_BUILD", + setup_cmd: "", + build_cmd: "", + execute_cmd: """cd .. && \ + find . -type f \\( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.cl' \\) \ + -not -path '*/build/*' -not -path '*/include/rapidjson/*' | \ + xargs -P 8 -I{} sh -c 'clang-format-18 -style=file {} | diff -u - {} || (echo "ERROR: {} needs formatting" && exit 1)'""" + ) +} + +def runClangFormatAndCppcheck() { + buildHipClangJobAndReboot( + setup_args: "NO_CK_BUILD", + setup_cmd: "", + build_cmd: "", + execute_cmd: """cd .. && \ + find . -type f \\( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.cl' \\) \ + -not -path '*/build/*' -not -path '*/include/rapidjson/*' | \ + xargs -P 8 -I{} sh -c 'clang-format-18 -style=file {} | diff -u - {} || (echo "ERROR: {} needs formatting" && exit 1)' && \ + /cppcheck/build/bin/cppcheck ../* -v -j \$(nproc) -I ../include -I ../profiler/include -I ../library/include \ + -D CK_ENABLE_FP64 -D CK_ENABLE_FP32 -D CK_ENABLE_FP16 -D CK_ENABLE_FP8 -D CK_ENABLE_BF16 -D CK_ENABLE_BF8 -D CK_ENABLE_INT8 \ + -D __gfx908__ -D __gfx90a__ -D __gfx942__ -D __gfx1030__ -D __gfx1100__ -D __gfx1101__ -D __gfx1102__ \ + -U __gfx803__ -U __gfx900__ -U __gfx906__ -U CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4 \ + --file-filter=*.cpp --force --enable=all --output-file=ck_cppcheck.log""" + ) +} + +def runFullGroupedConvTileTests() { + buildHipClangJobAndReboot( + setup_args: "NO_CK_BUILD", + build_type: 'Release', + execute_cmd: """ + python3 ../experimental/grouped_convolution_tile_instances/generate_instances.py --mode=profiler && \ + cmake .. --preset dev-gfx90a -D CK_EXPERIMENTAL_BUILDER=ON && \ + make -j64 test_grouped_convnd_fwd_tile test_grouped_convnd_bwd_weight_tile && \ + ./bin/test_grouped_convnd_bwd_weight_tile && \ + ./bin/test_grouped_convnd_fwd_tile""" + ) +} + +def runGroupedConvLargeCaseTests() { + buildHipClangJobAndReboot( + setup_args: "NO_CK_BUILD", + build_type: 'Release', + execute_cmd: """ + cmake .. --preset dev-gfx90a && \ + make -j64 test_grouped_convnd_fwd_large_cases test_grouped_convnd_bwd_data_large_cases test_grouped_convnd_fwd_bias_clamp_large_cases && \ + ./bin/test_grouped_convnd_fwd_large_cases && \ + ./bin/test_grouped_convnd_bwd_data_large_cases && \ + ./bin/test_grouped_convnd_fwd_bias_clamp_large_cases""" + ) +} + +def runComprehensiveConvDatasetTests() { + buildHipClangJobAndReboot( + setup_args: "NO_CK_BUILD", + build_type: 'Release', + execute_cmd: """ + cd ../build && \ + cmake .. --preset dev-gfx90a && \ + make -j64 test_grouped_convnd_fwd_dataset_xdl \ + test_grouped_convnd_bwd_data_dataset_xdl \ + test_grouped_convnd_bwd_weight_dataset_xdl && \ + cd ../test_data && \ + ./generate_test_dataset.sh small && \ + cd ../build && \ + ./bin/test_grouped_convnd_fwd_dataset_xdl && \ + ./bin/test_grouped_convnd_bwd_data_dataset_xdl && \ + ./bin/test_grouped_convnd_bwd_weight_dataset_xdl""" + ) +} + +def runTileEngineBasicTests(String compiler) { + buildHipClangJobAndReboot( + setup_args: "NO_CK_BUILD", + build_type: 'Release', + execute_cmd: """ + cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ + -D BUILD_CK_TILE_ENGINE="ON" \ + -D CMAKE_CXX_COMPILER="${compiler}" \ + -D CMAKE_BUILD_TYPE=Release \ + -D GPU_TARGETS="gfx942" \ + -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16" \ + -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ + -D GEMM_UNIVERSAL_CONFIG_FILE="default_ci_config.json" \ + -D GEMM_MULTI_D_DATATYPE="fp16" \ + -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ + -D GEMM_MULTI_D_CONFIG_FILE="default_ci_config.json" \ + -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ + -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ + -D GEMM_PRESHUFFLE_CONFIG_FILE="default_ci_config.json" .. && \ + ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all && \ + python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ + python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ + python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json""" + ) +} + +def runTileEngineGemmTests(String arch, String compiler) { + def execute_cmd + if (arch == "gfx942") { + execute_cmd = """ + cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ + -D BUILD_CK_TILE_ENGINE="ON" \ + -D CMAKE_CXX_COMPILER="${compiler}" \ + -D CMAKE_BUILD_TYPE=Release \ + -D GPU_TARGETS="gfx942" \ + -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16;bf8;bf16" \ + -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ + -D GEMM_STREAMK_DATATYPE="fp8;fp16" \ + -D GEMM_STREAMK_LAYOUT="rcr" \ + -D GEMM_MULTI_D_DATATYPE="fp16" \ + -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ + -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ + -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ + -D GROUPED_GEMM_DATATYPE="fp8;fp16" \ + -D GROUPED_GEMM_LAYOUT="rcr;rrr;crr;ccr" \ + -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ + ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all benchmark_gemm_streamk_all benchmark_grouped_gemm_all && \ + python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_universal_results.json && \ + python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ + python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ + python3 ../tile_engine/ops/gemm/grouped_gemm/grouped_gemm_benchmark.py . --problem-sizes "1024,1024,1024" --group-counts 8 --warmup 5 --repeat 5 --verbose --json grouped_gemm_results.json""" + } else if (arch == "gfx950") { + execute_cmd = """ + cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ + -D BUILD_CK_TILE_ENGINE="ON" \ + -D CMAKE_CXX_COMPILER="${compiler}" \ + -D CMAKE_BUILD_TYPE=Release \ + -D GPU_TARGETS="gfx950" \ + -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16" \ + -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ + -D GEMM_MULTI_D_DATATYPE="fp16" \ + -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ + -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ + -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ + -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ + ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all && \ + python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ + python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json && \ + python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json""" + } else if (arch == "gfx1201") { + execute_cmd = """ + cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ + -D BUILD_CK_TILE_ENGINE="ON" \ + -D CMAKE_CXX_COMPILER="${compiler}" \ + -D CMAKE_BUILD_TYPE=Release \ + -D GPU_TARGETS="gfx1201" \ + -D GEMM_UNIVERSAL_DATATYPE="fp16" \ + -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ + -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ + ninja -j${nthreads()} benchmark_gemm_universal_all && \ + python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json results.json""" + } + buildHipClangJobAndReboot(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: execute_cmd) +} + +def runBuildCKAndTests(String arch) { + def gpuTarget + def extraSetupArgs = "" + def execute_cmd = "" + def extraBuildArgs = [:] + + switch (arch) { + case "gfx90a": + gpuTarget = "gfx90a" + extraSetupArgs = " -DCK_CXX_STANDARD=\"17\"" + execute_cmd = build_client_examples_and_codegen_tests(gpuTarget) + break + case "gfx1250": + gpuTarget = "gfx1250" + extraSetupArgs = " -DDISABLE_DL_KERNELS=\"ON\"" + extraBuildArgs = [docker_name: "${env.CK_DOCKERHUB_PRIVATE}:npi-mi450-latest", no_reboot: true] + break + case "gfx10-1-generic": + case "gfx10-3-generic": + case "gfx11-generic": + case "gfx12-generic": + gpuTarget = arch + execute_cmd = build_client_examples(gpuTarget) + break + default: + gpuTarget = arch + execute_cmd = build_client_examples(gpuTarget) + } + + def setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="${gpuTarget}"${extraSetupArgs} """ + def buildArgs = [setup_args: setup_args, config_targets: "install", build_type: 'Release', prefixpath: '/usr/local'] + if (execute_cmd) { + buildArgs.execute_cmd = execute_cmd + } + buildArgs.putAll(extraBuildArgs) + Build_CK_and_Reboot(buildArgs) +} + +def runBuildInstancesOnly(String compiler) { + buildHipClangJobAndReboot( + setup_args: "NO_CK_BUILD", + build_cmd: "", + build_type: 'Release', + execute_cmd: """ + cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ + -DCMAKE_CXX_COMPILER="${compiler}" \ + -DCMAKE_HIP_COMPILER="${compiler}" \ + -DGPU_ARCHS="gfx908;gfx90a;gfx942;gfx950;gfx10-3-generic;gfx11-generic;gfx12-generic" \ + -D CMAKE_BUILD_TYPE=Release .. && ninja -j64""" + ) +} + +return this From b7c8fb164f9c61a5e1ad844b0d810e484d529505 Mon Sep 17 00:00:00 2001 From: Johannes Graner <67631091+johannes-graner@users.noreply.github.com> Date: Tue, 2 Jun 2026 08:51:17 +0000 Subject: [PATCH 010/143] [rocm-libraries] ROCm/rocm-libraries#7937 (commit abe276d) [CK Tile] Add conv Wavelet GEMM pipeline and bwd_weight instances (#7937) ## Motivation CK Tile had no pipeline competitive with old CK's wavelet on the RetinaNet K=36 C=256 3x3 conv bwd_weight class. This adds a wave-specialized "wavelet" GEMM pipeline so CK Tile has a competitive kernel for spatial small-K shapes. ## Technical Details - New wavelet GEMM pipeline (`gemm_pipeline_ag_bg_cr_wavelet.hpp`): workgroup split into math waves (LDS read + MFMA) and load waves (DRAM read + LDS write). - VGPR role-split: `operator()` has two top-level mutually-exclusive `is_math` branches so the allocator overlays both roles onto the same physical VGPRs, cutting arch VGPR ~33-40% and raising occupancy. Correctness depends on identical `block_sync_lds` counts on both arms plus a matching load-wave barrier stub in the epilogue (`cshuffle_epilogue.hpp`). - Kernel dispatch (`grouped_convolution_backward_weight_kernel.hpp`): `kIsWavelet` path, `LaunchBlockSize`, load-wave barrier stub. Uplift: wavelet is the fastest CK Tile pipeline on the RetinaNet K=36 C=256 3x3 family, beating the best non-wavelet CK Tile kernel by 10-27% (googlenet K=320 by 16-23%); the role-split roughly halves the parity gap vs old CK on the 13x13 fp16 shape. ## Test Plan - `ckProfiler grouped_conv_bwd_weight`, NHWGC layout, fp16/bf16, `split_k=all`, CPU verify on RetinaNet K=36 shapes (7x7, 13x13) and a broad 2D sweep. - Correctness: `-v=1` across `split_k` in {-1,1,2,4,8,16,32,64} (barrier-parity / deadlock check). - `test_grouped_convnd_bwd_weight` over the tests `.conf` wavelet instances. ## Test Result - All wavelet instances CPU-verify correct across the split-K sweep; no hangs (dual-arm barrier sequence matches). - Wavelet wins the RetinaNet K=36 C=256 3x3 family (10-27% over best non-wavelet CK Tile) and googlenet K=320 (16-23%); at parity-or-better vs old CK on the majority of spatial shapes. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- .../ck_tile/conv_tile_tuning_params.hpp | 10 + .../builder/include/ck_tile/builder/types.hpp | 4 +- .../backward_weight/profiler/nhwgc_bf16.conf | 3 + .../backward_weight/profiler/nhwgc_fp16.conf | 3 + .../backward_weight/tests/nhwgc_bf16.conf | 2 + .../backward_weight/tests/nhwgc_fp16.conf | 2 + .../generate_instances.py | 244 ++++++++---- include/ck_tile/core.hpp | 1 + include/ck_tile/core/arch/inst_prefetch.hpp | 182 ++++----- .../ops/epilogue/cshuffle_epilogue.hpp | 30 ++ include/ck_tile/ops/gemm.hpp | 1 + .../gemm_pipeline_ag_bg_cr_wavelet.hpp | 361 ++++++++++++++++++ .../ops/gemm/pipeline/gemm_pipelines.hpp | 3 +- ...ped_convolution_backward_weight_kernel.hpp | 82 +++- 14 files changed, 749 insertions(+), 179 deletions(-) create mode 100644 include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wavelet.hpp diff --git a/experimental/builder/include/ck_tile/builder/factory/helpers/ck_tile/conv_tile_tuning_params.hpp b/experimental/builder/include/ck_tile/builder/factory/helpers/ck_tile/conv_tile_tuning_params.hpp index 1296d7a0f9..d25e01812f 100644 --- a/experimental/builder/include/ck_tile/builder/factory/helpers/ck_tile/conv_tile_tuning_params.hpp +++ b/experimental/builder/include/ck_tile/builder/factory/helpers/ck_tile/conv_tile_tuning_params.hpp @@ -4,6 +4,7 @@ #pragma once #include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wavelet.hpp" #include "ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner.hpp" #include "ck_tile/builder/conv_algorithm_concepts.hpp" #include "ck_tile/builder/types.hpp" @@ -120,6 +121,14 @@ struct TilePipelineType using GemmPipeline = ck_tile::GemmPipelineAGmemBGmemCRegAsyncV1; }; +template <> +struct TilePipelineType +{ + template + using GemmPipeline = ck_tile:: + GemmPipelineAgBgCrWavelet; +}; + template consteval ck_tile::GemmPipeline SetTileBlockGemmPipelineVersion() { @@ -135,6 +144,7 @@ consteval ck_tile::GemmPipeline SetTileBlockGemmPipelineVersion() case PipelineVersion::V6: return ck_tile_pipeline::COMPUTE_V6; case PipelineVersion::ASYNC_V1: return ck_tile_pipeline::BASIC_ASYNC_V1; case PipelineVersion::ASYNC_V4: return ck_tile_pipeline::COMPUTE_ASYNC; + case PipelineVersion::WAVELET: return ck_tile_pipeline::WAVELET; case PipelineVersion::WEIGHT_ONLY: throw "PipelineVersion::WEIGHT_ONLY is not supported for block GEMM pipeline version."; default: throw "Unknown block GEMM PipelineVersion"; diff --git a/experimental/builder/include/ck_tile/builder/types.hpp b/experimental/builder/include/ck_tile/builder/types.hpp index 07ccd5e016..75aff9160b 100644 --- a/experimental/builder/include/ck_tile/builder/types.hpp +++ b/experimental/builder/include/ck_tile/builder/types.hpp @@ -160,7 +160,8 @@ enum class PipelineVersion V6, ASYNC_V1, ASYNC_V4, - WEIGHT_ONLY + WEIGHT_ONLY, + WAVELET }; // Enums for the GEMM specialization. @@ -355,6 +356,7 @@ inline std::string_view to_string(PipelineVersion ver) case ASYNC_V1: return "ASYNC_V1"; case ASYNC_V4: return "ASYNC_V4"; case WEIGHT_ONLY: return "WEIGHT_ONLY"; + case WAVELET: return "WAVELET"; default: return "Unknown"; } } diff --git a/experimental/grouped_convolution_tile_instances/configs/backward_weight/profiler/nhwgc_bf16.conf b/experimental/grouped_convolution_tile_instances/configs/backward_weight/profiler/nhwgc_bf16.conf index f2c7392641..6280889ca1 100644 --- a/experimental/grouped_convolution_tile_instances/configs/backward_weight/profiler/nhwgc_bf16.conf +++ b/experimental/grouped_convolution_tile_instances/configs/backward_weight/profiler/nhwgc_bf16.conf @@ -238,3 +238,6 @@ DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 128, BlkTile: 16x128x64, WaveTile: 16x16, K1: 4x4, WaveMap: 1x4, VmemReadVec: 2x1xSeq(1), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 256, BlkTile: 16x256x64, WaveTile: 16x16, K1: 2x4, WaveMap: 1x4, VmemReadVec: 2x1xSeq(1), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 256, BlkTile: 16x256x64, WaveTile: 16x16, K1: 2x2, WaveMap: 1x4, VmemReadVec: 2x1xSeq(1), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,4,8,8,1,0,0,16,64,64,1,4,1,16,16,32,bf16,bf16,WAVELET,Intrawave,0,1,fp32,bf16,EmptyTuple,PassThrough,0> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,8,1,0,0,64,32,64,1,2,1,16,16,16,bf16,bf16,WAVELET,Intrawave,0,1,fp32,bf16,EmptyTuple,PassThrough,0> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,4,8,8,1,0,0,64,64,64,2,2,1,16,16,16,bf16,bf16,WAVELET,Intrawave,0,1,fp32,bf16,EmptyTuple,PassThrough,0> diff --git a/experimental/grouped_convolution_tile_instances/configs/backward_weight/profiler/nhwgc_fp16.conf b/experimental/grouped_convolution_tile_instances/configs/backward_weight/profiler/nhwgc_fp16.conf index 563755f4de..321cceaa43 100644 --- a/experimental/grouped_convolution_tile_instances/configs/backward_weight/profiler/nhwgc_fp16.conf +++ b/experimental/grouped_convolution_tile_instances/configs/backward_weight/profiler/nhwgc_fp16.conf @@ -241,3 +241,6 @@ DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 128, BlkTile: 16x128x64, WaveTile: 16x16, K1: 4x4, WaveMap: 1x4, VmemReadVec: 1x8xSeq(4), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 256, BlkTile: 16x256x64, WaveTile: 16x16, K1: 2x4, WaveMap: 1x4, VmemReadVec: 1x8xSeq(4), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 256, BlkTile: 16x256x64, WaveTile: 16x16, K1: 2x2, WaveMap: 1x4, VmemReadVec: 1x2xSeq(4), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,4,8,8,1,0,0,16,64,64,1,4,1,16,16,32,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough,0> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,8,1,0,0,64,32,64,1,2,1,16,16,16,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough,0> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,4,8,8,1,0,0,64,64,64,2,2,1,16,16,16,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough,0> diff --git a/experimental/grouped_convolution_tile_instances/configs/backward_weight/tests/nhwgc_bf16.conf b/experimental/grouped_convolution_tile_instances/configs/backward_weight/tests/nhwgc_bf16.conf index e6430e3430..11abd0b218 100644 --- a/experimental/grouped_convolution_tile_instances/configs/backward_weight/tests/nhwgc_bf16.conf +++ b/experimental/grouped_convolution_tile_instances/configs/backward_weight/tests/nhwgc_bf16.conf @@ -46,3 +46,5 @@ DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 256, BlkTile: 256x16x64, WaveTile: 16x16, K1: 8x2, WaveMap: 4x1, VmemReadVec: 8x1xSeq(1), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 64, BlkTile: 16x16x64, WaveTile: 16x16, K1: 4x4, WaveMap: 1x1, VmemReadVec: 4x1xSeq(1), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 256, BlkTile: 16x256x64, WaveTile: 16x16, K1: 2x2, WaveMap: 1x4, VmemReadVec: 2x1xSeq(1), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,4,8,8,1,0,0,16,64,64,1,4,1,16,16,32,bf16,bf16,WAVELET,Intrawave,0,1,fp32,bf16,EmptyTuple,PassThrough,0> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,8,1,0,0,64,32,64,1,2,1,16,16,16,bf16,bf16,WAVELET,Intrawave,0,1,fp32,bf16,EmptyTuple,PassThrough,0> diff --git a/experimental/grouped_convolution_tile_instances/configs/backward_weight/tests/nhwgc_fp16.conf b/experimental/grouped_convolution_tile_instances/configs/backward_weight/tests/nhwgc_fp16.conf index 9aff77c6e5..f3a3f359af 100644 --- a/experimental/grouped_convolution_tile_instances/configs/backward_weight/tests/nhwgc_fp16.conf +++ b/experimental/grouped_convolution_tile_instances/configs/backward_weight/tests/nhwgc_fp16.conf @@ -46,3 +46,5 @@ DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 64, BlkTile: 16x16x64, WaveTile: 16x16, K1: 4x4, WaveMap: 1x1, VmemReadVec: 1x4xSeq(4), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v1, BlkGemmPipelinePrefetchStages: 1> DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 128, BlkTile: 128x16x64, WaveTile: 16x16, K1: 8x4, WaveMap: 4x1, VmemReadVec: 1x2xSeq(2), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> DeviceGroupedConvBwdWeight_Explicit_Xdl BlkSize: 128, BlkTile: 16x64x64, WaveTile: 16x16, K1: 4x4, WaveMap: 1x2, VmemReadVec: 1x8xSeq(4), BlkGemmPipelineScheduler: Intrawave, BlkGemmPipelineVersion: v2, BlkGemmPipelinePrefetchStages: 2> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,4,8,8,1,0,0,16,64,64,1,4,1,16,16,32,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough,0> +GroupedConvolutionBackwardWeightKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,8,1,0,0,64,32,64,1,2,1,16,16,16,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough,0> diff --git a/experimental/grouped_convolution_tile_instances/generate_instances.py b/experimental/grouped_convolution_tile_instances/generate_instances.py index 7ef4b458a6..e8898a854e 100755 --- a/experimental/grouped_convolution_tile_instances/generate_instances.py +++ b/experimental/grouped_convolution_tile_instances/generate_instances.py @@ -12,11 +12,10 @@ if str(_DISPATCHER_CODEGEN) not in sys.path: sys.path.insert(0, str(_DISPATCHER_CODEGEN)) -from grouped_config_rules import ( +from grouped_config_rules import ( # noqa E402 check_vectors as _shared_check_vectors, check_warp_coverage, check_bwd_data_vec_coverage, - WARP_SIZE, ) @@ -142,7 +141,9 @@ def check_vectors(a_scalar_per_vector, b_scalar_per_vector, c_scalar_per_vector) Delegates to the shared rule in grouped_config_rules.py. """ - return _shared_check_vectors(a_scalar_per_vector, b_scalar_per_vector, c_scalar_per_vector) + return _shared_check_vectors( + a_scalar_per_vector, b_scalar_per_vector, c_scalar_per_vector + ) def parse_instance_string(instance_string): @@ -267,6 +268,7 @@ def generate_conv_cpp( "COMPUTE_V6": "V6", "BASIC_ASYNC_V1": "ASYNC_V1", "COMPUTE_ASYNC": "ASYNC_V4", + "WAVELET": "WAVELET", } # Maps ck_tile StreamKReductionStrategy int values (from static_cast in instance string) @@ -492,8 +494,14 @@ def parse_bwd_weight_instances(instances, problem_name): continue native = try_parse_native_instance(instance, instance_id, problem_name) if native is not None: - if native.streamk_enabled and get_dtype(problem_name) == "float" and native.pipeline_version.find("ASYNC") != -1: - print(f"Skipping instance {instance_id} with streamk, async, float since it's not supported yet.") + if ( + native.streamk_enabled + and get_dtype(problem_name) == "float" + and native.pipeline_version.find("ASYNC") != -1 + ): + print( + f"Skipping instance {instance_id} with streamk, async, float since it's not supported yet." + ) continue convs.append(native) continue @@ -666,8 +674,11 @@ def parse_bwd_weight_instances(instances, problem_name): ) continue if not check_warp_coverage( - m_per_block, n_per_block, k_per_block, - a_scalar_per_vector, b_scalar_per_vector, + m_per_block, + n_per_block, + k_per_block, + a_scalar_per_vector, + b_scalar_per_vector, variant="bwd_weight", ): print( @@ -711,20 +722,24 @@ def parse_bwd_data_instances(instances, problem_name): convs.append(native) continue - start = instance.index('<') + 1 - end = instance.rindex('>') + start = instance.index("<") + 1 + end = instance.rindex(">") params_str = instance[start:end] args = parse_instance_string(params_str) is_v1_instance = instance.find("Xdl_CShuffle<") != -1 - + if is_v1_instance: if len(args) != 51: - raise RuntimeError(f"Wrong number of parameters in the V1 XDL CShuffle instance string: {instance}\n" + - f"Expected 51 parameters for V1 instance. Found {len(args)} parameters.") + raise RuntimeError( + f"Wrong number of parameters in the V1 XDL CShuffle instance string: {instance}\n" + + f"Expected 51 parameters for V1 instance. Found {len(args)} parameters." + ) else: - raise RuntimeError(f"Only V1 XDL CShuffle instances are supported for backward data. Found instance: {instance}") - + raise RuntimeError( + f"Only V1 XDL CShuffle instances are supported for backward data. Found instance: {instance}" + ) + spec = args[13] block_size = int(args[17]) m_per_block = int(args[18]) @@ -741,8 +756,10 @@ def parse_bwd_data_instances(instances, problem_name): c_scalar_per_vector = int(args[44]) if ak1 != bk1: - raise RuntimeError(f"Not supported instance {instance_id} since ak1 != bk1. ak1: {ak1}, bk1: {bk1} in instance: {instance}") - + raise RuntimeError( + f"Not supported instance {instance_id} since ak1 != bk1. ak1: {ak1}, bk1: {bk1} in instance: {instance}" + ) + k1 = min(ak1, bk1) # TODO: Do we need split image for 3D bwd data convs? @@ -768,9 +785,13 @@ def parse_bwd_data_instances(instances, problem_name): # Scheduler must be either Intrawave or Interwave. # Version must be from v1 to v5 if block_gemm_pipeline_scheduler not in ["Intrawave", "Interwave"]: - raise RuntimeError(f"Invalid Block GEMM pipeline scheduler: {block_gemm_pipeline_scheduler} in instance: {instance}") + raise RuntimeError( + f"Invalid Block GEMM pipeline scheduler: {block_gemm_pipeline_scheduler} in instance: {instance}" + ) if blk_gemm_pipeline_version not in ["v1", "v2", "v3", "v4", "v5"]: - raise RuntimeError(f"Invalid Block GEMM pipeline version: {blk_gemm_pipeline_version} in instance: {instance}") + raise RuntimeError( + f"Invalid Block GEMM pipeline version: {blk_gemm_pipeline_version} in instance: {instance}" + ) double_smem_buffer = blk_gemm_pipeline_version == "v4" scheduler = block_gemm_pipeline_scheduler @@ -798,25 +819,41 @@ def parse_bwd_data_instances(instances, problem_name): k_per_xdl = min(max(k1, get_k_mfma(dtype, m_per_xdl, n_per_xdl)), k_per_block) - # Skip irregular vector sizes — no HW vector load instructions for odd widths - if not check_vectors(a_scalar_per_vector, b_scalar_per_vector, c_scalar_per_vector): - print(f"Skipping instance {instance_id} with irregular load since it's not supported yet.") + # Skip irregular vector sizes -- no HW vector load instructions for odd widths + if not check_vectors( + a_scalar_per_vector, b_scalar_per_vector, c_scalar_per_vector + ): + print( + f"Skipping instance {instance_id} with irregular load since it's not supported yet." + ) continue # Skip multi-warp: single warp can't cover tile dim when it exceeds warp_size * vec if not check_warp_coverage( - m_per_block, n_per_block, k_per_block, - a_scalar_per_vector, b_scalar_per_vector, + m_per_block, + n_per_block, + k_per_block, + a_scalar_per_vector, + b_scalar_per_vector, variant="bwd_data", ): - print(f"Skipping instance {instance_id} with multiple warps per continous tile dim since it's not supported yet.") + print( + f"Skipping instance {instance_id} with multiple warps per continous tile dim since it's not supported yet." + ) continue if not check_bwd_data_vec_coverage( - m_per_block, n_per_block, k_per_block, - m_warp, n_warp, k_warp, - a_scalar_per_vector, b_scalar_per_vector, + m_per_block, + n_per_block, + k_per_block, + m_warp, + n_warp, + k_warp, + a_scalar_per_vector, + b_scalar_per_vector, ): - print(f"Skipping instance {instance_id} because current scalar per vector exceedes tile size") + print( + f"Skipping instance {instance_id} because current scalar per vector exceedes tile size" + ) continue conv = ConvInstanceTemplateParams( @@ -836,7 +873,7 @@ def parse_bwd_data_instances(instances, problem_name): instance_id, ) convs.append(conv) - + return convs @@ -945,18 +982,18 @@ def process_direction( DEPTHWISE_CONFIGS = [ { - "name": "ngchw_depthwise_fp32", - "conf": "ngchw_depthwise.conf", + "name": "ngchw_depthwise_fp32", + "conf": "ngchw_depthwise.conf", "signature": "SIGNATURE_NGCHW_FP32_FWD", }, { - "name": "ngchw_depthwise_fp16", - "conf": "ngchw_depthwise.conf", + "name": "ngchw_depthwise_fp16", + "conf": "ngchw_depthwise.conf", "signature": "SIGNATURE_NGCHW_FP16_FWD", }, { - "name": "ngchw_depthwise_bf16", - "conf": "ngchw_depthwise.conf", + "name": "ngchw_depthwise_bf16", + "conf": "ngchw_depthwise.conf", "signature": "SIGNATURE_NGCHW_BF16_FWD", }, ] @@ -989,36 +1026,53 @@ def parse_depthwise_config(conf_path: Path) -> list: return instances -def generate_depthwise_cpp(params: list, instance_name: str, signature: str, cpp_out: Path) -> None: - tile_h, tile_w, filt, str_h, str_w, pad_h, pad_w, nbatch, sub_h, sub_w, in_vec, out_vec = params +def generate_depthwise_cpp( + params: list, instance_name: str, signature: str, cpp_out: Path +) -> None: + ( + tile_h, + tile_w, + filt, + str_h, + str_w, + pad_h, + pad_w, + nbatch, + sub_h, + sub_w, + in_vec, + out_vec, + ) = params parent_dir = Path(__file__).resolve().parent template_file = parent_dir / "include/grouped_convolution_depthwise_tile.cpp.in" content = template_file.read_text() - content = content.replace("gen_signature", signature) + content = content.replace("gen_signature", signature) content = content.replace("gen_instance_name", instance_name) - content = content.replace("gen_block_size", "64") - content = content.replace("gen_tile_h", str(tile_h)) - content = content.replace("gen_tile_w", str(tile_w)) - content = content.replace("gen_filter_h", str(filt)) - content = content.replace("gen_filter_w", str(filt)) - content = content.replace("gen_stride_h", str(str_h)) - content = content.replace("gen_stride_w", str(str_w)) - content = content.replace("gen_dilation_h", "1") - content = content.replace("gen_dilation_w", "1") - content = content.replace("gen_pad_h", str(pad_h)) - content = content.replace("gen_pad_w", str(pad_w)) - content = content.replace("gen_nbatch", str(nbatch)) - content = content.replace("gen_subtile_h", str(sub_h)) - content = content.replace("gen_subtile_w", str(sub_w)) - content = content.replace("gen_in_vec", str(in_vec)) - content = content.replace("gen_out_vec", str(out_vec)) + content = content.replace("gen_block_size", "64") + content = content.replace("gen_tile_h", str(tile_h)) + content = content.replace("gen_tile_w", str(tile_w)) + content = content.replace("gen_filter_h", str(filt)) + content = content.replace("gen_filter_w", str(filt)) + content = content.replace("gen_stride_h", str(str_h)) + content = content.replace("gen_stride_w", str(str_w)) + content = content.replace("gen_dilation_h", "1") + content = content.replace("gen_dilation_w", "1") + content = content.replace("gen_pad_h", str(pad_h)) + content = content.replace("gen_pad_w", str(pad_w)) + content = content.replace("gen_nbatch", str(nbatch)) + content = content.replace("gen_subtile_h", str(sub_h)) + content = content.replace("gen_subtile_w", str(sub_w)) + content = content.replace("gen_in_vec", str(in_vec)) + content = content.replace("gen_out_vec", str(out_vec)) cpp_out.write_text(content) -def generate_depthwise_defs_inc(instances: list, config_name: str, signature: str, inc_path: Path) -> None: +def generate_depthwise_defs_inc( + instances: list, config_name: str, signature: str, inc_path: Path +) -> None: lines = [] for i in range(len(instances)): name = f"grouped_convolution_forward_tile_{config_name}_{i}" @@ -1032,7 +1086,9 @@ def generate_depthwise_defs_inc(instances: list, config_name: str, signature: st inc_path.write_text("\n".join(lines) + "\n") -def generate_depthwise_calls_inc(instances: list, config_name: str, calls_path: Path) -> None: +def generate_depthwise_calls_inc( + instances: list, config_name: str, calls_path: Path +) -> None: lines = [] for i in range(len(instances)): name = f"grouped_convolution_forward_tile_{config_name}_{i}" @@ -1044,11 +1100,11 @@ def process_depthwise_forward(configs_prefix: str, instances_path: str) -> None: """Generate all depthwise forward instances.""" generate_dir = Path(__file__).resolve().parent conf_dir = generate_dir / "configs/forward" / configs_prefix - inc_dir = generate_dir / "instances" / "forward" + inc_dir = generate_dir / "instances" / "forward" cpp_base = Path(instances_path) / "forward" for cfg in DEPTHWISE_CONFIGS: - name = cfg["name"] + name = cfg["name"] conf_path = conf_dir / cfg["conf"] signature = cfg["signature"] @@ -1064,24 +1120,33 @@ def process_depthwise_forward(configs_prefix: str, instances_path: str) -> None: for i, params in enumerate(instances): instance_name = f"grouped_convolution_forward_tile_{name}_{i}" - generate_depthwise_cpp(params, instance_name, signature, - cpp_dir / f"{instance_name}.cpp") + generate_depthwise_cpp( + params, instance_name, signature, cpp_dir / f"{instance_name}.cpp" + ) - generate_depthwise_defs_inc(instances, name, signature, - inc_dir / f"grouped_convolution_forward_tile_{name}.inc") - generate_depthwise_calls_inc(instances, name, - inc_dir / f"grouped_convolution_forward_tile_{name}_calls.inc") + generate_depthwise_defs_inc( + instances, + name, + signature, + inc_dir / f"grouped_convolution_forward_tile_{name}.inc", + ) + generate_depthwise_calls_inc( + instances, + name, + inc_dir / f"grouped_convolution_forward_tile_{name}_calls.inc", + ) print(f" -> {cpp_dir} ({len(instances)} .cpp files)") + fwd_configs = [ - "nhwgc_fp32", - "nhwgc_fp16", - "nhwgc_bf16", - "ndhwgc_fp32", - "ndhwgc_fp16", - "ndhwgc_bf16", - ] + "nhwgc_fp32", + "nhwgc_fp16", + "nhwgc_bf16", + "ndhwgc_fp32", + "ndhwgc_fp16", + "ndhwgc_bf16", +] bwd_weight_configs = [ "nhwgc_fp32", @@ -1153,7 +1218,14 @@ def process_depthwise_forward(configs_prefix: str, instances_path: str) -> None: copy_includes(args.instances_dir) match args.direction: case "forward": - process_direction(fwd_configs, args.direction, generate_instances_fwd, configs_prefix, args.filter_pattern, args.instances_dir) + process_direction( + fwd_configs, + args.direction, + generate_instances_fwd, + configs_prefix, + args.filter_pattern, + args.instances_dir, + ) process_depthwise_forward(configs_prefix, args.instances_dir) case "backward_weight": process_direction( @@ -1174,8 +1246,28 @@ def process_depthwise_forward(configs_prefix: str, instances_path: str) -> None: args.instances_dir, ) case "all": - process_direction(fwd_configs, "forward", generate_instances_fwd, configs_prefix, args.filter_pattern, args.instances_dir) + process_direction( + fwd_configs, + "forward", + generate_instances_fwd, + configs_prefix, + args.filter_pattern, + args.instances_dir, + ) process_depthwise_forward(configs_prefix, args.instances_dir) - process_direction(bwd_weight_configs, "backward_weight", generate_instances_bwd_weight, configs_prefix, args.filter_pattern, args.instances_dir) - process_direction(bwd_data_configs, "backward_data", generate_instances_bwd_data, configs_prefix, args.filter_pattern, args.instances_dir) - + process_direction( + bwd_weight_configs, + "backward_weight", + generate_instances_bwd_weight, + configs_prefix, + args.filter_pattern, + args.instances_dir, + ) + process_direction( + bwd_data_configs, + "backward_data", + generate_instances_bwd_data, + configs_prefix, + args.filter_pattern, + args.instances_dir, + ) diff --git a/include/ck_tile/core.hpp b/include/ck_tile/core.hpp index 47ee6c608c..4afba77d6a 100644 --- a/include/ck_tile/core.hpp +++ b/include/ck_tile/core.hpp @@ -17,6 +17,7 @@ #include "ck_tile/core/arch/arch.hpp" #include "ck_tile/core/arch/barrier.hpp" #include "ck_tile/core/arch/generic_memory_space_atomic.hpp" +#include "ck_tile/core/arch/inst_prefetch.hpp" #include "ck_tile/core/arch/mma/amdgcn_mma.hpp" #include "ck_tile/core/arch/mma/mfma/mfma.hpp" #include "ck_tile/core/arch/mma/mfma/mfma_gfx9.hpp" diff --git a/include/ck_tile/core/arch/inst_prefetch.hpp b/include/ck_tile/core/arch/inst_prefetch.hpp index d9f1cc4eb3..47170f63e5 100644 --- a/include/ck_tile/core/arch/inst_prefetch.hpp +++ b/include/ck_tile/core/arch/inst_prefetch.hpp @@ -1,91 +1,91 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. - -#pragma once - -#include "ck_tile/core.hpp" - -// ─── ISA label markers for two-pass instruction-prefetch offset patching ─── -// Used by script/patch_prefetch_offset.py to locate prefetch sites and targets -// in compiled GPU assembly and patch the koffset field. - -// Stringify helpers (shared by INST_PREFETCH_TARGET and INST_PREFETCH) -// CK_TILE_STR_ stringifies directly; CK_TILE_XSTR_ expands macros first. -#ifndef CK_TILE_STR_ -#define CK_TILE_STR_(x) #x -#define CK_TILE_XSTR_(x) CK_TILE_STR_(x) -#endif - -// INST_PREFETCH_TARGET(label) — default mode (mode=0): target is first instruction after -// comment. INST_PREFETCH_TARGET(label, mode) — mode=1 (BLOCK_ENTRY): script scans backward to -// nearest block -// label (.LBB*:) and uses the first instruction after that. -// Use when the compiler hoists ALU between the block entry -// and the asm comment. -#define CK_PLACE_MODE_DEFAULT 0 -#define CK_PLACE_MODE_BLOCK_ENTRY 1 - -#ifndef INST_PREFETCH_TARGET -#define INST_PREFETCH_TARGET_1(lbl) asm volatile("; [ck_label] name=" CK_TILE_STR_(lbl) " mode=0") -#define INST_PREFETCH_TARGET_2(lbl, mode) \ - asm volatile("; [ck_label] name=" CK_TILE_STR_(lbl) " mode=" CK_TILE_STR_(mode)) - -#define INST_PREFETCH_TARGET_GET_MACRO(_1, _2, NAME, ...) NAME -#define INST_PREFETCH_TARGET(...) \ - INST_PREFETCH_TARGET_GET_MACRO(__VA_ARGS__, INST_PREFETCH_TARGET_2, INST_PREFETCH_TARGET_1) \ - (__VA_ARGS__) -#endif - -// INST_PREFETCH(label, num_cachelines) -// INST_PREFETCH(label, num_cachelines, direction) -// INST_PREFETCH(label, num_cachelines, direction, offset_cachelines) -// Emit the [ck_prefetch] comment AND s_prefetch_inst_pc_rel with koffset=0. -// num_cachelines: number of 128B cache lines to prefetch (klength = num_cachelines - 1). -// direction: CK_PREFETCH_DIR_FORWARD (default) or CK_PREFETCH_DIR_BACKWARD. -// Forward: INST_PREFETCH_TARGET marks the first cacheline of the prefetch region. -// Backward: INST_PREFETCH_TARGET marks the last cacheline; the prefetch region extends -// backward by num_cachelines from INST_PREFETCH_TARGET. -// offset_cachelines: additional offset in cachelines added to the computed koffset. -// Allows multiple INST_PREFETCHes to share the same INST_PREFETCH_TARGET label but cover -// different sub-regions, e.g. INST_PREFETCH(lbl, 32, DIR_FORWARD, 0) and -// INST_PREFETCH(lbl, 32, DIR_FORWARD, 32) cover 64 cachelines total. -// The koffset is patched by script/patch_prefetch_offset.py in a second pass. -// Only emits code on gfx12+; on other targets it is a no-op. -#define CK_PREFETCH_DIR_FORWARD forward -#define CK_PREFETCH_DIR_BACKWARD backward - -#ifndef INST_PREFETCH -#if defined(__gfx12__) -#define INST_PREFETCH_4(lbl, num_cachelines, direction, offset_cachelines) \ - do \ - { \ - asm volatile( \ - "; [ck_prefetch] name=" CK_TILE_STR_(lbl) " dir=" CK_TILE_XSTR_( \ - direction) " offset=" CK_TILE_XSTR_(offset_cachelines) "\n\t" \ - "s_prefetch_inst_pc_rel " \ - "0, null, %0" \ - : \ - : "n"((num_cachelines) - 1)); \ - } while(false) -#define INST_PREFETCH_3(lbl, num_cachelines, direction) \ - INST_PREFETCH_4(lbl, num_cachelines, direction, 0) -#define INST_PREFETCH_2(lbl, num_cachelines) \ - INST_PREFETCH_3(lbl, num_cachelines, CK_PREFETCH_DIR_FORWARD) -#define INST_PREFETCH_GET_MACRO(_1, _2, _3, _4, NAME, ...) NAME -#define INST_PREFETCH(...) \ - INST_PREFETCH_GET_MACRO(__VA_ARGS__, INST_PREFETCH_4, INST_PREFETCH_3, INST_PREFETCH_2) \ - (__VA_ARGS__) -#else -#define INST_PREFETCH(lbl, ...) -#endif -#endif - -// Enable scalar prefetch in hardware (required on gfx12 before using s_prefetch) -__device__ __forceinline__ void enable_scalar_prefetch() -{ -#if defined(__gfx12__) - // SCALAR_PREFETCH_EN is bit 24 in MODE register (hwreg 1) - // Set 1 bit at offset 24 to value 1 - __builtin_amdgcn_s_setreg(1 | (24 << 6), 1); -#endif -} +// SPDX-License-Identifier: MIT +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + +#pragma once + +#include "ck_tile/core.hpp" + +// ─── ISA label markers for two-pass instruction-prefetch offset patching ─── +// Used by script/patch_prefetch_offset.py to locate prefetch sites and targets +// in compiled GPU assembly and patch the koffset field. + +// Stringify helpers (shared by INST_PREFETCH_TARGET and INST_PREFETCH) +// CK_TILE_STR_ stringifies directly; CK_TILE_XSTR_ expands macros first. +#ifndef CK_TILE_STR_ +#define CK_TILE_STR_(x) #x +#define CK_TILE_XSTR_(x) CK_TILE_STR_(x) +#endif + +// INST_PREFETCH_TARGET(label) — default mode (mode=0): target is first instruction after +// comment. INST_PREFETCH_TARGET(label, mode) — mode=1 (BLOCK_ENTRY): script scans backward to +// nearest block +// label (.LBB*:) and uses the first instruction after that. +// Use when the compiler hoists ALU between the block entry +// and the asm comment. +#define CK_PLACE_MODE_DEFAULT 0 +#define CK_PLACE_MODE_BLOCK_ENTRY 1 + +#ifndef INST_PREFETCH_TARGET +#define INST_PREFETCH_TARGET_1(lbl) asm volatile("; [ck_label] name=" CK_TILE_STR_(lbl) " mode=0") +#define INST_PREFETCH_TARGET_2(lbl, mode) \ + asm volatile("; [ck_label] name=" CK_TILE_STR_(lbl) " mode=" CK_TILE_STR_(mode)) + +#define INST_PREFETCH_TARGET_GET_MACRO(_1, _2, NAME, ...) NAME +#define INST_PREFETCH_TARGET(...) \ + INST_PREFETCH_TARGET_GET_MACRO(__VA_ARGS__, INST_PREFETCH_TARGET_2, INST_PREFETCH_TARGET_1) \ + (__VA_ARGS__) +#endif + +// INST_PREFETCH(label, num_cachelines) +// INST_PREFETCH(label, num_cachelines, direction) +// INST_PREFETCH(label, num_cachelines, direction, offset_cachelines) +// Emit the [ck_prefetch] comment AND s_prefetch_inst_pc_rel with koffset=0. +// num_cachelines: number of 128B cache lines to prefetch (klength = num_cachelines - 1). +// direction: CK_PREFETCH_DIR_FORWARD (default) or CK_PREFETCH_DIR_BACKWARD. +// Forward: INST_PREFETCH_TARGET marks the first cacheline of the prefetch region. +// Backward: INST_PREFETCH_TARGET marks the last cacheline; the prefetch region extends +// backward by num_cachelines from INST_PREFETCH_TARGET. +// offset_cachelines: additional offset in cachelines added to the computed koffset. +// Allows multiple INST_PREFETCHes to share the same INST_PREFETCH_TARGET label but cover +// different sub-regions, e.g. INST_PREFETCH(lbl, 32, DIR_FORWARD, 0) and +// INST_PREFETCH(lbl, 32, DIR_FORWARD, 32) cover 64 cachelines total. +// The koffset is patched by script/patch_prefetch_offset.py in a second pass. +// Only emits code on gfx12+; on other targets it is a no-op. +#define CK_PREFETCH_DIR_FORWARD forward +#define CK_PREFETCH_DIR_BACKWARD backward + +#ifndef INST_PREFETCH +#if defined(__gfx12__) +#define INST_PREFETCH_4(lbl, num_cachelines, direction, offset_cachelines) \ + do \ + { \ + asm volatile( \ + "; [ck_prefetch] name=" CK_TILE_STR_(lbl) " dir=" CK_TILE_XSTR_( \ + direction) " offset=" CK_TILE_XSTR_(offset_cachelines) "\n\t" \ + "s_prefetch_inst_pc_rel " \ + "0, null, %0" \ + : \ + : "n"((num_cachelines) - 1)); \ + } while(false) +#define INST_PREFETCH_3(lbl, num_cachelines, direction) \ + INST_PREFETCH_4(lbl, num_cachelines, direction, 0) +#define INST_PREFETCH_2(lbl, num_cachelines) \ + INST_PREFETCH_3(lbl, num_cachelines, CK_PREFETCH_DIR_FORWARD) +#define INST_PREFETCH_GET_MACRO(_1, _2, _3, _4, NAME, ...) NAME +#define INST_PREFETCH(...) \ + INST_PREFETCH_GET_MACRO(__VA_ARGS__, INST_PREFETCH_4, INST_PREFETCH_3, INST_PREFETCH_2) \ + (__VA_ARGS__) +#else +#define INST_PREFETCH(lbl, ...) +#endif +#endif + +// Enable scalar prefetch in hardware (required on gfx12 before using s_prefetch) +__device__ __forceinline__ void enable_scalar_prefetch() +{ +#if defined(__gfx12__) + // SCALAR_PREFETCH_EN is bit 24 in MODE register (hwreg 1) + // Set 1 bit at offset 24 to value 1 + __builtin_amdgcn_s_setreg(1 | (24 << 6), 1); +#endif +} diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index c9037a3408..ae9b698ae1 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -563,6 +563,33 @@ struct CShuffleEpilogue return lds_block_desc.get_element_space_size() * sizeof(ODataType); } + /// Number of block_sync_lds() calls in operator(). + /// Used by RunBarrierStub() to match barrier count for wavelet load waves. + /// IMPORTANT: Must be kept in sync with operator(). See RunBarrierStub(). + CK_TILE_HOST_DEVICE static constexpr index_t GetBarrierCount() + { + // operator() issues: + // 1x s_wait_tensorcnt_barrier() (counted as 1 barrier) + // num_access iterations x 2 block_sync_lds() each + constexpr index_t num_access = SFC::get_num_of_access(); + return 1 + 2 * num_access; + } + + /// Run matching barriers for wavelet load waves that don't participate + /// in the epilogue data path. Must issue the same number of barriers + /// as operator() to avoid deadlock. + CK_TILE_DEVICE static void RunBarrierStub() + { + constexpr index_t num_access = SFC::get_num_of_access(); + constexpr index_t count = GetBarrierCount(); + // Verify the barrier count formula matches the structural pattern in operator(): + // 1 x s_wait_tensorcnt_barrier + num_access x 2 block_sync_lds + static_assert(count == 1 + 2 * num_access, + "RunBarrierStub: barrier count mismatch with operator(). " + "If operator()'s barrier pattern changed, update GetBarrierCount()."); + static_for<0, count, 1>{}([&](auto) { block_sync_lds(); }); + } + template CK_TILE_DEVICE void scale_tile(LdsTile& lds_tile, ScaleM& scale_m_window, ScaleN& scale_n_window) @@ -784,6 +811,9 @@ struct CShuffleEpilogue } }(); + // NOTE: This barrier pattern must match GetBarrierCount(). + // Total barriers = 1 (s_wait_tensorcnt_barrier) + 2 * num_access (block_sync_lds pairs). + // If you add/remove barriers here, update GetBarrierCount() and RunBarrierStub(). s_wait_tensorcnt_barrier(); static_for<0, num_access, 1>{}([&](auto iAccess) { diff --git a/include/ck_tile/ops/gemm.hpp b/include/ck_tile/ops/gemm.hpp index b2b927b51f..8aae2bb351 100644 --- a/include/ck_tile/ops/gemm.hpp +++ b/include/ck_tile/ops/gemm.hpp @@ -66,6 +66,7 @@ #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_eight_waves_base.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wavelet.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_async_v1.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp" diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wavelet.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wavelet.hpp new file mode 100644 index 0000000000..f18acb241d --- /dev/null +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wavelet.hpp @@ -0,0 +1,361 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp" + +namespace ck_tile { + +/// @brief Wavelet GEMM pipeline: 2-way wave specialization for load/math overlap. +/// +/// Splits the workgroup into load waves and math waves: +/// - Math threads [0, MathBlockSize): LDS reads + MFMA accumulation only +/// - Load threads [MathBlockSize, LaunchBlockSize): global loads + LDS writes only +/// +/// Load threads use a remapped partition index ({warp_id - NumMathWarps, lane_id}) +/// so they appear as virtual threads [0, LoadBlockSize) to tile distributions. +/// This allows using standard Problem-sized distributions (MathBlockSize threads) +/// without modifying the core tile distribution infrastructure. +/// +/// The pipeline returns c_block_tile. Math threads hold the real accumulated result; +/// load threads return a zero-initialized c_block_tile. +/// +/// @tparam Problem The GEMM pipeline problem. Problem::kBlockSize == MathBlockSize +/// (the standard NumWarps * warp_size for BlockGemm compatibility). +/// @tparam Policy The universal pipeline policy. +/// @tparam NumLoadWaves_ Number of additional load waves (default 4). +template +struct GemmPipelineAgBgCrWavelet +{ + using PipelineImplBase = GemmPipelineAgBgCrImplBase; + + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; + + using BlockGemm = remove_cvref_t())>; + using I0 = number<0>; + using I1 = number<1>; + + // --- Thread group sizes --- + // MathBlockSize: the standard block size for BlockGemm (NumWarps * warp_size). + // This is Problem::kBlockSize. + static constexpr index_t MathBlockSize = Problem::kBlockSize; + + // LoadBlockSize: extra threads dedicated to loading. + static constexpr index_t NumLoadWaves = NumLoadWaves_; + static constexpr index_t LoadBlockSize = NumLoadWaves * get_warp_size(); + + // LaunchBlockSize: total threads launched per workgroup. + static constexpr index_t LaunchBlockSize = MathBlockSize + LoadBlockSize; + + // BlockSize exposed to the kernel for BlockGemm compatibility. + // The kernel uses this for MFMA wave-to-tile mapping. + static constexpr index_t BlockSize = MathBlockSize; + + // Standard pipeline interface members (required by builder instance traits). + static constexpr auto Scheduler = Problem::Scheduler; + static constexpr bool DoubleSmemBuffer = false; + static constexpr index_t NumWaveGroups = 1; + + // --- Wavelet traits --- + static constexpr bool IsWavelet = true; + + CK_TILE_DEVICE static bool IsMathWave() { return get_thread_local_1d_id() < MathBlockSize; } + + // --- Pipeline traits (matching standard pipeline interface) --- + static constexpr index_t PrefetchStages = 1; + static constexpr index_t PrefillStages = 1; + static constexpr index_t GlobalBufferNum = 1; + + static constexpr index_t MPerBlock = BlockGemmShape::kM; + static constexpr index_t NPerBlock = BlockGemmShape::kN; + static constexpr index_t KPerBlock = BlockGemmShape::kK; + + static constexpr bool kPadM = Problem::kPadM; + static constexpr bool kPadN = Problem::kPadN; + static constexpr bool kPadK = Problem::kPadK; + + static constexpr auto is_a_load_tr_v = bool_constant{}; + static constexpr auto is_b_load_tr_v = bool_constant{}; + + static constexpr bool Async = false; + static constexpr bool UsePersistentKernel = false; + + template + static constexpr index_t GetVectorSizeA() + { + return Policy::template GetVectorSizeA(); + } + template + static constexpr index_t GetVectorSizeB() + { + return Policy::template GetVectorSizeB(); + } + static constexpr index_t GetVectorSizeC() { return Policy::template GetVectorSizeC(); } + + [[nodiscard]] CK_TILE_HOST static const std::string GetPipelineName() { return "WAVELET"; } + + [[nodiscard]] CK_TILE_HOST static const std::string GetName() + { + constexpr index_t WaveNumM = BlockGemmShape::BlockWarps::at(I0{}); + constexpr index_t WaveNumN = BlockGemmShape::BlockWarps::at(I1{}); + return concat('_', + "pipeline_AgBgCrWavelet", + concat('x', MPerBlock, NPerBlock, KPerBlock), + BlockSize, + concat('x', GetVectorSizeA(), GetVectorSizeB(), GetVectorSizeC()), + concat('x', WaveNumM, WaveNumN), + concat('x', NumLoadWaves), + Problem::GetName()); + } + + CK_TILE_HOST_DEVICE static constexpr bool BlockHasHotloop(index_t num_loop) + { + return num_loop > 1; + } + + CK_TILE_HOST_DEVICE static constexpr TailNumber GetBlockLoopTailNum(index_t /*num_loop*/) + { + return TailNumber::Odd; + } + + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + return Policy::template GetSmemSize(); + } + + static constexpr index_t NumMathWarps = MathBlockSize / get_warp_size(); + + static constexpr index_t APackedSize = + ck_tile::numeric_traits>::PackedSize; + static constexpr index_t BPackedSize = + ck_tile::numeric_traits>::PackedSize; + + static constexpr index_t GetSmemPackA() { return Policy::template GetSmemPackA(); } + static constexpr index_t GetSmemPackB() { return Policy::template GetSmemPackB(); } + + // LoadProblem: same as Problem but with kBlockSize = LoadBlockSize. + // Used for DRAM tile distributions so that LoadBlockSize threads (which may + // differ from MathBlockSize) cover the full tile cooperatively. + struct LoadProblem : Problem + { + static constexpr index_t kBlockSize = LoadBlockSize; + }; + + // --- Main pipeline operator --- + template + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + // ---------------------------------------------------------------- + // Thread role. + // + // VGPR pressure: the body is split into two TOP-LEVEL, mutually + // exclusive branches by role. Each branch constructs ONLY its own + // windows / tiles / accumulator, so the register allocator can map the + // two roles' (disjoint) live ranges onto the same physical VGPRs -- the + // kernel's single per-wave VGPR count approaches max(load, math) rather + // than load + math. Constructing any role-specific object before the + // branch would force every wave to carry it and defeat this. + // + // BARRIER INVARIANT (critical -- mismatch deadlocks the workgroup): + // both branches execute exactly 1 + 2*(num_loop-1) block_sync_lds() + // calls, in the same order. AMD s_barrier pairs the Nth barrier each + // wave reaches, so identical counts (not identical PCs) are required. + // ---------------------------------------------------------------- + const bool is_math = IsMathWave(); + + // LDS views over shared smem -- cheap (descriptors), used by both roles. + auto&& [a_lds_block, b_lds_block] = PipelineImplBase{}.GetABLdsTensorViews(p_smem); + + constexpr bool is_a_col_major = std::is_same_v; + constexpr bool is_b_row_major = std::is_same_v; + + // LDS window distributions (compile-time, no registers). + constexpr auto a_lds_load_tile_distr = + make_static_tile_distribution(BlockGemm::MakeABlockDistributionEncode()); + constexpr auto b_lds_load_tile_distr = + make_static_tile_distribution(BlockGemm::MakeBBlockDistributionEncode()); + + // Accumulator type -- used so both branches return the same type. Load + // waves return a default-constructed (dead) tile that RunGemm discards. + using CBlockTile = decltype(BlockGemm().MakeCBlockTile()); + + if(is_math) + { + // ============================================================ + // MATH WAVES: LDS read + MFMA accumulate. No DRAM/load state. + // ============================================================ + auto [a_copy_lds_window_unused, a_lds_gemm_window] = + PipelineImplBase{}.MakeALdsWindows(a_lds_block, a_lds_load_tile_distr); + auto [b_copy_lds_window_unused, b_lds_gemm_window] = + PipelineImplBase{}.MakeBLdsWindows(b_lds_block, b_lds_load_tile_distr); + (void)a_copy_lds_window_unused; + (void)b_copy_lds_window_unused; + + auto block_gemm = BlockGemm(); + auto c_block_tile = block_gemm.MakeCBlockTile(); + tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + + // Prologue barrier (#0): wait for load waves to fill LDS iter 0. + block_sync_lds(); + + index_t i = 0; + while(i < num_loop - 1) + { + // MFMA on current LDS data (concurrent with load waves' DRAM fetch). + block_gemm.LocalPrefetch( + a_lds_gemm_window, b_lds_gemm_window, is_a_load_tr_v, is_b_load_tr_v); + block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); + + block_sync_lds(); // barrier A: math done reading LDS + block_sync_lds(); // barrier B: load done writing next LDS (math idle between) + + ++i; + } + + // Tail: last iteration's MFMA. + block_gemm.LocalPrefetch( + a_lds_gemm_window, b_lds_gemm_window, is_a_load_tr_v, is_b_load_tr_v); + block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); + + return c_block_tile; + } + else + { + // ============================================================ + // LOAD WAVES: DRAM read + LDS write. No accumulator/MFMA state. + // ============================================================ + // Remap warp ID so load threads appear as virtual threads + // [0, LoadBlockSize) to the LoadProblem tile distributions. + const auto load_partition = + array{get_warp_id() - NumMathWarps, get_lane_id()}; + + using YPerTileA = + std::conditional_t, number>; + using XPerTileA = + std::conditional_t, number>; + auto a_copy_dram_window = + make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), + make_tuple(YPerTileA{}, XPerTileA{}), + a_dram_block_window_tmp.get_window_origin(), + Policy::template MakeADramTileDistribution(), + load_partition); + + using YPerTileB = + std::conditional_t, number>; + using XPerTileB = + std::conditional_t, number>; + auto b_copy_dram_window = + make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), + make_tuple(YPerTileB{}, XPerTileB{}), + b_dram_block_window_tmp.get_window_origin(), + Policy::template MakeBDramTileDistribution(), + load_partition); + + auto [a_copy_lds_window, a_lds_gemm_window_unused] = + PipelineImplBase{}.MakeALdsWindows(a_lds_block, a_lds_load_tile_distr); + auto [b_copy_lds_window, b_lds_gemm_window_unused] = + PipelineImplBase{}.MakeBLdsWindows(b_lds_block, b_lds_load_tile_distr); + (void)a_lds_gemm_window_unused; + (void)b_lds_gemm_window_unused; + + auto a_block_tile = decltype(load_tile(a_copy_dram_window)){}; + auto b_block_tile = decltype(load_tile(b_copy_dram_window)){}; + + using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; + using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; + constexpr ADramTileWindowStep a_dram_step = + is_a_col_major ? make_array(KPerBlock, 0) : make_array(0, KPerBlock); + constexpr BDramTileWindowStep b_dram_step = + is_b_row_major ? make_array(KPerBlock, 0) : make_array(0, KPerBlock); + + // store_tile helper with optional in-register transpose (for + // architectures without hardware transpose-load, e.g. gfx942). + auto store_a_to_lds = [&]() { + if constexpr(is_a_col_major && !PipelineImplBase::is_a_load_tr) + { + auto a_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledARegTileDistribution()); + transpose_tile2d(a_shuffle_tmp, a_block_tile); + store_tile(a_copy_lds_window, a_shuffle_tmp, load_partition); + } + else + { + store_tile(a_copy_lds_window, a_block_tile, load_partition); + } + }; + auto store_b_to_lds = [&]() { + if constexpr(is_b_row_major && !PipelineImplBase::is_b_load_tr) + { + auto b_shuffle_tmp = make_static_distributed_tensor( + Policy::template MakeShuffledBRegTileDistribution()); + transpose_tile2d(b_shuffle_tmp, b_block_tile); + store_tile(b_copy_lds_window, b_shuffle_tmp, load_partition); + } + else + { + store_tile(b_copy_lds_window, b_block_tile, load_partition); + } + }; + + // Prologue: fetch iteration 0 from DRAM -> LDS. + a_block_tile = load_tile(a_copy_dram_window); + move_tile_window(a_copy_dram_window, a_dram_step); + b_block_tile = load_tile(b_copy_dram_window); + move_tile_window(b_copy_dram_window, b_dram_step); + store_a_to_lds(); + store_b_to_lds(); + + // Prologue barrier (#0): LDS iter 0 ready for math waves. + block_sync_lds(); + + index_t i = 0; + while(i < num_loop - 1) + { + // Fetch next iteration from DRAM (concurrent with math MFMA). + a_block_tile = load_tile(a_copy_dram_window); + move_tile_window(a_copy_dram_window, a_dram_step); + b_block_tile = load_tile(b_copy_dram_window); + move_tile_window(b_copy_dram_window, b_dram_step); + + block_sync_lds(); // barrier A: math done reading current LDS + + // Write next iteration's data to LDS. + store_a_to_lds(); + store_b_to_lds(); + + block_sync_lds(); // barrier B: LDS populated for next iteration + + ++i; + } + + // Load waves hold no result; return a dead tile (discarded by RunGemm). + return CBlockTile{}; + } + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipelines.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipelines.hpp index 963d8fd89c..54886e21ca 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipelines.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipelines.hpp @@ -23,7 +23,8 @@ enum struct GemmPipeline PRESHUFFLE_FLATMM, PRESHUFFLE_TDM, PRESHUFFLE_MX_TDM, - COMPUTE_MX_TDM + COMPUTE_MX_TDM, + WAVELET }; } // namespace ck_tile diff --git a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp index 2a4e6a062d..9d031a989e 100644 --- a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp +++ b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp @@ -447,7 +447,21 @@ struct GroupedConvolutionBackwardWeightKernel using GemmDsLayout = remove_cvref_t; static constexpr index_t NumDTensor = GroupedConvTraitsType_::NumDTensor; - static constexpr index_t kBlockSize = GemmPipeline::BlockSize; + // For wavelet, LaunchBlockSize > BlockSize. Use LaunchBlockSize for kernel launch. + template + struct has_launch_block_size : std::false_type + { + }; + template + struct has_launch_block_size> : std::true_type + { + }; + static constexpr index_t kBlockSize = []() { + if constexpr(has_launch_block_size::value) + return GemmPipeline::LaunchBlockSize; + else + return GemmPipeline::BlockSize; + }(); using OutDataType = remove_cvref_t; using InDataType = remove_cvref_t; @@ -995,6 +1009,22 @@ struct GroupedConvolutionBackwardWeightKernel {block_idx_k, block_idx_m}); } + // SFINAE helper: detect GemmPipeline::IsWavelet + template + struct has_is_wavelet : std::false_type + { + }; + template + struct has_is_wavelet> : std::true_type + { + }; + static constexpr bool kIsWavelet = []() { + if constexpr(has_is_wavelet::value) + return GemmPipeline::IsWavelet; + else + return false; + }(); + /** * @brief Runs single GEMM problem cooperatively by whole workgroup. * @@ -1027,24 +1057,56 @@ struct GroupedConvolutionBackwardWeightKernel const auto& c_block_tile = GemmPipeline{}.template operator()( a_block_window, b_block_window, num_loop, smem_ptr_0); - // Run Epilogue Pipeline with k_batch dispatching - if(kargs.k_batch == 1) + if constexpr(kIsWavelet) { - auto c_block_window = MakeCBlockWindow( - c_ptr, kargs, block_idx_m, block_idx_n); - - EpiloguePipeline{}(c_block_window, c_block_tile, d_block_window, smem_ptr_0); + // Wavelet: math waves run the epilogue, load waves run matching barriers + if(GemmPipeline::IsMathWave()) + { + if(kargs.k_batch == 1) + { + auto c_block_window = MakeCBlockWindow( + c_ptr, kargs, block_idx_m, block_idx_n); + EpiloguePipeline{}(c_block_window, c_block_tile, d_block_window, smem_ptr_0); + } + else + { + if constexpr(!(GroupedConvTraitsType_::VectorSizeC % 2 != 0 && + is_any_of::value)) + { + auto c_block_window = MakeCBlockWindow( + c_ptr, kargs, block_idx_m, block_idx_n); + EpiloguePipeline{}( + c_block_window, c_block_tile, d_block_window, smem_ptr_0); + } + } + } + else + { + // Load waves: match epilogue barrier count to avoid deadlock + EpiloguePipeline::RunBarrierStub(); + } } else { - if constexpr(!(GroupedConvTraitsType_::VectorSizeC % 2 != 0 && - is_any_of::value)) + // Standard (non-wavelet) path + if(kargs.k_batch == 1) { - auto c_block_window = MakeCBlockWindow( + auto c_block_window = MakeCBlockWindow( c_ptr, kargs, block_idx_m, block_idx_n); EpiloguePipeline{}(c_block_window, c_block_tile, d_block_window, smem_ptr_0); } + else + { + if constexpr(!(GroupedConvTraitsType_::VectorSizeC % 2 != 0 && + is_any_of::value)) + { + auto c_block_window = MakeCBlockWindow( + c_ptr, kargs, block_idx_m, block_idx_n); + + EpiloguePipeline{}(c_block_window, c_block_tile, d_block_window, smem_ptr_0); + } + } } } From 919096fde86e964099e3f333b9538e78e38f4995 Mon Sep 17 00:00:00 2001 From: Sami Remes <181322991+samremes@users.noreply.github.com> Date: Tue, 2 Jun 2026 13:08:46 +0000 Subject: [PATCH 011/143] [rocm-libraries] ROCm/rocm-libraries#7935 (commit 5c96097) [CK] Allow skipping split-K C-buffer zero-init in xdl_cshuffle blockscale GEMM (#7935) Add a `skip_zero_init` flag (default false) to the Problem/Argument of the xdl_cshuffle block-scale GEMM device ops (multiple_d ab_scale and blockscale b-preshuffle). When the flag is set, the device invoker skips the internal hipMemsetAsync that zeroes p_c_grid before the KBatch > 1 split-K atomic-accumulation path. The flag is declared on the gridwise Problem struct (inherited by Argument), so it is visible on both the rotating-cache (arg_) and the normal (arg) launch paths in each device op. Why: callers that already pre-zero the output buffer otherwise pay for a redundant device-wide memset before split-K atomic accumulation. Gating the memset behind an opt-in flag lets such callers avoid the duplicate work. Because the flag defaults to false, every existing call site is unaffected and the observable behavior is unchanged. ## Motivation ## Technical Details ## Test Plan ## Test Result ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. Co-authored-by: Cursor --- .../impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp | 4 ++-- ...emm_multiple_d_xdl_cshuffle_v3_blockscale_bpreshuffle.hpp | 4 ++-- .../grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp | 5 +++++ ..._gemm_xdl_cshuffle_v3_multi_d_blockscale_b_preshuffle.hpp | 5 +++++ 4 files changed, 14 insertions(+), 4 deletions(-) diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp index cdfdb2abca..c682c1b56b 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp @@ -227,7 +227,7 @@ struct DeviceGemmMultiD_ABScale_Xdl_CShuffle_V3 // rotating mem rotating_mem.Next(); // clear c mem - if(arg_.KBatch > 1) + if(arg_.KBatch > 1 && !arg_.skip_zero_init) hipGetErrorString(hipMemsetAsync(arg_.p_c_grid, 0, arg_.M * arg_.N * sizeof(CDataType), @@ -245,7 +245,7 @@ struct DeviceGemmMultiD_ABScale_Xdl_CShuffle_V3 } else { - if(arg.KBatch > 1) + if(arg.KBatch > 1 && !arg.skip_zero_init) hipGetErrorString(hipMemsetAsync(arg.p_c_grid, 0, arg.M * arg.N * sizeof(CDataType), diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_blockscale_bpreshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_blockscale_bpreshuffle.hpp index c2982dfa95..39f25e491b 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_blockscale_bpreshuffle.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_blockscale_bpreshuffle.hpp @@ -230,7 +230,7 @@ struct DeviceGemmMultiD_BlockScale_Xdl_CShuffle_V3_BPreshuffle // rotating mem rotating_mem.Next(); // clear c mem - if(arg_.KBatch > 1) + if(arg_.KBatch > 1 && !arg_.skip_zero_init) hipGetErrorString(hipMemsetAsync(arg_.p_c_grid, 0, arg_.M * arg_.N * sizeof(CDataType), @@ -248,7 +248,7 @@ struct DeviceGemmMultiD_BlockScale_Xdl_CShuffle_V3_BPreshuffle } else { - if(arg.KBatch > 1) + if(arg.KBatch > 1 && !arg.skip_zero_init) hipGetErrorString(hipMemsetAsync(arg.p_c_grid, 0, arg.M * arg.N * sizeof(CDataType), diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp index 27d5c4be61..e70a1d5bf1 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp @@ -650,6 +650,11 @@ struct GridwiseGemmMultiD_ABScale_xdl_cshuffle_v3 index_t BK0; index_t MBlock; index_t NBlock; + // When true, the caller guarantees p_c_grid is already zeroed before + // launch, so the device invoker skips its own hipMemsetAsync for the + // KBatch > 1 split-K atomic-accumulation path. Defaults to false, which + // preserves the original zero-init behavior. + bool skip_zero_init = false; }; // Argument diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_blockscale_b_preshuffle.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_blockscale_b_preshuffle.hpp index c0abfb107c..7767d2c484 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_blockscale_b_preshuffle.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_blockscale_b_preshuffle.hpp @@ -669,6 +669,11 @@ struct GridwiseGemmMultiD_blockscale_xdl_cshuffle_v3_b_preshuffle index_t BK0; index_t MBlock; index_t NBlock; + // When true, the caller guarantees p_c_grid is already zeroed before + // launch, so the device invoker skips its own hipMemsetAsync for the + // KBatch > 1 split-K atomic-accumulation path. Defaults to false, which + // preserves the original zero-init behavior. + bool skip_zero_init = false; }; // Argument From 99ab4c4ef7e23588b725a84e30023dfd801cc5ba Mon Sep 17 00:00:00 2001 From: Aviral Goel <191153937+AviralGoelAMD@users.noreply.github.com> Date: Tue, 2 Jun 2026 13:54:16 +0000 Subject: [PATCH 012/143] [rocm-libraries] ROCm/rocm-libraries#7830 (commit 590fe58) [CK_Tile][MI450] Add bf16 output wmma instruction (16x16x32) (#7830) Wire __builtin_amdgcn_wmma_bf16_16x16x32_bf16 into CK Tile for gfx1250, enabling bf16-input bf16-output WMMA at the warp GEMM level. - Add WmmaTraits specialization for - Add WarpGemmAttributeWmmaImpl typedef and WarpGemmWmma alias - Add Dispatcher entry for bf16->bf16 16x16x32 - Add warp_gemm test with reference GEMM validation ## Motivation ## Technical Details ## Test Plan ## Test Result ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- .../warp/warp_gemm_attribute_wmma_impl.hpp | 3 + ..._gemm_attribute_wmma_impl_16bit_traits.hpp | 24 +++ .../ops/gemm/warp/warp_gemm_dispatcher.hpp | 2 + .../ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp | 7 + test/ck_tile/warp_gemm/CMakeLists.txt | 3 +- .../test_wmma_bf16_16x16x32_gfx1250.cpp | 150 ++++++++++++++++++ 6 files changed, 188 insertions(+), 1 deletion(-) create mode 100644 test/ck_tile/warp_gemm/test_wmma_bf16_16x16x32_gfx1250.cpp diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp index 8fd185cb42..6f38199828 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp @@ -146,6 +146,9 @@ using WarpGemmAttributeWmmaImpl_f32_16x16x32_f16_f16 = using WarpGemmAttributeWmmaImpl_f32_16x16x32_bf16_bf16 = WarpGemmAttributeWmmaImpl>; +using WarpGemmAttributeWmmaImpl_bf16_16x16x32_bf16_bf16 = + WarpGemmAttributeWmmaImpl>; + using WarpGemmAttributeWmmaImpl_i32_16x16x64_i8_i8 = WarpGemmAttributeWmmaImpl>; diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp index b5d7365dad..e770efeaeb 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp @@ -142,4 +142,28 @@ struct WmmaTraits #endif } }; + +// bf16 -> bf16 specialization - GFX125 +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + using ArchType = gfx125_t; + + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx125__ + using P = WarpGemmParamsParser; + return __builtin_amdgcn_wmma_bf16_16x16x32_bf16( + 0, a_vec, 0, b_vec, 0, c_vec, P::reuse_a, P::reuse_b); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0}; +#endif + } +}; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp index 2e6fa605ba..b77f09fa9f 100644 --- a/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp @@ -149,6 +149,8 @@ template<> struct Dispatcher { using T #if defined(__gfx125__) template struct Dispatcher : WmmaTag { using Type = WarpGemmWmma_f32_16x16x32_bf16_bf16;}; +template struct Dispatcher + : WmmaTag { using Type = WarpGemmWmma_bf16_16x16x32_bf16_bf16;}; #else template<> struct Dispatcher { using Type = WarpGemmMfmaBf16Bf16F32M16N16K32<>; }; template<> struct Dispatcher { using Type = WarpGemmMfmaBf16Bf16F32M16N16K32TransposedCDistribution<>; }; diff --git a/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp b/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp index 1c522d07c1..224937dc6d 100644 --- a/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp @@ -68,6 +68,13 @@ using WarpGemmWmma_f32_16x16x32_bf16_bf16 = AttrNumAccess, AttrNumAccess>>; +template +using WarpGemmWmma_bf16_16x16x32_bf16_bf16 = + WarpGemmImpl>; + template using WarpGemmWmma_f32_16x16x16_f8_bf8 = WarpGemmImpl +#include "ck_tile/host.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp" + +using namespace ck_tile; + +template +struct WGDispCase +{ + using AType = A; + using BType = B; + using AccType = Acc; + static constexpr index_t MPerWave = M; + static constexpr index_t NPerWave = N; + static constexpr index_t KPerWave = K; + static constexpr bool kTransposeC = TransposeC; +}; + +using WGDispatcherTypesList = + ::testing::Types, + WGDispCase>; + +template +struct WarpGemmKernel +{ + static constexpr int kBlockSize = 32; + __device__ void operator()(void* A, void* B, void* C) const + { + using WarpGemm = WarpGemmDispatcher; + + const auto a_view = + make_naive_tensor_view(static_cast(A), + make_tuple(M, K), + make_tuple(K, number<1>{}), + number{}, + number<1>{}); + + const auto b_view = + make_naive_tensor_view(static_cast(B), + make_tuple(N, K), + make_tuple(K, number<1>{}), + number{}, + number<1>{}); + + const auto c_view = + make_naive_tensor_view(static_cast(C), + make_tuple(M, N), + make_tuple(N, number<1>{}), + number{}, + number<1>{}); + + using AWarpTensor = typename WarpGemm::AWarpTensor; + using BWarpTensor = typename WarpGemm::BWarpTensor; + using CWarpTensor = typename WarpGemm::CWarpTensor; + + constexpr auto a_len = AWarpTensor::get_tile_distribution().get_lengths(); + constexpr auto b_len = BWarpTensor::get_tile_distribution().get_lengths(); + constexpr auto c_len = CWarpTensor::get_tile_distribution().get_lengths(); + + auto a_win = make_tile_window( + a_view, a_len, make_multi_index(0, 0), AWarpTensor::get_tile_distribution()); + auto b_win = make_tile_window( + b_view, b_len, make_multi_index(0, 0), BWarpTensor::get_tile_distribution()); + auto c_win = make_tile_window( + c_view, c_len, make_multi_index(0, 0), CWarpTensor::get_tile_distribution()); + + AWarpTensor a_tile; + BWarpTensor b_tile; + load_tile(a_tile, a_win); + load_tile(b_tile, b_win); + + auto c_tile = WarpGemm{}(a_tile, b_tile); + + store_tile(c_win, c_tile); + } +}; + +template +static void RunWarpGemmCase(const HostTensor& A, + const HostTensor& B, + HostTensor& C) +{ + DeviceMem Ad(A), Bd(B), Cd(C); + + using Kernel = WarpGemmKernel; + dim3 grid(1), block{Kernel::kBlockSize}; + + (void)launch_kernel(stream_config{nullptr, true, 0, 0, 1}, + make_kernel(Kernel{}, + grid, + block, + 0, + Ad.GetDeviceBuffer(), + Bd.GetDeviceBuffer(), + Cd.GetDeviceBuffer())); + + Cd.FromDevice(C.mData.data()); +} + +template +class WGRuntimeTest : public ::testing::Test +{ +}; + +TYPED_TEST_SUITE(WGRuntimeTest, WGDispatcherTypesList); + +TYPED_TEST(WGRuntimeTest, Compare_Dispatcher_ReferenceGemm) +{ + using Case = TypeParam; + + using AType = typename Case::AType; + using BType = typename Case::BType; + using CType = typename Case::AccType; + + constexpr index_t M = Case::MPerWave; + constexpr index_t N = Case::NPerWave; + constexpr index_t K = Case::KPerWave; + + HostTensor A({M, K}); + HostTensor B({N, K}); + HostTensor C({M, N}); + + FillUniformDistribution{-1.f, 1.f, 11939}(A); + FillUniformDistribution{-1.f, 1.f, 11940}(B); + C.SetZero(); + + RunWarpGemmCase(A, B, C); + + HostTensor C_ref({M, N}); + C_ref.SetZero(); + reference_gemm(A, B.transpose(), C_ref); + + EXPECT_TRUE(check_err(C, C_ref, "Warp gemm bf16 result error.")); +} From d574cc47577d558c85a119b03f772ee087f7f077 Mon Sep 17 00:00:00 2001 From: "Maksim (Max) Podkorytov" Date: Tue, 2 Jun 2026 23:15:10 +0000 Subject: [PATCH 013/143] [rocm-libraries] ROCm/rocm-libraries#6696 (commit 9627b91) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace nested static_for lambdas with compile-time search helper (#6696) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Add `sequence_find_value` and `find_in_tuple_of_sequences` compile-time search helpers with O(1) template depth - Replace nested `static_for` lambdas in `TensorDescriptor::GetTransformAndItsUpperDimension` and `InitializeElementSize` - Apply same optimizations to `TensorAdaptor` Supersedes #4287. Conflict-resolved rebase of ROCm/composable_kernel#3600 onto current develop. ## Motivation The `TensorDescriptor` and `TensorAdaptor` classes had excessive template instantiation from: 1. Nested `static_for` loops with lambdas creating unique closure types at every call site 2. `generate_tuple` with lambdas causing per-type instantiation overhead The new helpers use constexpr array lookup and pack expansion instead of recursive template patterns, achieving O(1) template depth. ## Results (`example_grouped_conv_fwd_xdl_fp16`, n=10, interleaved, `-j1`, `-ftime-trace`) | TU | Baseline (mean) | New (mean) | Delta | Wilcoxon p | Mann-Whitney p | |----|-----------------|------------|-------|-----------|---------------| | `grouped_conv_fwd_xdl_fp16` (host) | 14,886 ms | 13,353 ms | **-10.3%** | **0.002** | **0.0002** | | `grouped_conv_fwd_xdl_fp16` (device) | 27,762 ms | 25,629 ms | **-7.7%** | **0.002** | **0.0002** | | **Total (all TUs)** | **57,732 ms** | **54,030 ms** | **-6.4%** | | | Unrelated TUs (`device_memory`, `host_tensor`, `convolution_parameter`) show no significant difference (p > 0.3), serving as negative controls. ### Methodology - 10 interleaved runs (baseline₁, new₁, baseline₂, new₂, ...) on the same node to eliminate ordering/warmup bias - Wilcoxon signed-rank test (paired, non-parametric) and Mann-Whitney U test (unpaired) - Built with patched clang (LLVM 22) on ctr2-alola-compile-11, `-j1` for accurate per-TU timing - Raw data available in Slurm job 275230 results ## Test plan - [x] 11 unit tests added (5 for `sequence_find_value`, 6 for `find_in_tuple_of_sequences`) - [x] Compile-time benchmark with statistical significance (p < 0.01) - [ ] Full CI Tracking issue: #4229 --- .../ck/tensor_description/tensor_adaptor.hpp | 59 +++----- .../tensor_description/tensor_descriptor.hpp | 62 ++++----- include/ck/utility/sequence_helper.hpp | 131 ++++++++++++++++++ test/util/unit_sequence_helper.cpp | 84 +++++++++++ 4 files changed, 264 insertions(+), 72 deletions(-) diff --git a/include/ck/tensor_description/tensor_adaptor.hpp b/include/ck/tensor_description/tensor_adaptor.hpp index 0804e07fc3..069f0ec794 100644 --- a/include/ck/tensor_description/tensor_adaptor.hpp +++ b/include/ck/tensor_description/tensor_adaptor.hpp @@ -48,28 +48,29 @@ struct TensorAdaptor return BottomDimensionHiddenIds{}; } - __host__ __device__ static constexpr auto InitializeElementSize(const Transforms& transforms) + // Helper to get length of a top dimension from transforms + template + __host__ __device__ static constexpr auto + GetTopDimLengthFromTransforms(const Transforms& transforms) { - const auto lengths = generate_tuple( - [&](auto idim_top) { - constexpr auto tmp = GetTransformAndItsUpperDimension(idim_top); - - constexpr index_t itran = tmp[Number<0>{}]; - constexpr index_t idim_up = tmp[Number<1>{}]; - constexpr bool found = tmp[Number<2>{}]; - - static_assert(found == true, - "wrong! not found matching transformation and upper-dimension"); - - const auto length = - transforms[Number{}].GetUpperLengths()[Number{}]; + constexpr auto result = find_in_tuple_of_sequences{})>( + UpperDimensionHiddenIdss{}); + static_assert(result.found, "wrong! not found matching transformation and upper-dimension"); + return transforms[Number{}].GetUpperLengths()[Number{}]; + } - return length; - }, - Number{}); + // Compute element size using pack expansion instead of generate_tuple with lambda + template + __host__ __device__ static constexpr auto ComputeElementSizeImpl(const Transforms& transforms, + Sequence) + { + return (GetTopDimLengthFromTransforms(transforms) * ...); + } - // TODO: make container_reduce support tuple of Number and index_t - return container_reduce(lengths, math::multiplies{}, Number<1>{}); + __host__ __device__ static constexpr auto InitializeElementSize(const Transforms& transforms) + { + return ComputeElementSizeImpl(transforms, + typename arithmetic_sequence_gen<0, ndim_top_, 1>::type{}); } template @@ -79,24 +80,10 @@ struct TensorAdaptor constexpr index_t idim_hidden = TopDimensionHiddenIds::At(idim_top); - index_t itran_found = 0; - index_t idim_up_found = 0; - bool found = false; - - static_for<0, ntransform_, 1>{}([&](auto itran) { - constexpr auto up_dim_ids = UpperDimensionHiddenIdss{}[itran]; - - static_for<0, up_dim_ids.Size(), 1>{}([&](auto idim_up) { - if constexpr(up_dim_ids[idim_up] == idim_hidden) - { - itran_found = itran; - idim_up_found = idim_up; - found = true; - } - }); - }); + // Use compile-time search helper instead of nested static_for with lambdas. + constexpr auto result = find_in_tuple_of_sequences(UpperDimensionHiddenIdss{}); - return make_tuple(itran_found, idim_up_found, found); + return make_tuple(result.itran, result.idim_up, result.found); } __host__ __device__ static constexpr index_t GetNumOfBottomDimension() diff --git a/include/ck/tensor_description/tensor_descriptor.hpp b/include/ck/tensor_description/tensor_descriptor.hpp index 3a5d258fae..7b20b299a9 100644 --- a/include/ck/tensor_description/tensor_descriptor.hpp +++ b/include/ck/tensor_description/tensor_descriptor.hpp @@ -53,28 +53,29 @@ struct TensorDescriptor return unique_sort_all_dim_ids::Size(); } - __host__ __device__ static constexpr auto InitializeElementSize(const Transforms& transforms) + // Helper to get length of a visible dimension from transforms + template + __host__ __device__ static constexpr auto + GetVisibleDimLengthFromTransforms(const Transforms& transforms) { - const auto lengths = generate_tuple( - [&](auto idim_visible) { - constexpr auto tmp = GetTransformAndItsUpperDimension(idim_visible); - - constexpr index_t itran = tmp[Number<0>{}]; - constexpr index_t idim_up = tmp[Number<1>{}]; - constexpr bool found = tmp[Number<2>{}]; - - static_assert(found == true, - "wrong! not found matching transformation and upper-dimension"); - - const auto length = - transforms[Number{}].GetUpperLengths()[Number{}]; + constexpr auto result = + find_in_tuple_of_sequences{})>(UpperDimensionIdss{}); + static_assert(result.found, "wrong! not found matching transformation and upper-dimension"); + return transforms[Number{}].GetUpperLengths()[Number{}]; + } - return length; - }, - Number{}); + // Compute element size using pack expansion instead of generate_tuple with lambda + template + __host__ __device__ static constexpr auto ComputeElementSizeImpl(const Transforms& transforms, + Sequence) + { + return (GetVisibleDimLengthFromTransforms(transforms) * ...); + } - // TODO: make container_reduce support tuple of Number and index_t - return container_reduce(lengths, math::multiplies{}, Number<1>{}); + __host__ __device__ static constexpr auto InitializeElementSize(const Transforms& transforms) + { + return ComputeElementSizeImpl( + transforms, typename arithmetic_sequence_gen<0, ndim_visible_, 1>::type{}); } template @@ -84,24 +85,13 @@ struct TensorDescriptor constexpr index_t idim_hidden = VisibleDimensionIds::At(idim_visible); - index_t itran_found = 0; - index_t idim_up_found = 0; - bool found = false; - - static_for<0, ntransform_, 1>{}([&](auto itran) { - constexpr auto up_dim_ids = UpperDimensionIdss{}[itran]; - - static_for<0, up_dim_ids.Size(), 1>{}([&](auto idim_up) { - if constexpr(up_dim_ids[idim_up] == idim_hidden) - { - itran_found = itran; - idim_up_found = idim_up; - found = true; - } - }); - }); + // Use compile-time search helper instead of nested static_for loops. + // This significantly reduces applier::operator() template instantiations + // by replacing nested lambda-based loops with a single constexpr search. + // See sequence_helper.hpp::find_in_tuple_of_sequences for details. + constexpr auto result = find_in_tuple_of_sequences(UpperDimensionIdss{}); - return make_tuple(itran_found, idim_up_found, found); + return make_tuple(result.itran, result.idim_up, result.found); } constexpr static index_t ntransform_ = GetNumOfTransform(); diff --git a/include/ck/utility/sequence_helper.hpp b/include/ck/utility/sequence_helper.hpp index 6f096aef74..fc1dc795d4 100644 --- a/include/ck/utility/sequence_helper.hpp +++ b/include/ck/utility/sequence_helper.hpp @@ -52,4 +52,135 @@ __host__ __device__ constexpr auto unpack_and_merge_sequences(TupleOfSequences t return unpack(merge_sequences_functor{}, tuple_of_sequences); } +// sequence_find_value - O(1) template depth constexpr search +// +// Optimization: Constexpr loop with array lookup instead of recursive template pattern +// +// Why this approach: +// - Recursive template (OLD): template instantiation for each recursion level → O(N) instantiations +// Example: Finding value in Sequence<1,2,3,4,5> requires 5 recursive instantiations +// +// - Constexpr loop (NEW): Single function instantiation with runtime loop → O(1) instantiation +// Same search requires only 1 function instantiation, loop executes at compile-time +// +// Implementation details: +// 1. Pack expansion creates constexpr array: {(Is == Target)...} +// 2. Constexpr for loop searches the array +// 3. Entire function evaluates at compile-time (no runtime cost) +// +// Impact: +// - Significantly reduces template instantiation depth for sequence search operations +// - Dramatically improves compilation time vs recursive template approach +// - Pattern applies to any compile-time search/lookup operation +// +// Trade-off: Uses constexpr evaluation instead of pure template metaprogramming. +// Requires C++14 constexpr but results in dramatically better compile times. +// +template +__host__ __device__ constexpr index_t sequence_find_value(Sequence) +{ + if constexpr(sizeof...(Is) == 0) + { + return -1; + } + else + { + constexpr bool matches[] = {(Is == Target)...}; + for(index_t i = 0; i < static_cast(sizeof...(Is)); ++i) + { + if(matches[i]) + return i; + } + return -1; + } +} + +// Result type for find_in_tuple_of_sequences +template +struct FindTransformResult +{ + static constexpr index_t itran = ITran; + static constexpr index_t idim_up = IDimUp; + static constexpr bool found = Found; +}; + +// find_in_tuple_of_sequences - finds which sequence contains a target value +// +// Optimization: Pack expansion with constexpr search instead of nested static_for loops +// +// Why this approach: +// - Nested static_for (OLD): Creates lambda closure for each iteration level +// Example: Searching Tuple, Seq<2,3>, Seq<4,5>> creates multiple applier::operator() +// instantiations. Result: Many applier instantiations for typical tensor descriptor operations. +// +// - Pack expansion + constexpr (NEW): Single function with compile-time array search +// Example: Same search creates constexpr array, single search function. +// Result: 1 function instantiation regardless of tuple size. +// +// Implementation: +// 1. Pack expansion: sequence_find_value(Seqs{})... applies search to each sequence +// 2. Results collected in constexpr array +// 3. Linear search finds first non-negative result (sequence containing target) +// +// Impact: +// - Significantly reduces applier::operator() instantiations in tensor descriptor transforms +// - O(1) template depth instead of O(N*M) for N sequences of length M +// +// Use case: Finding which dimension index contains a specific value (common in tensor reordering) +// +template +struct FindInTupleOfSequencesCompute +{ + private: + // Result struct for constexpr computation + struct ResultData + { + index_t itran; + index_t idim_up; + bool found; + }; + + // Compute result using constexpr function with array lookup + static constexpr ResultData compute() + { + if constexpr(sizeof...(Seqs) == 0) + { + return {0, 0, false}; + } + else + { + // Pack expansion creates array - O(1) template depth + constexpr index_t indices[] = {sequence_find_value(Seqs{})...}; + + // Find first matching sequence + for(index_t i = 0; i < static_cast(sizeof...(Seqs)); ++i) + { + if(indices[i] >= 0) + { + return {i, indices[i], true}; + } + } + return {0, 0, false}; + } + } + + static constexpr ResultData result_ = compute(); + + public: + static constexpr index_t itran = result_.itran; + static constexpr index_t idim_up = result_.idim_up; + static constexpr bool found = result_.found; + + using type = FindTransformResult; +}; + +// Find target value in a tuple of sequences +// Returns FindTransformResult +// Uses O(1) template depth via pack expansion (no recursion) +template +__host__ __device__ constexpr auto find_in_tuple_of_sequences(Tuple) +{ + return typename FindInTupleOfSequencesCompute::type{}; +} + } // namespace ck diff --git a/test/util/unit_sequence_helper.cpp b/test/util/unit_sequence_helper.cpp index 4f8740f799..a082979b77 100644 --- a/test/util/unit_sequence_helper.cpp +++ b/test/util/unit_sequence_helper.cpp @@ -63,3 +63,87 @@ TEST(UnpackAndMergeSequences, TwoSequences) auto expected = Sequence<100, 200, 300>{}; EXPECT_TRUE((is_same::value)); } + +// Tests for sequence_find_value (PR #3600) +TEST(SequenceFindValue, FindExistingElement) +{ + constexpr auto result = sequence_find_value<17>(Sequence<5, 11, 17, 23, 29>{}); + EXPECT_EQ(result, 2); // 17 is at index 2 +} + +TEST(SequenceFindValue, FindFirstElement) +{ + constexpr auto result = sequence_find_value<7>(Sequence<7, 13, 19, 31>{}); + EXPECT_EQ(result, 0); +} + +TEST(SequenceFindValue, FindLastElement) +{ + constexpr auto result = sequence_find_value<41>(Sequence<3, 11, 23, 41>{}); + EXPECT_EQ(result, 3); +} + +TEST(SequenceFindValue, ElementNotFound) +{ + constexpr auto result = sequence_find_value<50>(Sequence<2, 8, 14, 26>{}); + EXPECT_EQ(result, -1); +} + +TEST(SequenceFindValue, EmptySequence) +{ + constexpr auto result = sequence_find_value<1>(Sequence<>{}); + EXPECT_EQ(result, -1); +} + +// Tests for find_in_tuple_of_sequences (PR #3600) +TEST(FindInTupleOfSequences, FindInFirstSequence) +{ + constexpr auto tuple_of_seqs = + make_tuple(Sequence<5, 11>{}, Sequence<17, 23>{}, Sequence<29, 37>{}); + constexpr auto result = find_in_tuple_of_sequences<11>(tuple_of_seqs); + EXPECT_EQ(result.itran, 0); // Found in first sequence (index 0) + EXPECT_EQ(result.idim_up, 1); // At position 1 within that sequence + EXPECT_TRUE(result.found); +} + +TEST(FindInTupleOfSequences, FindInMiddleSequence) +{ + constexpr auto tuple_of_seqs = + make_tuple(Sequence<2, 4, 6>{}, Sequence<8, 10>{}, Sequence<12>{}); + constexpr auto result = find_in_tuple_of_sequences<10>(tuple_of_seqs); + EXPECT_EQ(result.itran, 1); // Found in second sequence (index 1) + EXPECT_EQ(result.idim_up, 1); // At position 1 within that sequence + EXPECT_TRUE(result.found); +} + +TEST(FindInTupleOfSequences, FindInLastSequence) +{ + constexpr auto tuple_of_seqs = make_tuple(Sequence<3>{}, Sequence<7>{}, Sequence<13, 19, 31>{}); + constexpr auto result = find_in_tuple_of_sequences<31>(tuple_of_seqs); + EXPECT_EQ(result.itran, 2); // Found in third sequence (index 2) + EXPECT_EQ(result.idim_up, 2); // At position 2 within that sequence + EXPECT_TRUE(result.found); +} + +TEST(FindInTupleOfSequences, NotFound) +{ + constexpr auto tuple_of_seqs = make_tuple(Sequence<1, 3>{}, Sequence<5, 7, 9>{}); + constexpr auto result = find_in_tuple_of_sequences<100>(tuple_of_seqs); + EXPECT_FALSE(result.found); +} + +TEST(FindInTupleOfSequences, EmptyTuple) +{ + constexpr auto tuple_of_seqs = make_tuple(); + constexpr auto result = find_in_tuple_of_sequences<1>(tuple_of_seqs); + EXPECT_FALSE(result.found); +} + +TEST(FindInTupleOfSequences, SingleSequence) +{ + constexpr auto tuple_of_seqs = make_tuple(Sequence<41, 43, 47, 53>{}); + constexpr auto result = find_in_tuple_of_sequences<47>(tuple_of_seqs); + EXPECT_EQ(result.itran, 0); + EXPECT_EQ(result.idim_up, 2); + EXPECT_TRUE(result.found); +} From 57205893112d7dddc3f7555ce233f60f4a595015 Mon Sep 17 00:00:00 2001 From: Illia Silin <98187287+illsilin@users.noreply.github.com> Date: Wed, 3 Jun 2026 01:58:59 +0000 Subject: [PATCH 014/143] [rocm-libraries] ROCm/rocm-libraries#7960 (commit ddac5cf) [CK] Upgrade to new gfx1250 compiler and fix build issues (#7960) ## Motivation The docker image we've been using to build for gfx1250 is a few months old, so we need to upgrade. Some of the changes in the latest compiler version require changes in the code. TDM is temporarily disabled due to changes in the lds load/store intrinsics. ## Technical Details ## Test Plan ## Test Result ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- .../core/arch/amd_buffer_addressing.hpp | 49 ++++++--------- .../arch/amd_buffer_addressing_builtins.hpp | 48 ++++++-------- .../ck_tile/core/arch/amd_tdm_descriptor.hpp | 62 +++++++++---------- test/synchronization/monitor_mwait.cpp | 27 ++++---- vars/ck.groovy | 2 +- 5 files changed, 81 insertions(+), 107 deletions(-) diff --git a/include/ck_tile/core/arch/amd_buffer_addressing.hpp b/include/ck_tile/core/arch/amd_buffer_addressing.hpp index 4134b077db..ed1ddd5648 100644 --- a/include/ck_tile/core/arch/amd_buffer_addressing.hpp +++ b/include/ck_tile/core/arch/amd_buffer_addressing.hpp @@ -3073,21 +3073,15 @@ amd_tdm_load(const TDMDescriptor& descriptor static constexpr auto I1 = number<1>{}; static constexpr auto I2 = number<2>{}; static constexpr auto I3 = number<3>{}; - if constexpr(TensorRank == 2 && !IsGatherMode) - { - auto tdm_desc_grp = descriptor.getResourceDescriptorGroup2(); - __builtin_amdgcn_tensor_load_to_lds_d2( - tdm_desc_grp.get(I0), tdm_desc_grp.get(I1), static_cast(coherence)); - } - else - { - auto tdm_desc_grp = descriptor.getResourceDescriptorGroup4(); - __builtin_amdgcn_tensor_load_to_lds(tdm_desc_grp.get(I0), - tdm_desc_grp.get(I1), - tdm_desc_grp.get(I2), - tdm_desc_grp.get(I3), - static_cast(coherence)); - } + static constexpr auto I4 = number<4>{}; + + auto tdm_desc_grp = descriptor.getResourceDescriptorGroup(); + __builtin_amdgcn_tensor_load_to_lds(tdm_desc_grp.get(I0), + tdm_desc_grp.get(I1), + tdm_desc_grp.get(I2), + tdm_desc_grp.get(I3), + tdm_desc_grp.get(I4), + static_cast(coherence)); #else ignore = descriptor; #endif @@ -3105,21 +3099,16 @@ amd_tdm_store(const TDMDescriptor& descripto static constexpr auto I1 = number<1>{}; static constexpr auto I2 = number<2>{}; static constexpr auto I3 = number<3>{}; - if constexpr(TensorRank == 2 && !IsGatherMode) - { - auto tdm_desc_grp = descriptor.getResourceDescriptorGroup2(); - __builtin_amdgcn_tensor_store_from_lds_d2( - tdm_desc_grp.get(I0), tdm_desc_grp.get(I1), static_cast(coherence)); - } - else - { - auto tdm_desc_grp = descriptor.getResourceDescriptorGroup4(); - __builtin_amdgcn_tensor_store_from_lds(tdm_desc_grp.get(I0), - tdm_desc_grp.get(I1), - tdm_desc_grp.get(I2), - tdm_desc_grp.get(I3), - static_cast(coherence)); - } + static constexpr auto I4 = number<4>{}; + + auto tdm_desc_grp = descriptor.getResourceDescriptorGroup(); + __builtin_amdgcn_tensor_store_from_lds(tdm_desc_grp.get(I0), + tdm_desc_grp.get(I1), + tdm_desc_grp.get(I2), + tdm_desc_grp.get(I3), + tdm_desc_grp.get(I4), + static_cast(coherence)); +} #else ignore = descriptor; #endif diff --git a/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp b/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp index de15dacab0..55597b3723 100644 --- a/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp +++ b/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp @@ -3413,21 +3413,15 @@ amd_tdm_load(const TDMDescriptor& descriptor static constexpr auto I1 = number<1>{}; static constexpr auto I2 = number<2>{}; static constexpr auto I3 = number<3>{}; - if constexpr(TensorRank == 2 && !IsGatherMode) - { - auto tdm_desc_grp = descriptor.getResourceDescriptorGroup2(); - __builtin_amdgcn_tensor_load_to_lds_d2( - tdm_desc_grp.get(I0), tdm_desc_grp.get(I1), static_cast(coherence)); - } - else - { - auto tdm_desc_grp = descriptor.getResourceDescriptorGroup4(); - __builtin_amdgcn_tensor_load_to_lds(tdm_desc_grp.get(I0), - tdm_desc_grp.get(I1), - tdm_desc_grp.get(I2), - tdm_desc_grp.get(I3), - static_cast(coherence)); - } + static constexpr auto I4 = number<4>{}; + + auto tdm_desc_grp = descriptor.getResourceDescriptorGroup(); + __builtin_amdgcn_tensor_load_to_lds(tdm_desc_grp.get(I0), + tdm_desc_grp.get(I1), + tdm_desc_grp.get(I2), + tdm_desc_grp.get(I3), + tdm_desc_grp.get(I4), + static_cast(coherence)); #else ignore = descriptor; #endif @@ -3445,21 +3439,15 @@ amd_tdm_store(const TDMDescriptor& descripto static constexpr auto I1 = number<1>{}; static constexpr auto I2 = number<2>{}; static constexpr auto I3 = number<3>{}; - if constexpr(TensorRank == 2 && !IsGatherMode) - { - auto tdm_desc_grp = descriptor.getResourceDescriptorGroup2(); - __builtin_amdgcn_tensor_store_from_lds_d2( - tdm_desc_grp.get(I0), tdm_desc_grp.get(I1), static_cast(coherence)); - } - else - { - auto tdm_desc_grp = descriptor.getResourceDescriptorGroup4(); - __builtin_amdgcn_tensor_store_from_lds(tdm_desc_grp.get(I0), - tdm_desc_grp.get(I1), - tdm_desc_grp.get(I2), - tdm_desc_grp.get(I3), - static_cast(coherence)); - } + static constexpr auto I4 = number<4>{}; + + auto tdm_desc_grp = descriptor.getResourceDescriptorGroup(); + __builtin_amdgcn_tensor_store_from_lds(tdm_desc_grp.get(I0), + tdm_desc_grp.get(I1), + tdm_desc_grp.get(I2), + tdm_desc_grp.get(I3), + tdm_desc_grp.get(I4), + static_cast(coherence)); #else ignore = descriptor; #endif diff --git a/include/ck_tile/core/arch/amd_tdm_descriptor.hpp b/include/ck_tile/core/arch/amd_tdm_descriptor.hpp index c2e927143a..2b6b761f95 100644 --- a/include/ck_tile/core/arch/amd_tdm_descriptor.hpp +++ b/include/ck_tile/core/arch/amd_tdm_descriptor.hpp @@ -411,12 +411,11 @@ class TDMDescriptor CK_TILE_DEVICE TDMGatherIndexSize getTDMGatherIndexSize() const { return m_rowIdxSize; } CK_TILE_DEVICE const void* getRowIndex() const { return m_rowIndex; } - // currently llvm gives two builtins for TDM descriptor - // __builtin_amdgcn_tensor_load_to_lds and __builtin_amdgcn_tensor_load_to_lds_d2 - CK_TILE_DEVICE auto getResourceDescriptorGroup2() const -> tuple + // currently llvm gives unified builtins for TDM descriptor + // __builtin_amdgcn_tensor_load_to_lds (2D uses zero vectors for unused args) + CK_TILE_DEVICE auto getResourceDescriptorGroup() const + -> tuple { - static_assert(TensorRank <= 2, "TensorRank must be less than or equal to 2"); - static_assert(!IsGatherMode, "Gather mode not supported for getResourceDescriptorGroup2"); TDM_GROUP0 group0{reinterpret_cast(m_localAddress), reinterpret_cast(m_globalAddress), static_cast(m_rowIdxSize), @@ -425,34 +424,31 @@ class TDMDescriptor TDM_GROUP1 group1; configureGroup1(group1); - // generate tuples with 2 elements; first is int32x4_t, second is int32x8_t - return make_tuple(amd_wave_read_first_lane(group0.bitfield), - amd_wave_read_first_lane(group1.bitfield)); - } - - CK_TILE_DEVICE auto - getResourceDescriptorGroup4() const -> tuple - { - TDM_GROUP0 group0{reinterpret_cast(m_localAddress), - reinterpret_cast(m_globalAddress), - static_cast(m_rowIdxSize), - IsGatherMode ? 1u : 0u}; - - TDM_GROUP1 group1; - configureGroup1(group1); - - // generate tuples with 4 elements; first is int32x4_t, second is int32x8_t, third is - // int32x4_t, fourth is int32x4_t - TDM_GROUP2 group2; - configureGroup2(group2); - - TDM_GROUP3 group3; - configureGroup3(group3); - - return make_tuple(amd_wave_read_first_lane(group0.bitfield), - amd_wave_read_first_lane(group1.bitfield), - amd_wave_read_first_lane(group2.bitfield), - amd_wave_read_first_lane(group3.bitfield)); + if constexpr(TensorRank <= 2 && !IsGatherMode) + { + int32x4_t v4i_zeros = {}; + int32x8_t v8i_zeros = {}; + return make_tuple(amd_wave_read_first_lane(group0.bitfield), + amd_wave_read_first_lane(group1.bitfield), + v4i_zeros, + v4i_zeros, + v8i_zeros); + } + else + { + TDM_GROUP2 group2; + configureGroup2(group2); + + TDM_GROUP3 group3; + configureGroup3(group3); + + int32x8_t v8i_zeros = {}; + return make_tuple(amd_wave_read_first_lane(group0.bitfield), + amd_wave_read_first_lane(group1.bitfield), + amd_wave_read_first_lane(group2.bitfield), + amd_wave_read_first_lane(group3.bitfield), + v8i_zeros); + } } private: diff --git a/test/synchronization/monitor_mwait.cpp b/test/synchronization/monitor_mwait.cpp index b3490bb08a..f8fba1e571 100644 --- a/test/synchronization/monitor_mwait.cpp +++ b/test/synchronization/monitor_mwait.cpp @@ -13,7 +13,8 @@ using ::ck::DeviceMem; using F8DataType = ck::f8_t; #if defined(__gfx125__) -__device__ constexpr int hint_and_scope = 2 << 3; // temporal + Device +__device__ constexpr int hint = __ATOMIC_RELAXED; +__device__ constexpr int scope = __MEMORY_SCOPE_SYSTEM; // BUG: duration = 0x8000 (sleep-forever) should not be used as the wave might never wake up if the // s_monitor_sleep(duration) is called when MWAIT=0 __device__ constexpr short duration = static_cast(1 << 15) - 1; // forever - 1 clock cycle @@ -27,8 +28,8 @@ __global__ void gpu_ping(F8DataType* ptr, const int Num, int* runNum, bool ck_lo ptr[0] = F8DataType{0x38}; while(run++ < Num && !ck::fp8_is_nan(ptr[0]) && !ck::fp8_is_nan(ptr[1])) { - while((__builtin_amdgcn_flat_load_monitor_b32(static_cast(static_cast(ptr)), - hint_and_scope) & + while((__builtin_amdgcn_flat_load_monitor_b32( + static_cast(static_cast(ptr)), hint, scope) & 0xFF) == 0x38) { __builtin_amdgcn_s_monitor_sleep(duration); @@ -63,8 +64,8 @@ __global__ void gpu_pong(F8DataType* ptr, const int Num, int* runNum, bool ck_lo int run = 0; while(run++ < Num && !ck::fp8_is_nan(ptr[0]) && !ck::fp8_is_nan(ptr[1])) { - while((__builtin_amdgcn_flat_load_monitor_b32(static_cast(static_cast(ptr)), - hint_and_scope) & + while((__builtin_amdgcn_flat_load_monitor_b32( + static_cast(static_cast(ptr)), hint, scope) & 0xFF) == 0) { // Wait for the ping thread to set the value to 0x38 @@ -198,15 +199,15 @@ __global__ void gpu_ping(int* ptrA, while(run++ < Num) { - while(__builtin_amdgcn_flat_load_monitor_b128(reinterpret_cast(ptrA), - hint_and_scope)[tid] != expectedA0) + while(__builtin_amdgcn_flat_load_monitor_b128( + reinterpret_cast(ptrA), hint, scope)[tid] != expectedA0) { __builtin_amdgcn_s_monitor_sleep(duration); } ptrB[tid] = toUpdateB1; - while(__builtin_amdgcn_flat_load_monitor_b128(reinterpret_cast(ptrA), - hint_and_scope)[tid] != expectedA1) + while(__builtin_amdgcn_flat_load_monitor_b128( + reinterpret_cast(ptrA), hint, scope)[tid] != expectedA1) { __builtin_amdgcn_s_monitor_sleep(duration); } @@ -237,15 +238,15 @@ __global__ void gpu_pong(int* ptrB, int run = 0; while(run++ < Num) { - while(__builtin_amdgcn_flat_load_monitor_b128(reinterpret_cast(ptrB), - hint_and_scope)[tid] != expectedB0) + while(__builtin_amdgcn_flat_load_monitor_b128( + reinterpret_cast(ptrB), hint, scope)[tid] != expectedB0) { __builtin_amdgcn_s_monitor_sleep(duration); } ptrA[tid] = toUpdateA0; - while(__builtin_amdgcn_flat_load_monitor_b128(reinterpret_cast(ptrB), - hint_and_scope)[tid] != expectedB1) + while(__builtin_amdgcn_flat_load_monitor_b128( + reinterpret_cast(ptrB), hint, scope)[tid] != expectedB1) { __builtin_amdgcn_s_monitor_sleep(duration); } diff --git a/vars/ck.groovy b/vars/ck.groovy index e2caf8007f..c71564d898 100644 --- a/vars/ck.groovy +++ b/vars/ck.groovy @@ -1374,7 +1374,7 @@ def runBuildCKAndTests(String arch) { case "gfx1250": gpuTarget = "gfx1250" extraSetupArgs = " -DDISABLE_DL_KERNELS=\"ON\"" - extraBuildArgs = [docker_name: "${env.CK_DOCKERHUB_PRIVATE}:npi-mi450-latest", no_reboot: true] + extraBuildArgs = [docker_name: "${env.CK_DOCKERHUB_PRIVATE}:ck_ub24.04_gfx1250", no_reboot: true] break case "gfx10-1-generic": case "gfx10-3-generic": From 01bd52bdb5fdf55f4047cce5e4abc1ebbf302f5f Mon Sep 17 00:00:00 2001 From: Yi DING <28386673+DDEle@users.noreply.github.com> Date: Wed, 3 Jun 2026 02:09:05 +0000 Subject: [PATCH 015/143] [rocm-libraries] ROCm/rocm-libraries#7925 (commit a8f0845) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [CK] Fix gfx950 AITER Sync Regressions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Fixes three gfx950 regressions in the AITER downstream CI that surfaced after the internal/gfx1250 re-sync (ROCm/rocm-libraries#6978): > **Companion aiter PR:** ROCm/aiter#3392 — host-side adaptations (`Kernel::BlockSize()` `constexpr` drops, blockscale `KBatch=1` clamp) plus the CK submodule bump used to validate these fixes together. - **FlyDSL MoE AOT cache miss** — the AITER MoE tests run with `check_aot_cache=True` and fail on any FlyDSL JIT cache miss, but the CI never pre-compiles the FlyDSL MoE kernels, so gfx950 always misses. Pre-compile them at the start of the AITER test stage. - **`buffer.load.lds.v4i32` link error** — ROCm/rocm-libraries#6978 reintroduced a clang-version guard mapping `llvm.amdgcn.raw.buffer.load.lds` to a `.v4i32`-suffixed name. That name exists in no LLVM (the rsrc operand is a fixed, non-overloaded `<4 x i32>`, so the intrinsic is never type-mangled), so gfx950 4-DWORD direct-to-LDS (e.g. fp4 MoE bpreshuffle) fails to link with `lld: undefined symbol: llvm.amdgcn.raw.buffer.load.lds.v4i32`. Use the canonical plain name unconditionally. - **mixed-precision flatmm warp-GEMM call** — ROCm/rocm-libraries#6978 generalized the scaled `WarpGemmImpl::operator()` from a fixed `` signature to a variadic `` one and updated the `mx_flatmm` pipeline to pass the op-selectors as `OpSelA<>`/`OpSelB<>` types, but missed the mixed-precision flatmm pipeline (`F8xMXF4`/`F16xMXF4`), which still passed raw integer op-selectors. These no longer bind to `typename... Params` (`error: no matching member function for call to 'operator()'`), breaking compilation of the fp8/bf16 × fp4 cktile MoE gemm1 instances on gfx950 (aiter `test_moe_2stage`). Wrap the op-selectors in `OpSelA<>`/`OpSelB<>`. ## Changes - `Jenkinsfile`: pre-compile the FlyDSL MoE AOT cache (`python3 aiter/aot/flydsl/moe.py`) before the AITER tests. - `include/ck/utility/amd_buffer_addressing_builtins.hpp` and `include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp`: drop the `__clang_major__` guard and always use `__asm("llvm.amdgcn.raw.buffer.load.lds")`. The plain name is the canonical one for all sizes including the gfx950 16-byte form, as the upstream LLVM gfx950 tests confirm. - `include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp`: wrap the warp-GEMM op-selectors in `OpSelA<>`/`OpSelB<>` at the five call sites, matching the `mx_flatmm` pipeline. ## Test plan Validated via CI. --- .../ck/utility/amd_buffer_addressing_builtins.hpp | 11 ----------- .../core/arch/amd_buffer_addressing_builtins.hpp | 11 ----------- ...ixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp | 12 +++++++----- vars/ck.groovy | 2 ++ 4 files changed, 9 insertions(+), 27 deletions(-) diff --git a/include/ck/utility/amd_buffer_addressing_builtins.hpp b/include/ck/utility/amd_buffer_addressing_builtins.hpp index e35986177a..a8e2cc4881 100644 --- a/include/ck/utility/amd_buffer_addressing_builtins.hpp +++ b/include/ck/utility/amd_buffer_addressing_builtins.hpp @@ -830,16 +830,6 @@ amd_buffer_atomic_max(const typename vector_type_maker::type::type src_thr } // Direct loads from global to LDS. -#if __clang_major__ >= 21 && __clang_major__ < 23 -__device__ void -llvm_amdgcn_raw_buffer_load_lds(int32x4_t rsrc, - __attribute__((address_space(3))) uint32_t* lds_ptr, - index_t size, - index_t voffset, - index_t soffset, - index_t offset, - index_t aux) __asm("llvm.amdgcn.raw.buffer.load.lds.v4i32"); -#else __device__ void llvm_amdgcn_raw_buffer_load_lds(int32x4_t rsrc, __attribute__((address_space(3))) uint32_t* lds_ptr, @@ -848,7 +838,6 @@ llvm_amdgcn_raw_buffer_load_lds(int32x4_t rsrc, index_t soffset, index_t offset, index_t aux) __asm("llvm.amdgcn.raw.buffer.load.lds"); -#endif #ifndef __HIPCC_RTC__ template diff --git a/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp b/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp index 55597b3723..448394dd43 100644 --- a/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp +++ b/include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp @@ -1381,16 +1381,6 @@ CK_TILE_DEVICE_EXTERN double llvm_amdgcn_raw_buffer_atomic_max_fp64( int glc_slc) __asm("llvm.amdgcn.raw.buffer.atomic.fmax.f64.v4i32"); // Direct loads from global to LDS. -#if __clang_major__ >= 21 && __clang_major__ < 23 -CK_TILE_DEVICE_EXTERN void -llvm_amdgcn_raw_buffer_load_lds(int32x4_t rsrc, - as3_uint32_ptr lds_ptr, - index_t size, - index_t voffset, - index_t soffset, - index_t offset, - index_t aux) __asm("llvm.amdgcn.raw.buffer.load.lds.v4i32"); -#else CK_TILE_DEVICE_EXTERN void llvm_amdgcn_raw_buffer_load_lds(int32x4_t rsrc, as3_uint32_ptr lds_ptr, @@ -1399,7 +1389,6 @@ llvm_amdgcn_raw_buffer_load_lds(int32x4_t rsrc, index_t soffset, index_t offset, index_t aux) __asm("llvm.amdgcn.raw.buffer.load.lds"); -#endif template CK_TILE_DEVICE void async_buffer_load_dwordxn_v(void* smem, diff --git a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp index eb1df36ea6..f75cd25bd6 100644 --- a/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp @@ -1982,7 +1982,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 // warp GEMM WG{}.template // operator()( - operator()( + operator(), OpSelB>( c_warp_tensor, a_warp_tensor(number{}), b_warp_tensor_ping(nIter_pack * number{} + @@ -2092,7 +2092,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 // warp GEMM WG{}.template // operator()( - operator()( + operator(), OpSelB>( c_warp_tensor, a_warp_tensor(number{}), b_warp_tensor_pong(nIter_pack * number{} + @@ -2214,7 +2214,8 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); // warp GEMM - WG{}.template operator()( + WG{}.template + operator(), OpSelB>( c_warp_tensor, a_warp_tensor(number{}), b_warp_tensor_ping(nIter_pack * number{} + @@ -2283,7 +2284,8 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); // warp GEMM - WG{}.template operator()( + WG{}.template + operator(), OpSelB>( // operator()( c_warp_tensor, a_warp_tensor(number{}), @@ -2346,7 +2348,7 @@ struct F8xMXF4FlatmmPipelineAGmemBGmemCRegV1 // warp GEMM WG{}.template // operator()( - operator()( + operator(), OpSelB>( c_warp_tensor, a_warp_tensor(number{}), b_warp_tensor_ping(nIter_pack * number{} + diff --git a/vars/ck.groovy b/vars/ck.groovy index c71564d898..c97cfd81e3 100644 --- a/vars/ck.groovy +++ b/vars/ck.groovy @@ -1180,6 +1180,8 @@ def getPytorchTestsCmds() { } def getAiterTestsCmds() { return [ + // Pre-compile FlyDSL MoE AOT cache before the tests. + "cd /home/jenkins/workspace/aiter && python3 aiter/aot/flydsl/moe.py", "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8.py", "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8_blockscale.py", "python3 /home/jenkins/workspace/aiter/op_tests/test_mha.py", From 7ecbf82708905697c7764b89e9656de9eaa9aae6 Mon Sep 17 00:00:00 2001 From: Anton Gorenko Date: Wed, 3 Jun 2026 06:16:10 +0000 Subject: [PATCH 016/143] [rocm-libraries] ROCm/rocm-libraries#7500 (commit f5cd4fd) [CK_TILE][FMHA] Optimize long-context decoding on gfx11/12 (#7500) ## Motivation Relevant issue: ROCM-22065 FMHA has less-than-optimal performance of long-context decoding (i.e. when seqlen_q = 1) on gfx11/12. This PR optimizes the splitkv pipeline and configs for such scenarios. ## Technical Details Optimizations applied in this PR: 1. use tiles with smaller M0 (16 vs 64), these tiles are used when seqlen_q <= 16 2. adapt qr_nwarp_sshuffle pipeline for gfx11, it allows to use more warps even for M0 = 16 (the qr pipeline parallelizes work between warps in M dim so with M0 = 16 it allows to use only 1 warp) 3. enable kMergeNumHeadGroupsSeqLenQ (an optimization that merges one group of heads in GQA) for all hdim values, not only 128 4. increase the number of splits (multiply by the number of head groups) if (3) is used 5. increase the number of splits for RDNAs (`multiProcessorCount` is the number of WGPs on RDNAs, not CUs, so it should be doubled to have meaning similar to CDNAs) Performance on gfx1151: | Case | develop (GB/s) | This PR (GB/s) | |:-------|-------:|-------:| | [fp16\|group\|bshd] b:1, h:32/32, s:1/45056, d:64/64 | 127.58 | 183.11 | | [fp16\|group\|bhsd] b:1, h:32/32, s:1/45056, d:64/64 | 153.64 | 215.02 | | [fp16\|group\|bshd] b:1, h:16/8, s:1/77184, d:128/128 | 120.51 | 225.76 | | [fp16\|group\|bhsd] b:1, h:16/8, s:1/77184, d:128/128 | 130.62 | 223.84 | | [fp16\|group\|bshd] b:1, h:32/32, s:1/9600, d:128/128 | 82.65 | 138.44 | | [fp16\|group\|bhsd] b:1, h:32/32, s:1/9600, d:128/128 | 105.75 | 220.45 | | [fp16\|group\|bshd] b:1, h:8/1, s:1/401024, d:256/256 | 16.27 | 187.89 | | [fp16\|group\|bhsd] b:1, h:8/1, s:1/401024, d:256/256 | 16.28 | 188.19 | ## Test Plan An additional test case is added to the exiting test. It uses seqlen_q = 1, GQA, no mask to trigger the changes ``` ninja test_ck_tile_fmha_fwd_fp16 && bin/test_ck_tile_fmha_fwd_fp16 --gtest_filter="*SplitKV* ninja test_ck_tile_fmha_fwd_bf16 && bin/test_ck_tile_fmha_fwd_bf16 --gtest_filter="*SplitKV* ``` Manual testing can be done with these commands: ``` bin/tile_example_fmha_fwd -prec=fp16 -mode=1 -page_block_size=128 -b=1 -h=32 -h_k=32 -d=64 -s=1 -s_k=$((352 * 128)) -lse=1 -mask=0 -num_splits=0 -kname=1 -v=1 bin/tile_example_fmha_fwd -prec=fp16 -mode=1 -page_block_size=128 -b=1 -h=16 -h_k=8 -d=128 -s=1 -s_k=$((603 * 128)) -lse=1 -mask=0 -num_splits=0 -kname=1 -v=1 bin/tile_example_fmha_fwd -prec=fp16 -mode=1 -page_block_size=128 -b=1 -h=32 -h_k=32 -d=128 -s=1 -s_k=$((75 * 128)) -lse=1 -mask=0 -num_splits=0 -kname=1 -v=1 bin/tile_example_fmha_fwd -prec=fp16 -mode=1 -page_block_size=128 -b=1 -h=8 -h_k=1 -d=256 -s=1 -s_k=$((3133 * 128)) -lse=1 -mask=0 -num_splits=0 -kname=1 -v=1 ``` ## Test Result All the tests must pass. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- .../01_fmha/codegen/ops/fmha_fwd_splitkv.py | 55 ++++++++++++++----- example/ck_tile/01_fmha/fmha_fwd_runner.hpp | 16 +++++- ...litkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp | 10 +++- ...nwarp_sshuffle_qr_ks_vs_default_policy.hpp | 23 +++++++- test/ck_tile/fmha/test_fmha_fwd.cpp | 1 + 5 files changed, 84 insertions(+), 21 deletions(-) diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py index 849f463afa..ed025dcf5f 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py @@ -128,7 +128,7 @@ namespace {{ template void run_instance(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a) {{ - if constexpr ({F_hdim} == 128 && {F_bias} == ck_tile::BlockAttentionBiasEnum::NO_BIAS + if constexpr ({F_bias} == ck_tile::BlockAttentionBiasEnum::NO_BIAS && (std::is_same_v<{F_mask}, ck_tile::SimplifiedGenericAttentionMask> || std::is_same_v<{F_mask}, FmhaMasks::NoMask>)) {{ if (a.max_seqlen_q == 1 && a.nhead_k < a.nhead_q) {{ @@ -283,7 +283,7 @@ """ FMHA_FWD_SPLITKV_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.do_fp8_static_quant == {F_squant}) && - ((a.block_table_ptr != nullptr) == {F_pagedkv}) && (t.has_sink == {F_sink}) && ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ + ((a.block_table_ptr != nullptr) == {F_pagedkv}) && (t.has_sink == {F_sink}) && ({F_scheck}) && ({F_seqtune}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{ using traits_ = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, true, {F_squant}, {F_pagedkv},{F_sink}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>; // get combine kernel tile sizes @@ -364,6 +364,14 @@ def scheck(self) -> str: else: assert False + def seqtune(self, max_bm0: int) -> str: + if self.bm0 == max_bm0: + return "true/*fall back to largest tile*/" + else: + if self.mode == "group": + return f"a.max_seqlen_q <= {self.bm0}" + return f"a.seqlen_q <= {self.bm0}" + @property def skcheck(self) -> str: if self.mode == "group": @@ -561,6 +569,7 @@ def api(self) -> str: for i_dtype, (dtype, pool_by_dtype) in enumerate(pool_by_arch.items()): per_hdim_case = str() for i_hdim, (hdim, pool_by_hdim) in enumerate(pool_by_dtype.items()): + max_bm0 = max((t.bm0 for t in pool_by_hdim), default=0) inners = str() for i_trait, trait in enumerate(pool_by_hdim): inners += FMHA_FWD_SPLITKV_API_INNER_DISPATCH.format( @@ -579,6 +588,7 @@ def api(self) -> str: F_pagedkv=BOOL_MAP[trait.pagedkv], F_sink=BOOL_MAP[trait.sink], F_scheck=trait.scheck, + F_seqtune=trait.seqtune(max_bm0), F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck, @@ -763,6 +773,7 @@ def get_pipelines(dtype, hdim, mask_impl) -> List[FmhaFwdSplitKVPipeline]: pipelines.append(Pipeline("qr", "row", "t", "f", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip pipelines.append(Pipeline("qr", "row", "t", "t", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip pipelines.append(Pipeline("qr", "row", "t", "t", "t", "t", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip + pipelines.append(Pipeline("qr_nwarp_sshuffle", "row", "t", "t", "f", "f", logits, bias, "t", squant, pagedkv, sink, mask)) # fmt: skip elif dtype in ["fp8", "bf8"]: for logits, mask, bias in itertools.product( ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys() @@ -846,11 +857,15 @@ class KernelComponentFactoryGfx11(KernelComponentFactoryBase): def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: if dtype in ["fp16", "bf16"]: return { - # bm0, bn0, bk0, bn1, bk1, - "32" : FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "64" : FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "128": FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "256": FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + # bm0, bn0, bk0, bn1, bk1, + "32" : [FmhaFwdTileSize( 16, 64, 16, 32, 32, 32, 1, 2, 1, 1, 2, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "64" : [FmhaFwdTileSize( 16, 64, 32, 64, 32, 64, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "128": [FmhaFwdTileSize( 16, 64, 32, 128, 32, 128, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "256": [FmhaFwdTileSize( 16, 64, 32, 256, 32, 256, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], } # fmt: skip else: return None @@ -863,11 +878,15 @@ class KernelComponentFactoryGfx12(KernelComponentFactoryBase): def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: if dtype in ["fp16", "bf16"]: return { - # bm0, bn0, bk0, bn1, bk1, - "32" : FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "64" : FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "128": FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), - "256": FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1), + # bm0, bn0, bk0, bn1, bk1, + "32" : [FmhaFwdTileSize( 16, 64, 16, 32, 32, 32, 1, 2, 1, 1, 2, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "64" : [FmhaFwdTileSize( 16, 64, 32, 64, 32, 64, 1, 4, 1, 1, 4, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "128": [FmhaFwdTileSize( 16, 128, 32, 128, 32, 128, 1, 8, 1, 1, 8, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], + "256": [FmhaFwdTileSize( 16, 128, 32, 256, 32, 256, 1, 8, 1, 1, 8, 1, 16, 16, 16, 16, 16, 16, -1), + FmhaFwdTileSize( 64, 64, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 16, 16, 16, 16, 16, 16, -1)], } # fmt: skip elif dtype in ["fp8", "bf8"]: return { @@ -930,11 +949,17 @@ def get_fwd_splitkv_blobs( d = factory.get_hdim_tile_size_dict(dtype) if d is None: continue - # for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]): for hdim_str, mode in itertools.product(d.keys(), MODE_MAP.keys()): - tile = d[hdim_str] + tiles = d[hdim_str] + if not isinstance(tiles, list): + tiles = [tiles] hdim = int(hdim_str) - for pipeline in factory.get_pipelines(dtype, hdim, mask_impl): + for tile, pipeline in itertools.product( + tiles, factory.get_pipelines(dtype, hdim, mask_impl) + ): + # Use qr_nwarp_sshuffle with multiple N warps and qr otherwise + if (tile.F_rn0 != 1) != (pipeline.tag == "qr_nwarp_sshuffle"): + continue if mode == "group": if pipeline.F_spad != "t" or pipeline.F_skpad != "t": # in group mode, spad/skpad must be true, since we can't predict if seqlen of current batch need pad or not diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index 0b51dffa46..243ff87faa 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -165,8 +165,10 @@ int override_num_splits_if_necessary( if(num_splits < 1 && p_drop == 0.0f) { + // props.multiProcessorCount for >=gfx10 is the number of WGPs (each has 2 CUs) + const int num_blocks_per_SM = props.warpSize == 32 ? 4 : 2; return num_splits_heuristic( - batch * nhead * num_m_blocks, props.multiProcessorCount * 2, 128); + batch * nhead * num_m_blocks, props.multiProcessorCount * num_blocks_per_SM, 128); } return num_splits; @@ -648,8 +650,18 @@ fwd_result fmha_fwd_run(mode_enum mode, // legalize num_splits according to other options if(num_splits < 1) { + int nhead_merged = nhead; + int max_seqlen_q_merged = max_seqlen_q; + // When max_seqlen_q == 1 and multiple head groups are merged (kMergeNumHeadGroupsSeqLenQ) + // then more splits are required + if(bias.type == bias_enum::no_bias && mask.type == mask_enum::no_mask && + max_seqlen_q == 1 && nhead_k < nhead) + { + nhead_merged = nhead_k; + max_seqlen_q_merged = max_seqlen_q * (nhead / nhead_k); + } num_splits = override_num_splits_if_necessary( - batch, nhead, max_seqlen_q, hdim_v, p_drop, num_splits); + batch, nhead_merged, max_seqlen_q_merged, hdim_v, p_drop, num_splits); } if(128 < num_splits) { diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp index adc8ea5a90..bdc598f754 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp @@ -6,6 +6,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp" +#include "ck_tile/ops/gemm/warp/warp_wmma_gemm_gfx11_utils.hpp" #include "ck_tile/ops/reduce/block/block_reduce.hpp" namespace ck_tile { @@ -257,7 +258,7 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS clear_tile(o_acc); if((__builtin_isinf_sign(sink_v) >= 0) && i_split == 0) { - set_tile(m, SMPLComputeDataType{sink_v * C_LOG2E}); + set_tile(m, SMPLComputeDataType{sink_v * static_cast(C_LOG2E)}); set_tile(l, SMPLComputeDataType{1.0f}); } else @@ -698,8 +699,15 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS block_tile_reduce_sync(rowsum_p, f_sum, bool_constant{}); +#if defined(__gfx11__) + auto p = make_static_distributed_tensor( + decltype(gemm_1)::template MakeABlockTileDistribution()); + PermuteWarpGemmCToA( + p, cast_tile(tile_elementwise_in(p_compute_element_func, p_compute))); +#else const auto p = cast_tile(tile_elementwise_in(p_compute_element_func, p_compute)); +#endif // l{j}, Oacc{j} constexpr auto o_spans = decltype(o_acc)::get_distributed_spans(); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp index c5af751cd5..316720ac22 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs_default_policy.hpp @@ -5,8 +5,6 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qx_ks_vs_custom_policy.hpp" -#include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_custom_policy.hpp" -#include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1.hpp" namespace ck_tile { @@ -163,6 +161,25 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVSDefaultPolicy constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; constexpr index_t kTileK = Problem::BlockFmhaShape::kN0; +#if defined(__gfx11__) + // Keep C distribution and replicate data for NWarp to prevent doubling registers + // PermuteWarpGemmCToA will convert C distribution to A for matrix P later + constexpr index_t K1 = kKPerBlock / WG::kM; + constexpr index_t K0 = kTileK / kKPerBlock; + constexpr index_t M1 = MWarp; + constexpr index_t M0 = kMPerBlock / WG::kN; + + constexpr auto s2_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, sequence>, + tuple>, + tuple>, + sequence<1, 2, 2>, + sequence<0, 0, 1>>{}; + + constexpr auto s2_block_dstr_encoding = detail::make_embed_tile_distribution_encoding( + s2_block_outer_dstr_encoding, typename WG::CWarpDstrEncoding{}); +#else // K2 is equal to Impl::kABKPerLane * kKIterPerWarpGemm constexpr index_t K3 = WG::kK / WG::WarpGemmAttribute::Impl::kABKLane; constexpr index_t K2 = WG::WarpGemmAttribute::Impl::kABKLane; @@ -179,7 +196,7 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVSDefaultPolicy tuple, sequence<2, 2>>, sequence<1, 2, 2, 2>, sequence<0, 0, 1, 3>>{}; - +#endif constexpr auto s2_block_dstr = make_static_tile_distribution(s2_block_dstr_encoding); return s2_block_dstr; diff --git a/test/ck_tile/fmha/test_fmha_fwd.cpp b/test/ck_tile/fmha/test_fmha_fwd.cpp index 6ae33da30f..bdfc2d17c4 100644 --- a/test/ck_tile/fmha/test_fmha_fwd.cpp +++ b/test/ck_tile/fmha/test_fmha_fwd.cpp @@ -735,6 +735,7 @@ INSTANTIATE_TEST_SUITE_P(TestCkTileFmhaFwd, Values(3, 4), Values(std::tuple{4, 3, 1, 200, 1024, "0"}, std::tuple{2, 2, -1, 512, 2000, "0"}, + std::tuple{2, 8, 2, 1, 1024, "0"}, std::tuple{3, 2, -1, 230, 899, "t:128,128"}))); TEST_P(SplitKV, DataTypeConfig) From 88f8d24c344b12f091bcf4585e74071f511e2abf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ville=20Pietil=C3=A4?= <188998872+vpietila-amd@users.noreply.github.com> Date: Wed, 3 Jun 2026 08:40:03 +0000 Subject: [PATCH 017/143] [rocm-libraries] ROCm/rocm-libraries#7936 (commit 3dc91e6) [CK Tile] Fix V6 pipeline applicability and split-image initialization (#7936) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation After adding code generation via CK Tile Dispatcher, some fwd and bwd weight tests for CK Tile convolutions are failing. This PR introduced correct applicability checks and fixes the split-image parameter initialization such that non-applicable instances are not invoked during test execution and split-image instances are correctly initialized. ## Technical Details Investigation revealed two distinct problems 1. For bwd weight, the compute V3 uses prefetch of 3 distinct tiles, which works incorrectly when the number of K-slices addressed by the workgroup is 1. This occurs when a large split-K value is used for a problem that results in a small Gemm-K value. 2. For fwd direction, the current CK Profiler/test infrastructure doesn't initialize the split-image parameters for instance where split-image is enable. Uninitialized split-image values result in non-deterministic behavior where the tests might randomly fail. Fixed problem 1. by adding a check in `IsSupportedArgument` that marks the instance invalid if the `num_loops = ceil(GemmK / (k_batch * KPerBlock)) < 4` for V6 pipeline kernel instances. The check is compile-time eliminated for other kernels. Fixed problem 2. by adding initialization of split-image parameters when split-image is enabled. The default initialization corresponds to full image with no split, i.e., the number of splits is 1 and it has the size of the full image. Added unit tests for the added logic. ## Test Plan Running the following test suites cover the logic added in this PR - test_grouped_convnd_fwd_tile - test_ck_tile_grouped_conv_fwd - test_grouped_convnd_bwd_weight_tile - test_ck_tile_grouped_conv_bwd_weight All test suites above are included in the automated test runs. ## Test Result ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- ...ped_convolution_backward_weight_kernel.hpp | 38 +- .../grouped_convolution_forward_kernel.hpp | 42 +- test/ck_tile/grouped_conv/CMakeLists.txt | 1 + .../test_ck_tile_grouped_conv_bwd_weight.cpp | 138 +++++- .../test_ck_tile_grouped_conv_fwd.cpp | 400 ++++++++++++++++++ .../test_grouped_convnd_bwd_weight_tile.cpp | 2 +- 6 files changed, 613 insertions(+), 8 deletions(-) create mode 100644 test/ck_tile/grouped_conv/test_ck_tile_grouped_conv_fwd.cpp diff --git a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp index 9d031a989e..f9db3412e7 100644 --- a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp +++ b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp @@ -35,6 +35,15 @@ struct is_streamk_partitioner> : std::true_t { }; +template +struct is_compute_v6_pipeline : std::false_type +{ +}; +template +struct is_compute_v6_pipeline> : std::true_type +{ +}; + template CK_TILE_HOST void LogInfo(Args&&... args) noexcept { @@ -468,8 +477,9 @@ struct GroupedConvolutionBackwardWeightKernel using DsDataType = remove_cvref_t; using WeiDataType = remove_cvref_t; - static constexpr bool IsSplitKSupported = true; - static constexpr bool IsStreamK = is_streamk_partitioner::value; + static constexpr bool IsSplitKSupported = true; + static constexpr bool IsStreamK = is_streamk_partitioner::value; + static constexpr bool IsComputeV6Pipeline = is_compute_v6_pipeline::value; using GroupedConvBwdWeightKernelArgsSpecialized = std::conditional_t= PrefetchStages + 1 = 4 + // Otherwise it produces incorrect results (num_loop=1) or is just inefficient (num_loop=2 + // or 3). + if constexpr(IsComputeV6Pipeline) + { + const index_t num_loop = + integer_divide_ceil(kargs.GemmK, kargs.k_batch * TilePartitioner::KPerBlock); + constexpr int num_loop_threashold = GemmPipeline_::PrefetchStages + 1; + if(num_loop < num_loop_threashold) + { + LogInfo("For V6 pipeline, GemmK / (k_batch * KPerBlock) must be >= ", + num_loop_threashold, + ". Now GemmK is ", + kargs.GemmK, + ", k_batch is ", + kargs.k_batch, + ", KPerBlock is ", + number{}, + ", num_loop is ", + num_loop); + return false; + } + } + if constexpr(!std::is_same_v && !std::is_same_v) { diff --git a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp index 14c4356aa1..2e3edeea59 100644 --- a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp +++ b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp @@ -438,12 +438,15 @@ struct GroupedConvFwdKernelArgs index_t num_d_pieces = 1, num_h_pieces = 1, num_w_pieces = 1; // Split factors // Minimal per-piece data (only unique values) + // Default-initialized to 0 so that uninitialized pieces are detectable + // (the invoker sets these after MakeKernelArgs). struct PieceInfo { - index_t block_start; // Starting block index for this piece - index_t block_end; // Ending block index (exclusive) - index_t d_start, h_start, w_start; // Piece starting position in OUTPUT space - index_t d_size, h_size, w_size; // Piece size in OUTPUT space + index_t block_start = -1; // Starting block index for this piece + index_t block_end = -1; // Ending block index (exclusive) + index_t d_start = -1, h_start = -1, + w_start = -1; // Piece starting position in OUTPUT space + index_t d_size = -1, h_size = -1, w_size = -1; // Piece size in OUTPUT space }; static constexpr index_t MaxPieces = 64; // Max pieces: 4 (1D), 16 (2D), 64 (3D) @@ -767,6 +770,37 @@ struct GroupedConvolutionForwardKernel MakeKernelArgs(const GroupedConvFwdHostArgs& hostArgs) { auto kargs = GroupedConvFwdKernelArgsSpecialized(hostArgs); + + // Initialize split-image with a single piece covering the entire output. + // The invoker may later override this with multi-piece data for large + // tensors. Without this default, the split-image kernel path would use + // uninitialized piece data and produce wrong results. + if constexpr(EnableSplitImage) + { + constexpr index_t ndim = GroupedConvTraitsType_::NDimSpatial; + constexpr index_t off = GroupedConvFwdKernelArgsSpecialized::NonSpatialDims; + + const index_t total_w = kargs.out_g_n_k_wos_lengths[off + ndim - 1]; + const index_t total_h = (ndim >= 2) ? kargs.out_g_n_k_wos_lengths[off + ndim - 2] : 1; + const index_t total_d = (ndim >= 3) ? kargs.out_g_n_k_wos_lengths[off + ndim - 3] : 1; + + kargs.split_image.total_d = total_d; + kargs.split_image.total_h = total_h; + kargs.split_image.total_w = total_w; + kargs.split_image.total_spatial = total_d * total_h * total_w; + + kargs.num_spatial_pieces = 1; + kargs.split_image.pieces[0].block_start = 0; + kargs.split_image.pieces[0].block_end = + TilePartitioner::GridSize(kargs.GemmM, kargs.GemmN); + kargs.split_image.pieces[0].d_start = 0; + kargs.split_image.pieces[0].h_start = 0; + kargs.split_image.pieces[0].w_start = 0; + kargs.split_image.pieces[0].d_size = total_d; + kargs.split_image.pieces[0].h_size = total_h; + kargs.split_image.pieces[0].w_size = total_w; + } + return kargs; } diff --git a/test/ck_tile/grouped_conv/CMakeLists.txt b/test/ck_tile/grouped_conv/CMakeLists.txt index b1c1d77205..375e270f2d 100644 --- a/test/ck_tile/grouped_conv/CMakeLists.txt +++ b/test/ck_tile/grouped_conv/CMakeLists.txt @@ -4,6 +4,7 @@ # Currently ck_tile is only built on gfx9 if(GPU_TARGETS MATCHES "gfx9|gfx11|gfx12") add_gtest_executable(test_ck_tile_grouped_conv_bwd_weight test_ck_tile_grouped_conv_bwd_weight.cpp) + add_gtest_executable(test_ck_tile_grouped_conv_fwd test_ck_tile_grouped_conv_fwd.cpp) endif() # StreamK requires cross-CU coherence via StreamKCoherency, which only has diff --git a/test/ck_tile/grouped_conv/test_ck_tile_grouped_conv_bwd_weight.cpp b/test/ck_tile/grouped_conv/test_ck_tile_grouped_conv_bwd_weight.cpp index 7f37ddc6f7..4f420b557b 100644 --- a/test/ck_tile/grouped_conv/test_ck_tile_grouped_conv_bwd_weight.cpp +++ b/test/ck_tile/grouped_conv/test_ck_tile_grouped_conv_bwd_weight.cpp @@ -36,6 +36,46 @@ struct TestConvConfig static constexpr auto Scheduler = GemmPipelineScheduler::Intrawave; }; +// ============================================================================ +// V6 pipeline: num_loop check +// +// GemmPipelineAgBgCrCompV6 uses a 3-stage prefetch (PrefetchStages=3). The +// hot loop only executes when num_loop > 3 (i.e. num_loop >= 4). When +// num_loop == 1 the Odd-tail branch unconditionally processes all 3 prefetch +// buffers, including two that contain K-data from neighbouring workgroups' +// K-partitions, causing an ~3x over-count in the accumulator. The kernel +// must therefore reject configurations where: +// +// num_loop = ceil(GemmK / (k_batch * KPerBlock)) < 4 +// +// In the 2D bwd-weight tests below, GemmK = N * Ho * Wo. +// The V6 config tile has KPerBlock = 32. +// ============================================================================ +struct TestConvConfigV6 +{ + static constexpr index_t VectorSizeA = 4; + static constexpr index_t VectorSizeB = 8; + static constexpr index_t VectorSizeC = 8; + + static constexpr index_t M_Tile = 256; + static constexpr index_t N_Tile = 256; + static constexpr index_t K_Tile = 32; // KPerBlock = 32 + + static constexpr index_t M_Warp = 2; + static constexpr index_t N_Warp = 2; + static constexpr index_t K_Warp = 1; + + static constexpr index_t M_Warp_Tile = 32; + static constexpr index_t N_Warp_Tile = 32; + static constexpr index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr GemmPipeline Pipeline = GemmPipeline::COMPUTE_V6; + static constexpr index_t NumWaveGroups = 1; + static constexpr index_t NumGroupsToMerge = 1; + static constexpr auto Scheduler = GemmPipelineScheduler::Intrawave; +}; + // Helper to build full kernel type template ; - using GemmPipeline = GemmPipelineAgBgCrCompV3; + using GemmPipeline = std::conditional_t, + GemmPipelineAgBgCrCompV3>; using EpilogueProblem = CShuffleEpilogueProblem accepted. +// With k_batch=2: num_loop = ceil(98/64) = 2 -> rejected (< 4). +// With k_batch=3: num_loop = ceil(98/96) = 2 -> rejected. +// With k_batch=4: num_loop = ceil(98/128)= 1 -> rejected. + +TEST_F(GroupedConvBwdWeightV6PipelineTest, AcceptsWhenNumLoopAtLeast4) +{ + using Kernel = typename BuildKernel::type; + + // k_batch=1: num_loop=ceil(98/32)=4. Exactly on the boundary -> must pass. + auto host_args = create_2d_host_args(1); + auto kargs = typename Kernel::GroupedConvBwdWeightKernelArgsSpecialized(host_args); + EXPECT_TRUE(Kernel::IsSupportedArgument(kargs)) + << "V6 kernel must accept k_batch=1 (num_loop=4 >= 4)"; +} + +TEST_F(GroupedConvBwdWeightV6PipelineTest, RejectsWhenNumLoopIs2) +{ + using Kernel = typename BuildKernel::type; + + // k_batch=2: num_loop=ceil(98/64)=2 < 4. Must be rejected. + auto host_args = create_2d_host_args(2); + auto kargs = typename Kernel::GroupedConvBwdWeightKernelArgsSpecialized(host_args); + EXPECT_FALSE(Kernel::IsSupportedArgument(kargs)) + << "V6 kernel must reject k_batch=2 (num_loop=2 < 4) to avoid " + "incorrect Odd-tail prefetch over-read"; +} + +TEST_F(GroupedConvBwdWeightV6PipelineTest, RejectsWhenNumLoopIs1) +{ + using Kernel = typename BuildKernel::type; + + // k_batch=4: num_loop=ceil(98/128)=1 < 4. Must be rejected. + auto host_args = create_2d_host_args(4); + auto kargs = typename Kernel::GroupedConvBwdWeightKernelArgsSpecialized(host_args); + EXPECT_FALSE(Kernel::IsSupportedArgument(kargs)) + << "V6 kernel must reject k_batch=4 (num_loop=1 < 4)"; +} + +TEST_F(GroupedConvBwdWeightV6PipelineTest, AcceptsLargeSpatialWithSmallKBatch) +{ + using Kernel = typename BuildKernel::type; + + // Large spatial: N=2, Hi=Wi=70 -> Ho=Wo=70, GemmK=2*70*70=9800. + // k_batch=1: num_loop=ceil(9800/32)=307 >= 4 -> accepted. + auto host_args = create_large_2d_host_args(1); + auto kargs = typename Kernel::GroupedConvBwdWeightKernelArgsSpecialized(host_args); + EXPECT_TRUE(Kernel::IsSupportedArgument(kargs)) + << "V6 kernel must accept large spatial (num_loop=307) with k_batch=1"; +} + +TEST_F(GroupedConvBwdWeightV6PipelineTest, RejectsLargeSpatialWithLargeKBatch) +{ + using Kernel = typename BuildKernel::type; + + // GemmK=9800. With k_batch=64: num_loop=ceil(9800/(64*32))=ceil(9800/2048)=5 >= 4 -> pass. + // With k_batch=128: num_loop=ceil(9800/4096)=3 < 4 -> rejected. + auto host_args_pass = create_large_2d_host_args(64); + auto kargs_pass = typename Kernel::GroupedConvBwdWeightKernelArgsSpecialized(host_args_pass); + EXPECT_TRUE(Kernel::IsSupportedArgument(kargs_pass)) + << "V6 kernel must accept k_batch=64 (num_loop=5 >= 4)"; + + auto host_args_fail = create_large_2d_host_args(128); + auto kargs_fail = typename Kernel::GroupedConvBwdWeightKernelArgsSpecialized(host_args_fail); + EXPECT_FALSE(Kernel::IsSupportedArgument(kargs_fail)) + << "V6 kernel must reject k_batch=128 (num_loop=3 < 4)"; +} diff --git a/test/ck_tile/grouped_conv/test_ck_tile_grouped_conv_fwd.cpp b/test/ck_tile/grouped_conv/test_ck_tile_grouped_conv_fwd.cpp new file mode 100644 index 0000000000..d9ad9559bf --- /dev/null +++ b/test/ck_tile/grouped_conv/test_ck_tile_grouped_conv_fwd.cpp @@ -0,0 +1,400 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "gtest/gtest.h" +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp" +#include "ck_tile/ops/grouped_convolution/pipeline/grouped_conv_universal_pipeline_ag_bg_cr_policy.hpp" +#include "ck_tile/ops/grouped_convolution/utils/grouped_convolution_utils.hpp" + +using namespace ck_tile; + +// ============================================================================ +// Minimal conv config with unit vector sizes (VectorSizeA/B/C = 1) and +// BASIC_V1 pipeline, mirroring the depthwise-like instances that appear in +// the tile dispatcher and triggered the split-image bug. +// ============================================================================ +struct TestFwdConvConfigBasicV1UnitVec +{ + static constexpr index_t VectorSizeA = 1; + static constexpr index_t VectorSizeB = 1; + static constexpr index_t VectorSizeC = 1; + + static constexpr index_t M_Tile = 16; + static constexpr index_t N_Tile = 64; + static constexpr index_t K_Tile = 64; + + static constexpr index_t M_Warp = 1; + static constexpr index_t N_Warp = 4; + static constexpr index_t K_Warp = 1; + + static constexpr index_t M_Warp_Tile = 16; + static constexpr index_t N_Warp_Tile = 16; + static constexpr index_t K_Warp_Tile = 32; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr GemmPipeline Pipeline = GemmPipeline::BASIC_V1; + static constexpr auto Scheduler = GemmPipelineScheduler::Intrawave; + static constexpr index_t NumWaveGroups = 1; + static constexpr index_t NumGroupsToMerge = 1; +}; + +// Standard config with larger vector sizes (passes all checks for normal problem sizes) +struct TestFwdConvConfigStandard +{ + static constexpr index_t VectorSizeA = 4; + static constexpr index_t VectorSizeB = 8; + static constexpr index_t VectorSizeC = 8; + + static constexpr index_t M_Tile = 128; + static constexpr index_t N_Tile = 128; + static constexpr index_t K_Tile = 32; + + static constexpr index_t M_Warp = 2; + static constexpr index_t N_Warp = 2; + static constexpr index_t K_Warp = 1; + + static constexpr index_t M_Warp_Tile = 16; + static constexpr index_t N_Warp_Tile = 16; + static constexpr index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr GemmPipeline Pipeline = GemmPipeline::BASIC_V1; + static constexpr auto Scheduler = GemmPipelineScheduler::Intrawave; + static constexpr index_t NumWaveGroups = 1; + static constexpr index_t NumGroupsToMerge = 1; +}; + +// ============================================================================ +// Helper to assemble the full forward kernel type. +// EnableSplitImage_ corresponds to the compile-time flag on GroupedConvTraits. +// ============================================================================ +template +struct BuildFwdKernel +{ + using GemmShape = TileGemmShape< + sequence, + sequence, + sequence>; + + using ConvTraits = GroupedConvTraits, + OutLayout, + ConvConfig::VectorSizeA, + ConvConfig::VectorSizeB, + ConvConfig::VectorSizeC, + ConvConfig::NumGroupsToMerge, + EnableSplitImage>; + + using TilePartitioner = + GemmSpatiallyLocalTilePartitioner; + + using GemmUniversalTraits = + TileGemmUniversalTraits; + + using UniversalGemmProblem = + UniversalGemmPipelineProblem; + + using GemmPipeline = GemmPipelineAGmemBGmemCRegV1; + + using EpilogueProblem = CShuffleEpilogueProblem, + float, + PrecType, + typename ConvTraits::ImplicitGemmDsLayout, + typename ConvTraits::FixedGemmParams::ELayout, + element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + ConvConfig::M_Warp, + ConvConfig::N_Warp, + ConvConfig::M_Warp_Tile, + ConvConfig::N_Warp_Tile, + ConvConfig::K_Warp_Tile, + ConvTraits::FixedGemmParams::TransposeC, + ConvConfig::NumWaveGroups, + ConvTraits::FixedGemmParams::FixedVectorSize, + ConvTraits::VectorSizeC>; + + using Epilogue = CShuffleEpilogue; + + using type = + GroupedConvolutionForwardKernel; +}; + +// ============================================================================ +// Helper to create 2D forward host args (null device pointers, host-only). +// ============================================================================ +static GroupedConvFwdHostArgs<> create_2d_fwd_host_args(index_t G, + index_t N, + index_t K, + index_t C, + index_t Y, + index_t X, + index_t Hi, + index_t Wi, + index_t stride_y = 1, + index_t stride_x = 1, + index_t dilation_y = 1, + index_t dilation_x = 1, + index_t lpad_y = 0, + index_t lpad_x = 0, + index_t rpad_y = 0, + index_t rpad_x = 0, + index_t k_batch = 1) +{ + auto conv_param = conv::ConvParam{2, + G, + N, + K, + C, + {Y, X}, + {Hi, Wi}, + {stride_y, stride_x}, + {dilation_y, dilation_x}, + {lpad_y, lpad_x}, + {rpad_y, rpad_x}}; + + return GroupedConvFwdHostArgs<>{conv_param, nullptr, nullptr, {}, nullptr, k_batch}; +} + +// ============================================================================ +// Tests +// ============================================================================ + +class GroupedConvFwdIsSupportedArgumentTest : public ::testing::Test +{ +}; + +// --------------------------------------------------------------------------- +// Split-image default initialization in MakeKernelArgs (full-image path) +// --------------------------------------------------------------------------- +// MakeKernelArgs() initializes split_image.pieces[0] as a single piece covering +// the entire output. This ensures the split-image kernel path works correctly +// even without the large-tensor invoker. The invoker can override with +// multi-piece data for large tensors. + +// MakeKernelArgs initializes pieces so split-image is accepted for any valid problem. +TEST_F(GroupedConvFwdIsSupportedArgumentTest, SplitImageFullImageAfterMakeKernelArgs) +{ + using Kernel = typename BuildFwdKernel::type; + + // K=64, C=64 — MakeKernelArgs should set up a single-piece split_image. + // 3x3 filter, stride 1, no padding => output H/W = 5x5. + auto host_args = create_2d_fwd_host_args(1, 2, 64, 64, 3, 3, 7, 7); + auto kargs = Kernel::MakeKernelArgs(host_args); + + // Verify piece[0] was initialized by MakeKernelArgs (full-image, single piece). + EXPECT_EQ(kargs.num_spatial_pieces, 1); + EXPECT_EQ(kargs.split_image.pieces[0].block_start, 0); + EXPECT_GT(kargs.split_image.pieces[0].block_end, 0) + << "MakeKernelArgs must initialize split_image.pieces[0]"; + EXPECT_EQ(kargs.split_image.pieces[0].h_start, 0); + EXPECT_EQ(kargs.split_image.pieces[0].w_start, 0); + EXPECT_EQ(kargs.split_image.pieces[0].h_size, 5) << "Output H = (7 - 3)/1 + 1 = 5"; + EXPECT_EQ(kargs.split_image.pieces[0].w_size, 5) << "Output W = (7 - 3)/1 + 1 = 5"; + + // Unused pieces retain the sentinel default (-1) from PieceInfo. + EXPECT_EQ(kargs.split_image.pieces[1].block_start, -1) + << "Unused pieces must retain sentinel default"; + EXPECT_EQ(kargs.split_image.pieces[1].block_end, -1) + << "Unused pieces must retain sentinel default"; + + EXPECT_TRUE(Kernel::IsSupportedArgument(kargs)) + << "Split-image instance must be accepted after MakeKernelArgs initializes pieces"; +} + +// Split-image with depthwise-like K=1 problem is also accepted (pieces are initialized). +TEST_F(GroupedConvFwdIsSupportedArgumentTest, SplitImageFullImageSmallK) +{ + using Kernel = typename BuildFwdKernel::type; + + // K=1, C=1 + auto host_args = create_2d_fwd_host_args(1, 2, 1, 1, 3, 3, 7, 7); + auto kargs = Kernel::MakeKernelArgs(host_args); + + EXPECT_EQ(kargs.num_spatial_pieces, 1); + EXPECT_EQ(kargs.split_image.pieces[0].block_start, 0); + EXPECT_GT(kargs.split_image.pieces[0].block_end, 0) + << "MakeKernelArgs must initialize pieces even for small K"; +} + +// Large K is accepted with proper initialization. +TEST_F(GroupedConvFwdIsSupportedArgumentTest, SplitImageFullImageLargeK) +{ + using Kernel = typename BuildFwdKernel::type; + + // K=96 — the case that caused flaky failures + // 1x1 filter, stride 1, no padding => output H/W = 73x128. + auto host_args = create_2d_fwd_host_args(3, 5, 96, 200, 1, 1, 73, 128); + auto kargs = Kernel::MakeKernelArgs(host_args); + + EXPECT_EQ(kargs.num_spatial_pieces, 1); + EXPECT_EQ(kargs.split_image.pieces[0].block_start, 0); + EXPECT_GT(kargs.split_image.pieces[0].block_end, 0) + << "MakeKernelArgs must initialize pieces for K=96"; + EXPECT_EQ(kargs.split_image.pieces[0].h_size, 73); + EXPECT_EQ(kargs.split_image.pieces[0].w_size, 128); +} + +// --------------------------------------------------------------------------- +// Multi-piece split-image path (large-tensor invoker) +// --------------------------------------------------------------------------- +// The large-tensor invoker (grouped_convolution_forward_large_tensor_invoker.hpp in the examples +// code) calls MakeKernelArgs() first, then overrides split_image.pieces[] with multi-piece data +// computed by calculate_spatial_piece(). + +TEST_F(GroupedConvFwdIsSupportedArgumentTest, SplitImageMultiPieceInvokerOverride) +{ + using Build = BuildFwdKernel; + using Kernel = typename Build::type; + using TilePartitioner = typename Build::TilePartitioner; + + // Large problem: G=1, N=4, K=64, C=64, filter=1x1, input=128x128 + // Output = 128x128. Split H into 2 pieces: H=[0..64) and H=[64..128). + const index_t G = 1, N = 4, K = 64, C = 64; + const index_t Hi = 128, Wi = 128; + const index_t Ho = 128, Wo = 128; // 1x1 filter, stride 1, no padding + + auto host_args = create_2d_fwd_host_args(G, N, K, C, 1, 1, Hi, Wi); + auto kargs = Kernel::MakeKernelArgs(host_args); + + // Verify MakeKernelArgs set up the default single-piece first. + ASSERT_EQ(kargs.num_spatial_pieces, 1); + ASSERT_GT(kargs.split_image.pieces[0].block_end, 0); + + // Now simulate what the invoker does: split H into 2 pieces. + const index_t num_h_pieces = 2; + const index_t num_w_pieces = 1; + const index_t num_d_pieces = 1; + const index_t total_pieces = num_d_pieces * num_h_pieces * num_w_pieces; + const index_t base_piece_h = Ho / num_h_pieces; // 64 + const index_t base_piece_w = Wo; // 128 + const index_t base_piece_d = 1; + + index_t total_blocks = 0; + std::array temp_pieces{}; + for(index_t piece = 0; piece < total_pieces; piece++) + { + temp_pieces[piece] = calculate_spatial_piece(piece, + num_d_pieces, + num_h_pieces, + num_w_pieces, + base_piece_d, + base_piece_h, + base_piece_w, + 1, // total_d + Ho, + Wo, + N, + K, + total_blocks); + total_blocks = temp_pieces[piece].block_end; + } + + // Override kargs with multi-piece data. + kargs.num_spatial_pieces = total_pieces; + kargs.split_image.num_h_pieces = num_h_pieces; + kargs.split_image.num_w_pieces = num_w_pieces; + kargs.split_image.num_d_pieces = num_d_pieces; + for(index_t i = 0; i < total_pieces; i++) + { + kargs.split_image.pieces[i] = {temp_pieces[i].block_start, + temp_pieces[i].block_end, + temp_pieces[i].d_start, + temp_pieces[i].h_start, + temp_pieces[i].w_start, + temp_pieces[i].d_size, + temp_pieces[i].h_size, + temp_pieces[i].w_size}; + } + + // Verify piece 0: covers H=[0..64), W=[0..128) + EXPECT_EQ(kargs.split_image.pieces[0].block_start, 0); + EXPECT_GT(kargs.split_image.pieces[0].block_end, 0); + EXPECT_EQ(kargs.split_image.pieces[0].h_start, 0); + EXPECT_EQ(kargs.split_image.pieces[0].h_size, 64); + EXPECT_EQ(kargs.split_image.pieces[0].w_start, 0); + EXPECT_EQ(kargs.split_image.pieces[0].w_size, 128); + + // Verify piece 1: covers H=[64..128), W=[0..128) + EXPECT_EQ(kargs.split_image.pieces[1].block_start, kargs.split_image.pieces[0].block_end); + EXPECT_GT(kargs.split_image.pieces[1].block_end, kargs.split_image.pieces[1].block_start); + EXPECT_EQ(kargs.split_image.pieces[1].h_start, 64); + EXPECT_EQ(kargs.split_image.pieces[1].h_size, 64); + EXPECT_EQ(kargs.split_image.pieces[1].w_start, 0); + EXPECT_EQ(kargs.split_image.pieces[1].w_size, 128); + + // Pieces must be contiguous: piece1.block_start == piece0.block_end + EXPECT_EQ(kargs.split_image.pieces[1].block_start, kargs.split_image.pieces[0].block_end); + + // Total blocks across pieces must equal the full grid size. + const index_t full_grid = TilePartitioner::GridSize(kargs.GemmM, kargs.GemmN); + EXPECT_EQ(kargs.split_image.pieces[total_pieces - 1].block_end, full_grid); + + EXPECT_EQ(kargs.num_spatial_pieces, 2); + EXPECT_TRUE(Kernel::IsSupportedArgument(kargs)) + << "Split-image instance must be accepted after invoker populates multi-piece data"; +} diff --git a/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_tile.cpp b/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_tile.cpp index 9cf60d7146..41a6eee79a 100644 --- a/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_tile.cpp +++ b/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_tile.cpp @@ -48,7 +48,7 @@ class TestGroupedConvndBwdWeightTile : public ::testing::Test .output = {.config = {.layout = SignatureDetailsType::out_layout}}}; std::vector> conv_args; - std::vector split_ks{"-1", "1", "2"}; + std::vector split_ks{"-1", "1", "2", "64"}; template void Run() From db05d611368c8ce935005ead4a7a404390a19afa Mon Sep 17 00:00:00 2001 From: chris-tsiaousis-hpc <253485634+chris-tsiaousis-hpc@users.noreply.github.com> Date: Wed, 3 Jun 2026 14:35:18 +0000 Subject: [PATCH 018/143] [rocm-libraries] ROCm/rocm-libraries#6212 (commit ccee58d) =?UTF-8?q?[CK=20TILE]=20Unification=20Work=20=E2=80=93=20?= =?UTF-8?q?More=20accurate=20tests=20for=20MmaPipelines=20(#6212)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation This PR solves several issues: #### More accurate tests for MmaPipelines The current tests for the MmaPipelines (test_amdgcn_sparse_mma, test_amdgcn_wavewise_mma) use explicit input fragment vectors filled with 1s, and only check the output of a single lane. We should have tests that actually use the MmaPipelines with non-trivial input matrices and verify the complete output. Some other aspects of the current MmaPipelines tests that I noticed and deserve some attention: 1. There is sometimes iteration over K outside of the pipeline, which is then included in WaveTileK or FragK, which is not correct. We should remove it, move K iteration inside of the pipeline, or be more clear about this outer-K loop size and how it propagates downwards. 2. There is very tight coupling between the kernel, gtest code, and test_pipeline helper, requiring a lot of information and functions to be passed back and forth. 3. The test_pipeline helper is doing a bunch of register-related logic on the host (related to point 1) 4. Without this register logic the only thing it does is check the device, call the kernel, and check the output, but with a lot of boilerplate. #### Test helper for detecting target arch at HOST runtime There is a really apparent issue we faced while writing tests: Scenario: 1. Compile a test that supports both gfx950 and gfx1201 for gfx950 2. Run the test on a server that only has gfx1201 GPU Actual: Segmentation fault Expected: The test can correctly detect from HOST runtime that the DEVICE target_id was different and skips the test. Notes: The only way of detecting the COMPILER_TARGET_ID in the existing "arch" framework is launching a kernel and calling `get_compiler_target()` (so, from a DEVICE code). This will create a segmentation fault if the current arch differs from the target arch. To cope with this issue, we propose to export the compiler target(s) (note they can be many) through `projects/composablekernel/test/ck_tile/core/arch/CMakeLists.txt` and define a test helper to deal with such cases. #### Add composition support to Transforms We have a small number of Transforms which act on MmaOp input and output data, before and after the MmaOp call respectively. These are currently implemented to work on an MmaTile level, but in theory they are also supposed to work at a WaveTile level, i.e. after composition of multiple MmaTiles to create larger effective MNK dimensions. Currently the composed MmaTiles look like 2D C-style arrays of the individual MmaTile level register vectors (see WaveWiseMmaPipeline). The transforms should be able to take these and perform the proper transforms to the whole WaveTile at once. This might allow for better performing transformations. Note: This PR handles the SparseTransform case and if we don't end up doing scale as a transformation, there isn't really much left to do. If we end up having only the sparse transform as a non-trivial transform, then we could also consider removing the Transform framework. --- .../example_tile_distr_enc_calc.cpp | 2 +- include/ck_tile/core.hpp | 1 + include/ck_tile/core/arch/mma/amdgcn_mma.hpp | 7 - .../core/arch/mma/mfma/mfma_selector.hpp | 1 + .../core/arch/mma/mfma/mfma_transforms.hpp | 1 + include/ck_tile/core/arch/mma/mma.hpp | 8 + .../ck_tile/core/arch/mma/mma_pipeline.hpp | 39 +- .../ck_tile/core/arch/mma/mma_wavewise.hpp | 2 +- .../arch/mma/scale/scale_mma_pipeline.hpp | 101 +++- .../arch/mma/sparse/sparse_mma_pipeline.hpp | 170 +++++- .../arch/mma/sparse/sparse_transforms.hpp | 99 ++- test/ck_tile/core/arch/mma/CMakeLists.txt | 116 +++- .../arch/mma/get_cmake_targets_helper.hpp | 87 +++ .../core/arch/mma/get_wave_size_helper.hpp | 34 -- .../mma/pipeline/pipeline_tests_helper.hpp | 564 ++++++++++++------ .../mma/pipeline/test_amdgcn_mma_pipeline.cpp | 2 +- .../mma/pipeline/test_amdgcn_scale_mma.cpp | 397 +++++++++--- .../mma/pipeline/test_amdgcn_sparse_mma.cpp | 377 ++++++++---- .../mma/pipeline/test_amdgcn_wavewise_mma.cpp | 180 ++++-- .../ck_tile/core/arch/mma/test_amdgcn_mma.cpp | 9 +- 20 files changed, 1627 insertions(+), 570 deletions(-) create mode 100644 include/ck_tile/core/arch/mma/mma.hpp create mode 100644 test/ck_tile/core/arch/mma/get_cmake_targets_helper.hpp delete mode 100644 test/ck_tile/core/arch/mma/get_wave_size_helper.hpp diff --git a/example/ck_tile/51_tile_distr_enc_reg_map/example_tile_distr_enc_calc.cpp b/example/ck_tile/51_tile_distr_enc_reg_map/example_tile_distr_enc_calc.cpp index a491c0d2b9..9e62f6e939 100644 --- a/example/ck_tile/51_tile_distr_enc_reg_map/example_tile_distr_enc_calc.cpp +++ b/example/ck_tile/51_tile_distr_enc_reg_map/example_tile_distr_enc_calc.cpp @@ -5,7 +5,7 @@ #include #include #include "ck_tile/core/arch/arch.hpp" -#include "ck_tile/core/arch/mma/amdgcn_mma.hpp" +#include "ck_tile/core/arch/mma/mma.hpp" #include "ck_tile/core/arch/mma/utility/tile_distribution_encoding_register_mapper.hpp" #include "ck_tile/core/arch/mma/utility/tile_distribution_encoding_calculator.hpp" #include "ck_tile/core/container/tuple.hpp" diff --git a/include/ck_tile/core.hpp b/include/ck_tile/core.hpp index 4afba77d6a..2b7066cabf 100644 --- a/include/ck_tile/core.hpp +++ b/include/ck_tile/core.hpp @@ -24,6 +24,7 @@ #include "ck_tile/core/arch/mma/mfma/mfma_selector.hpp" #include "ck_tile/core/arch/mma/mfma/mfma_traits.hpp" #include "ck_tile/core/arch/mma/mfma/mfma_transforms.hpp" +#include "ck_tile/core/arch/mma/mma.hpp" #include "ck_tile/core/arch/mma/mma_op_family.hpp" #include "ck_tile/core/arch/mma/mma_pipeline.hpp" #include "ck_tile/core/arch/mma/mma_selector.hpp" diff --git a/include/ck_tile/core/arch/mma/amdgcn_mma.hpp b/include/ck_tile/core/arch/mma/amdgcn_mma.hpp index 5985f63440..4cb28762ba 100644 --- a/include/ck_tile/core/arch/mma/amdgcn_mma.hpp +++ b/include/ck_tile/core/arch/mma/amdgcn_mma.hpp @@ -393,10 +393,3 @@ CK_TILE_HOST_DEVICE void print(amdgcn_mma= 23 #pragma clang diagnostic pop #endif - -// Include the implementations -#include "wmma/wmma.hpp" // should be included before the below headers - -#include "mfma/mfma.hpp" -#include "scale/scale.hpp" -#include "sparse/sparse.hpp" diff --git a/include/ck_tile/core/arch/mma/mfma/mfma_selector.hpp b/include/ck_tile/core/arch/mma/mfma/mfma_selector.hpp index 2140e3317a..0fa1bada78 100644 --- a/include/ck_tile/core/arch/mma/mfma/mfma_selector.hpp +++ b/include/ck_tile/core/arch/mma/mfma/mfma_selector.hpp @@ -6,6 +6,7 @@ #include "ck_tile/core/config.hpp" #include "ck_tile/core/arch/arch.hpp" #include "ck_tile/core/arch/mma/amdgcn_mma.hpp" +#include "ck_tile/core/arch/mma/mma_selector.hpp" #include "ck_tile/core/arch/mma/mma_traits.hpp" #include "ck_tile/core/numeric/vector_type.hpp" diff --git a/include/ck_tile/core/arch/mma/mfma/mfma_transforms.hpp b/include/ck_tile/core/arch/mma/mfma/mfma_transforms.hpp index 5a3fc9a7e4..9609ed5116 100644 --- a/include/ck_tile/core/arch/mma/mfma/mfma_transforms.hpp +++ b/include/ck_tile/core/arch/mma/mfma/mfma_transforms.hpp @@ -4,6 +4,7 @@ #pragma once #include "ck_tile/core/arch/arch.hpp" +#include "ck_tile/core/arch/mma/mma_transforms.hpp" namespace ck_tile::core::arch::mma { diff --git a/include/ck_tile/core/arch/mma/mma.hpp b/include/ck_tile/core/arch/mma/mma.hpp new file mode 100644 index 0000000000..ec38fe78e3 --- /dev/null +++ b/include/ck_tile/core/arch/mma/mma.hpp @@ -0,0 +1,8 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "wmma/wmma.hpp" +#include "mfma/mfma.hpp" +#include "sparse/sparse.hpp" diff --git a/include/ck_tile/core/arch/mma/mma_pipeline.hpp b/include/ck_tile/core/arch/mma/mma_pipeline.hpp index de01760620..0994f2d178 100644 --- a/include/ck_tile/core/arch/mma/mma_pipeline.hpp +++ b/include/ck_tile/core/arch/mma/mma_pipeline.hpp @@ -270,12 +270,20 @@ struct MmaPipelineBase { if constexpr(MmaOpTraits::IsSupported) { - auto transformed_inputs = applyTransformsToInputs( - hasFlag() ? std::forward(b) - : std::forward(a), - hasFlag() ? std::forward(a) - : std::forward(b), - std::forward(accum)); + constexpr bool swap_a_and_b = hasFlag(); + + auto transformed_inputs = [&]() { + if constexpr(swap_a_and_b) + { + return applyTransformsToInputs( + std::forward(b), std::forward(a), std::forward(accum)); + } + else + { + return applyTransformsToInputs( + std::forward(a), std::forward(b), std::forward(accum)); + } + }(); Derived::execImpl(transformed_inputs); @@ -302,17 +310,18 @@ struct MmaPipelineBase { if constexpr(MmaOpTraits::IsSupported) { - // TODO: c++20: Call template functions with MmaPipelineOptionFlags directly + static_assert(MmaOpTraits::IsScale, + "This exec variant is intended for scale policy structs"); + constexpr bool swap_a_and_b = hasFlag(); + auto transformed_inputs = applyTransformsToInputs( - hasFlag() ? std::forward(b) - : std::forward(a), - hasFlag() ? std::forward(a) - : std::forward(b), + swap_a_and_b ? std::forward(b) : std::forward(a), + swap_a_and_b ? std::forward(a) : std::forward(b), std::forward(accum), - hasFlag() ? std::forward(scale_B) - : std::forward(scale_A), - hasFlag() ? std::forward(scale_A) - : std::forward(scale_B)); + swap_a_and_b ? std::forward(scale_B) + : std::forward(scale_A), + swap_a_and_b ? std::forward(scale_A) + : std::forward(scale_B)); Derived::execImpl(transformed_inputs); diff --git a/include/ck_tile/core/arch/mma/mma_wavewise.hpp b/include/ck_tile/core/arch/mma/mma_wavewise.hpp index 9fbbab411e..bc7e383f6d 100644 --- a/include/ck_tile/core/arch/mma/mma_wavewise.hpp +++ b/include/ck_tile/core/arch/mma/mma_wavewise.hpp @@ -169,7 +169,7 @@ struct WaveWiseMmaPipeline : public MmaPipelineBase::SelectedOp, typename MmaTransforms = // TODO: c++20 MmaTransformsI MmaTransforms = typename MmaTransformsDefaultSelector::SelectedTransforms> // clang-format off -struct ScaleMmaPipeline : public MmaPipelineBase(MmaPipelineOptionFlag::NONE), ScaleMmaPipeline> +struct ScaleMmaPipeline : public MmaPipelineBase(MmaPipelineOptionFlag::NONE), ScaleMmaPipeline> { - using Base = MmaPipelineBase(MmaPipelineOptionFlag::NONE), ScaleMmaPipeline>; + using Base = MmaPipelineBase(MmaPipelineOptionFlag::NONE), ScaleMmaPipeline>; // clang-format on using MmaOp = MmaOp_; // Expose the selected MmaOp - // Expose caller-side vector types - using AVecType = typename MmaOp::AVecType; - using BVecType = typename MmaOp::BVecType; - using CVecType = typename MmaOp::CVecType; + // Fragment dimensions (from the hardware MmaOp) + constexpr static uint32_t FragM = MmaOp::kM; + constexpr static uint32_t FragN = MmaOp::kN; + constexpr static uint32_t FragK = MmaOp::kK; - // Expose internal vector types + // Fragment counts for decomposition + constexpr static uint32_t FragsM = WaveTileM / FragM; + constexpr static uint32_t FragsN = WaveTileN / FragN; + constexpr static uint32_t FragsK = WaveTileK / FragK; + + // Vector types for packed registers in each fragment using InternalAVecT = typename MmaOp::AVecType; using InternalBVecT = typename MmaOp::BVecType; using InternalCVecT = typename MmaOp::CVecType; + // Buffer types for WaveTiles + using AVecType = InternalAVecT[FragsM][FragsK]; + using BVecType = InternalBVecT[FragsN][FragsK]; + using CVecType = InternalCVecT[FragsM][FragsN]; + // Transforms using ATransform = typename MmaTransforms::ATransform; using BTransform = typename MmaTransforms::BTransform; using CTransform = typename MmaTransforms::CTransform; using DTransform = typename MmaTransforms::DTransform; + // Sanity checks + static_assert(WaveTileM >= FragM, "WaveTileM must be >= FragM"); + static_assert(WaveTileN >= FragN, "WaveTileN must be >= FragN"); + static_assert(WaveTileK >= FragK, "WaveTileK must be >= FragK"); + static_assert(WaveTileM % FragM == 0u, "WaveTileM must be a multiple of FragM"); + static_assert(WaveTileN % FragN == 0u, "WaveTileN must be a multiple of FragN"); + static_assert(WaveTileK % FragK == 0u, "WaveTileK must be a multiple of FragK"); + template (MmaPipelineOpt CK_TILE_DEVICE static void execImpl(std::tuple& vecs) { - auto& [a_vec, b_vec, c_vec, scale_A, scale_B] = vecs; - c_vec = MmaOp::exec(a_vec, b_vec, c_vec, scale_A, scale_B); + auto& [a_frag, b_frag, c_frag, scale_A, scale_B] = vecs; + + if constexpr(AccumPolicy == MmaAccumPolicy::ROW_MAJOR) + { + for(uint32_t bm = 0u; bm < FragsM; ++bm) + { + for(uint32_t bn = 0u; bn < FragsN; ++bn) + { + for(uint32_t bk = 0u; bk < FragsK; ++bk) + { + c_frag[bm][bn] = MmaOp::exec( + a_frag[bm][bk], b_frag[bn][bk], c_frag[bm][bn], scale_A, scale_B); + } + } + } + } + else if constexpr(AccumPolicy == MmaAccumPolicy::COL_MAJOR) + { + for(uint32_t bn = 0u; bn < FragsN; ++bn) + { + for(uint32_t bm = 0u; bm < FragsM; ++bm) + { + for(uint32_t bk = 0u; bk < FragsK; ++bk) + { + c_frag[bm][bn] = MmaOp::exec( + a_frag[bm][bk], b_frag[bn][bk], c_frag[bm][bn], scale_A, scale_B); + } + } + } + } + else + { + static_assert(false, "Invalid accumulation policy"); + } } }; diff --git a/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp b/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp index d57f544a41..7b2f24dea8 100644 --- a/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp +++ b/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp @@ -5,10 +5,10 @@ #include "ck_tile/core/arch/mma/mma_pipeline.hpp" #include "ck_tile/core/arch/mma/mma_selector.hpp" #include "ck_tile/core/arch/mma/mma_traits.hpp" +#include "ck_tile/core/arch/mma/mma_wavewise.hpp" #include "ck_tile/core/arch/mma/sparse/sparse_transforms.hpp" #include "ck_tile/core/numeric/vector_type.hpp" #include -#include namespace ck_tile::core::arch::mma { @@ -20,12 +20,33 @@ constexpr inline int getPipelineFlags() } } // namespace sparse::detail +/** + * @class SparseMmaPipeline + * @brief Driver for the wave-tile sparse Mma operation. Given a backend MmaOp implementation + * (e.g., smfmac), this class performs fragment-wise (MmaTile) decomposition to matrix-multiply + * input WaveTiles of (A: WaveTileM x WaveTileK) x (B: WaveTileK x WaveTileN) and accumulates + * results into output WaveTile (C: WaveTileM x WaveTileN). + * Like WaveWiseMmaPipeline, this decomposes WaveTile dimensions into fragments and iterates + * internally over FragsM × FragsN × FragsK. The A operand is provided in uncompressed form; + * 2:4 structured sparsity compression (SparseCompressTransform) is applied. + * @tparam ADataType Data type of input WaveTile A + * @tparam BDataType Data type of input WaveTile B + * @tparam CDataType Data type of input/output WaveTile C (accumulator) + * @tparam WaveTileM Mma WaveTile M dimension + * @tparam WaveTileN Mma WaveTile N dimension + * @tparam WaveTileK Mma WaveTile K dimension + * @tparam AccumPolicy The fragment order of the accum. registers (row or col major frag order) + * @tparam CompilerTarget The compiler target + * @tparam MmaOp_ Backend wrapper class that will perform the mma op + * @tparam MmaTransforms The set of transforms to be applied to input/output WaveTiles + */ template ::SelectedOp, typename MmaTransforms = // TODO: c++20 MmaTransformsI MmaTransforms = typename MmaTransformsDefaultSelector::SelectedTransforms> // clang-format off -struct SparseMmaPipeline : public MmaPipelineBase> +struct SparseMmaPipeline : public MmaPipelineBase> { - using Base = MmaPipelineBase>; + using Base = MmaPipelineBase>; // clang-format on static_assert(!Base::template hasFlag(), @@ -52,48 +73,153 @@ struct SparseMmaPipeline : public MmaPipelineBase; static constexpr index_t ASize = AVecTraits::vector_size * MmaOp::kCompressionRatio; using AVecType = ext_vector_t; }; + using ExternalAFragVecT = typename ExternalAVecCalculator::AVecType; + + // Scalar type of A + using AScalarT = typename ExternalAVecCalculator::AVecTraits::scalar_type; + + // Per-fragment sizes + static constexpr uint32_t ExternalAFragSize = ExternalAVecCalculator::ASize; + static constexpr uint32_t InternalAFragSize = + vector_traits::vector_size; - // Expose caller-side vector types - using AVecType = typename ExternalAVecCalculator::AVecType; - using BVecType = typename MmaOp::BVecType; - using CVecType = typename MmaOp::CVecType; + // Full wave-tile sizes (all fragments combined) + static constexpr uint32_t TotalUncompressedElems = FragsM * FragsK * ExternalAFragSize; + static constexpr uint32_t TotalCompressedElems = + TotalUncompressedElems / MmaOp::kCompressionRatio; - // Expose internal vector types - using InternalAVecT = typename MmaOp::AVecType; + // Variable-length idx type for the whole wave-tile (spans multiple int32_t words if needed) + static constexpr index_t IdxNumWords = sparse::detail::idx_words_needed; + using IdxType = sparse::detail::SparseIdxPack; + + // Per-fragment compressed vector type (for individual MmaOp::exec calls) + using FragAVecT = typename MmaOp::AVecType; + + // Internal vector types used by the base class formatBuffer. + // InternalAVecT matches the full compressed wave-tile so the base class can + // format the SparseCompressTransform result via formatBuffer. + using InternalAVecT = ext_vector_t; using InternalBVecT = typename MmaOp::BVecType; using InternalCVecT = typename MmaOp::CVecType; + // Buffer types for WaveTiles (caller-facing). + // A is a single flat uncompressed vector covering the whole wave-tile. + // The base class compresses it in one pass via ATransform. + using AVecType = ext_vector_t; + using BVecType = InternalBVecT[FragsN][FragsK]; + using CVecType = InternalCVecT[FragsM][FragsN]; + // Transforms using ATransform = typename MmaTransforms::ATransform; using BTransform = typename MmaTransforms::BTransform; using CTransform = typename MmaTransforms::CTransform; using DTransform = typename MmaTransforms::DTransform; + // Sanity checks + static_assert(WaveTileM >= FragM, "WaveTileM must be >= FragM"); + static_assert(WaveTileN >= FragN, "WaveTileN must be >= FragN"); + static_assert(WaveTileK >= FragK, "WaveTileK must be >= FragK"); + static_assert(WaveTileM % FragM == 0u, "WaveTileM must be a multiple of FragM"); + static_assert(WaveTileN % FragN == 0u, "WaveTileN must be a multiple of FragN"); + static_assert(WaveTileK % FragK == 0u, "WaveTileK must be a multiple of FragK"); + template CK_TILE_DEVICE static void - execImpl(std::tuple& vecs) + execImpl(std::tuple& transformedInputs) { + auto& [a, b_frag, c_frag] = transformedInputs; + auto& [a_compressed_whole, idx] = a; + + // Validate that the ATransform result and per-fragment reinterpretation are correct checkATransformResult(); - auto& [a_result, b_vec, c_vec] = vecs; - auto& [a_vec, idx] = a_result; - c_vec = MmaOp::exec(a_vec, b_vec, c_vec, idx); + + // Reinterpret the full compressed vector as per-fragment arrays + auto* a_frags = ck_tile::bit_cast(&a_compressed_whole); + + // Accumulation loop with per-fragment idx extraction + if constexpr(AccumPolicy == MmaAccumPolicy::ROW_MAJOR) + { + for(uint32_t bm = 0u; bm < FragsM; ++bm) + { + for(uint32_t bn = 0u; bn < FragsN; ++bn) + { + for(uint32_t bk = 0u; bk < FragsK; ++bk) + { + c_frag[bm][bn] = MmaOp::exec( + a_frags[bm][bk], + b_frag[bn][bk], + c_frag[bm][bn], + sparse::detail::extract_fragment_idx( + idx, bm, bk)); + } + } + } + } + else if constexpr(AccumPolicy == MmaAccumPolicy::COL_MAJOR) + { + for(uint32_t bn = 0u; bn < FragsN; ++bn) + { + for(uint32_t bm = 0u; bm < FragsM; ++bm) + { + for(uint32_t bk = 0u; bk < FragsK; ++bk) + { + c_frag[bm][bn] = MmaOp::exec( + a_frags[bm][bk], + b_frag[bn][bk], + c_frag[bm][bn], + sparse::detail::extract_fragment_idx( + idx, bm, bk)); + } + } + } + } + else + { + static_assert(false, "Invalid accumulation policy"); + } } private: - // Type check helper - not a device function, so std::declval is available + // Compile-time validation of ATransform result and per-fragment reinterpretation. + // Ensures the compressed vector returned by ATransform::exec can be safely + // reinterpreted as FragAVecT[FragsM][FragsK] for per-fragment MmaOp dispatch. template static constexpr void checkATransformResult() { using ExternalAvecRef = std::add_lvalue_reference_t; static_assert(std::is_same_v()))>); + decltype(ATransform::exec(std::declval()))>, + "ATransformResult must match the return type of ATransform::exec"); + + using CompressedVecType = + std::remove_reference_t>; + static_assert(sizeof(CompressedVecType) == sizeof(FragAVecT) * FragsM * FragsK, + "Compressed A vector size must equal sizeof(FragAVecT[FragsM][FragsK])"); + + static_assert(alignof(CompressedVecType) >= alignof(FragAVecT), + "Compressed vector alignment must be >= FragAVecT alignment " + "for safe reinterpret_cast to per-fragment array"); + + using ActualIdxType = std::tuple_element_t<1, ATransformResult>; + static_assert(std::is_same_v, + "Sparsity index type must match SparseIdxPack"); } }; diff --git a/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp b/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp index 4b0effc2bf..f89a062240 100644 --- a/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp +++ b/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp @@ -13,6 +13,30 @@ namespace ck_tile::core::arch::mma { namespace sparse::detail { + +/// Number of int32_t words needed to store CompressedSize 2-bit idx fields. +template +static constexpr index_t idx_words_needed = (CompressedSize * 2 + 31) / 32; + +/** + * @class SparseIdxPack + * @brief Variable-length container for 2:4 structured sparsity index metadata. + * + * Each compressed element produces a 2-bit index field encoding the original + * position (0–3) within its group of 4. When composing multiple MMA fragments + * in M and K dimensions within a WaveTile, the total number of index bits can + * exceed 32. This struct packs the index fields into an array of int32_t words, + * sized at compile time. + * + * @tparam NumWords Number of int32_t words needed to store all index fields. + */ +template +struct SparseIdxPack +{ + static_assert(NumWords > 0, "SparseIdxPack requires at least 1 word"); + int32_t words[NumWords] = {}; +}; + /** * @brief Compress A vector for 2:4 structured sparsity instruction by moving all non-zero * elements into lower part of a_vec to half its effective size. @@ -20,21 +44,29 @@ namespace sparse::detail { * @tparam ADataType The data type of a_vec * @tparam CompressedSize The target compression size * @tparam AVec The vector type of a_vec (deduced) - * @return Packed 32‑bit word containing **CompressedSize** 2‑bit fields. - * Each field encodes the original position (0–3) of the corresponding - * non‑zero element in the input. If fewer than CompressedSize - * non‑zeros are found, remaining fields default to 2 (see below). + * @return SparseIdxPack containing **CompressedSize** 2‑bit fields packed + * across one or more int32_t words. Each field encodes the original + * position (0–3) of the corresponding non‑zero element in the input. + * If fewer than CompressedSize non‑zeros are found, remaining fields + * default to 2 (see below). */ template -static CK_TILE_DEVICE int32_t compress_a_impl(AVec& a_vec) +static CK_TILE_DEVICE auto compress_a_impl(AVec& a_vec) { - // idx holds one 2‑bit index per output element (total CompressedSize entries). + static constexpr index_t NumIdxWords = idx_words_needed; + // idx holds one 2‑bit index per output element (total CompressedSize entries), + // packed across NumIdxWords int32_t words. // It is initialized to the pattern 0b10 for every field. This matches // what the hardware expects when there are fewer than two non‑zero values // in a 4‑element group – the unused output is treated as coming from slot 2. // The loop below will clear and set each field as real non‑zeros are seen. - int32_t idx = 0; - static_for<0, CompressedSize, 1>{}([&](auto k) { idx |= (2u << (2u * k)); }); + SparseIdxPack idx{}; + static_for<0, CompressedSize, 1>{}([&](auto k) { + constexpr uint32_t bit_pos = static_cast(k) * 2u; + constexpr uint32_t word = bit_pos / 32u; + constexpr uint32_t shift = bit_pos % 32u; + idx.words[word] |= static_cast(2u << shift); + }); static_for<0, CompressedSize / 2, 1>{}([&](auto i) { ADataType nonzero_elems[2] = {a_vec[i * 4 + 2], a_vec[i * 4 + 3]}; @@ -45,8 +77,13 @@ static CK_TILE_DEVICE int32_t compress_a_impl(AVec& a_vec) { nonzero_elems[non_zero_pos] = a_vec[i * 4 + j]; // clear the two‑bit field for this output and insert j - idx &= ~(0b11u << (2u * (i * 2 + non_zero_pos))); - idx |= static_cast(j) << (2u * (i * 2 + non_zero_pos)); + const uint32_t field_idx = + static_cast(i) * 2u + static_cast(non_zero_pos); + const uint32_t bit_pos = field_idx * 2u; + const uint32_t word = bit_pos / 32u; + const uint32_t shift = bit_pos % 32u; + idx.words[word] &= ~static_cast(0b11u << shift); + idx.words[word] |= static_cast(static_cast(j) << shift); ++non_zero_pos; } }); @@ -56,6 +93,40 @@ static CK_TILE_DEVICE int32_t compress_a_impl(AVec& a_vec) return idx; } +/** + * @brief Extract the per-fragment sparsity index from a packed idx pack. + * After whole-wave-tile compression, the returned idx packs 2-bit fields for + * every compressed output element across one or more int32_t words. + * @return A single int32_t with this fragment's 2-bit fields at the + * least-significant positions, suitable for passing to the MMA builtin. + */ +template +static CK_TILE_DEVICE int32_t extract_fragment_idx(const SparseIdxPack& idx, + uint32_t m, + uint32_t k) +{ + static constexpr uint32_t IdxBitsPerFrag = FragCompressedSize * 2; + const auto fragLinearIdx = m * FragsK + k; + const auto totalBitOffset = fragLinearIdx * IdxBitsPerFrag; + const auto wordIdx = totalBitOffset / 32u; + const auto bitInWord = totalBitOffset % 32u; + + uint32_t result = static_cast(idx.words[wordIdx]) >> bitInWord; + + // If fragment bits span a word boundary, stitch in bits from the next word. + // (This is a safety measure; it should not occur when IdxBitsPerFrag is a + // power-of-2 divisor of 32, which is always the case for current MMA ops.) + if constexpr(NumIdxWords > 1) + { + if(bitInWord != 0 && bitInWord + IdxBitsPerFrag > 32u) + { + result |= static_cast(idx.words[wordIdx + 1]) << (32u - bitInWord); + } + } + + return static_cast(result); +} + } // namespace sparse::detail /** @@ -75,15 +146,15 @@ struct SparseCompressTransform static constexpr auto VecN = VecTraits::vector_size; static constexpr index_t CompressedSize = VecN / CompressionRatio; using VecCompressed = ext_vector_t; + using IdxType = + sparse::detail::SparseIdxPack>; static_assert(VecN % CompressionRatio == 0, "VecN must be divisible by CompressionRatio"); static_assert(CompressedSize > 0, "CompressedSize must be > 0"); - const auto idx = sparse::detail::compress_a_impl(v); + auto idx = sparse::detail::compress_a_impl(v); - // TODO c++20: Use bit_cast - return std::tuple( - *std::launder(reinterpret_cast(&v)), idx); + return std::tuple(*ck_tile::bit_cast(&v), idx); } }; diff --git a/test/ck_tile/core/arch/mma/CMakeLists.txt b/test/ck_tile/core/arch/mma/CMakeLists.txt index e65a76c134..7f7817b5bf 100644 --- a/test/ck_tile/core/arch/mma/CMakeLists.txt +++ b/test/ck_tile/core/arch/mma/CMakeLists.txt @@ -7,19 +7,109 @@ if(CK_USE_OCP_FP8) list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8) endif() -if(GPU_TARGETS MATCHES "gfx9|gfx120") - add_gtest_executable(test_amdgcn_sparse_mma pipeline/test_amdgcn_sparse_mma.cpp) +# --------------------------------------------------------------------------- +# Map GPU target strings to hex amdgcn_target_id values (arch.hpp). +# Builds a -DCK_CMAKE_GPU_TARGET_IDS=0xHHHH,... definition that host-side +# test code can consume without launching a device kernel. +# --------------------------------------------------------------------------- +function(_ck_gpu_target_string_to_id TARGET_STR OUT_VAR) + string(TOLOWER "${TARGET_STR}" _tgt) + string(REGEX REPLACE ":.*" "" _tgt "${_tgt}") + # GFX9 + if(_tgt STREQUAL "gfx908") + set(${OUT_VAR} "0x0908" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx90a") + set(${OUT_VAR} "0x090A" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx942") + set(${OUT_VAR} "0x0942" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx950") + set(${OUT_VAR} "0x0950" PARENT_SCOPE) + # GFX10.3 + elseif(_tgt STREQUAL "gfx1030") + set(${OUT_VAR} "0x1030" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1031") + set(${OUT_VAR} "0x1031" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1032") + set(${OUT_VAR} "0x1032" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1033") + set(${OUT_VAR} "0x1033" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1034") + set(${OUT_VAR} "0x1034" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1035") + set(${OUT_VAR} "0x1035" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1036") + set(${OUT_VAR} "0x1036" PARENT_SCOPE) + elseif(_tgt MATCHES "^gfx10-3-generic$") + set(${OUT_VAR} "0x103F" PARENT_SCOPE) + # GFX11 + elseif(_tgt STREQUAL "gfx1100") + set(${OUT_VAR} "0x1100" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1101") + set(${OUT_VAR} "0x1101" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1102") + set(${OUT_VAR} "0x1102" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1103") + set(${OUT_VAR} "0x1103" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1150") + set(${OUT_VAR} "0x1150" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1151") + set(${OUT_VAR} "0x1151" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1152") + set(${OUT_VAR} "0x1152" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1153") + set(${OUT_VAR} "0x1153" PARENT_SCOPE) + elseif(_tgt MATCHES "^gfx11-generic$") + set(${OUT_VAR} "0x11FF" PARENT_SCOPE) + # GFX12 + elseif(_tgt STREQUAL "gfx1200") + set(${OUT_VAR} "0x1200" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1201") + set(${OUT_VAR} "0x1201" PARENT_SCOPE) + elseif(_tgt MATCHES "^gfx12-generic$") + set(${OUT_VAR} "0x12FF" PARENT_SCOPE) + elseif(_tgt STREQUAL "gfx1250") + set(${OUT_VAR} "0x1250" PARENT_SCOPE) + else() + message(WARNING "_ck_gpu_target_string_to_id: unknown GPU target '${TARGET_STR}', skipping") + set(${OUT_VAR} "" PARENT_SCOPE) + endif() +endfunction() + +function(_ck_add_gpu_target_ids_define TARGET_NAME) + get_property(_archs TARGET ${TARGET_NAME} PROPERTY HIP_ARCHITECTURES) + string(REPLACE "," ";" _archs "${_archs}") + set(_hex_ids) + foreach(_tgt IN LISTS _archs) + _ck_gpu_target_string_to_id("${_tgt}" _hex) + if(_hex AND NOT _hex STREQUAL "0x0000") + list(APPEND _hex_ids "${_hex}") + endif() + endforeach() + list(JOIN _hex_ids "," _hex_str) + if(_hex_str) + target_compile_definitions(${TARGET_NAME} PRIVATE "CK_CMAKE_GPU_TARGET_IDS=${_hex_str}") + endif() +endfunction() +# Convenience: add_gtest_executable + inject CK_CMAKE_GPU_TARGET_IDS +macro(_add_mma_gtest TEST_NAME) + add_gtest_executable(${TEST_NAME} ${ARGN}) + _ck_add_gpu_target_ids_define(${TEST_NAME}) +endmacro() +# --------------------------------------------------------------------------- + +if(GPU_TARGETS MATCHES "gfx9|gfx120") + _add_mma_gtest(test_amdgcn_sparse_mma pipeline/test_amdgcn_sparse_mma.cpp) target_compile_options(test_amdgcn_sparse_mma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) endif() if(GPU_TARGETS MATCHES "gfx950") - add_gtest_executable(test_amdgcn_scale_mma pipeline/test_amdgcn_scale_mma.cpp) + _add_mma_gtest(test_amdgcn_scale_mma pipeline/test_amdgcn_scale_mma.cpp) target_compile_options(test_amdgcn_scale_mma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) endif() if(GPU_TARGETS MATCHES "gfx9") - add_gtest_executable(test_amdgcn_mma test_amdgcn_mma.cpp) + _add_mma_gtest(test_amdgcn_mma test_amdgcn_mma.cpp) target_compile_options(test_amdgcn_mma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) - add_gtest_executable(test_amdgcn_wavewise_mma pipeline/test_amdgcn_wavewise_mma.cpp) + _add_mma_gtest(test_amdgcn_wavewise_mma pipeline/test_amdgcn_wavewise_mma.cpp) target_compile_options(test_amdgcn_wavewise_mma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) else() message(DEBUG "Skipping ck_tile_gemm tests for current target") @@ -37,45 +127,45 @@ macro(set_mma_test_arch_define target_name) endmacro() if(GPU_TARGETS MATCHES "gfx9") - add_gtest_executable(test_amdgcn_mma_layout_gfx9 test_amdgcn_mma_layout_gfx9.cpp) + _add_mma_gtest(test_amdgcn_mma_layout_gfx9 test_amdgcn_mma_layout_gfx9.cpp) target_compile_options(test_amdgcn_mma_layout_gfx9 PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) set_mma_test_arch_define(test_amdgcn_mma_layout_gfx9) endif() if(GPU_TARGETS MATCHES "gfx908|gfx90a") - add_gtest_executable(test_amdgcn_mma_layout_gfx908_and_gfx90a test_amdgcn_mma_layout_gfx908_and_gfx90a.cpp) + _add_mma_gtest(test_amdgcn_mma_layout_gfx908_and_gfx90a test_amdgcn_mma_layout_gfx908_and_gfx90a.cpp) target_compile_options(test_amdgcn_mma_layout_gfx908_and_gfx90a PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) set_mma_test_arch_define(test_amdgcn_mma_layout_gfx908_and_gfx90a) endif() if(GPU_TARGETS MATCHES "gfx90a|gfx942|gfx950") - add_gtest_executable(test_amdgcn_mma_layout_gfx90a_and_higher test_amdgcn_mma_layout_gfx90a_and_higher.cpp) + _add_mma_gtest(test_amdgcn_mma_layout_gfx90a_and_higher test_amdgcn_mma_layout_gfx90a_and_higher.cpp) target_compile_options(test_amdgcn_mma_layout_gfx90a_and_higher PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) set_mma_test_arch_define(test_amdgcn_mma_layout_gfx90a_and_higher) endif() if(GPU_TARGETS MATCHES "gfx942|gfx950") - add_gtest_executable(test_amdgcn_mma_layout_gfx942_and_higher test_amdgcn_mma_layout_gfx942_and_higher.cpp) + _add_mma_gtest(test_amdgcn_mma_layout_gfx942_and_higher test_amdgcn_mma_layout_gfx942_and_higher.cpp) target_compile_options(test_amdgcn_mma_layout_gfx942_and_higher PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS} -Wno-header-hygiene) set_mma_test_arch_define(test_amdgcn_mma_layout_gfx942_and_higher) endif() if(GPU_TARGETS MATCHES "gfx950") - add_gtest_executable(test_amdgcn_mma_layout_gfx950 test_amdgcn_mma_layout_gfx950.cpp) + _add_mma_gtest(test_amdgcn_mma_layout_gfx950 test_amdgcn_mma_layout_gfx950.cpp) target_compile_options(test_amdgcn_mma_layout_gfx950 PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) set_mma_test_arch_define(test_amdgcn_mma_layout_gfx950) endif() if(GPU_TARGETS MATCHES "gfx11") - add_gtest_executable(test_amdgcn_mma_layout_gfx11 test_amdgcn_mma_layout_gfx11.cpp) + _add_mma_gtest(test_amdgcn_mma_layout_gfx11 test_amdgcn_mma_layout_gfx11.cpp) target_compile_options(test_amdgcn_mma_layout_gfx11 PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) endif() if(GPU_TARGETS MATCHES "gfx120") - add_gtest_executable(test_amdgcn_mma_layout_gfx12 test_amdgcn_mma_layout_gfx12.cpp) + _add_mma_gtest(test_amdgcn_mma_layout_gfx12 test_amdgcn_mma_layout_gfx12.cpp) target_compile_options(test_amdgcn_mma_layout_gfx12 PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) endif() -add_gtest_executable(test_amdgcn_mma_pipeline pipeline/test_amdgcn_mma_pipeline.cpp) +_add_mma_gtest(test_amdgcn_mma_pipeline pipeline/test_amdgcn_mma_pipeline.cpp) target_compile_options(test_amdgcn_mma_pipeline PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) diff --git a/test/ck_tile/core/arch/mma/get_cmake_targets_helper.hpp b/test/ck_tile/core/arch/mma/get_cmake_targets_helper.hpp new file mode 100644 index 0000000000..eeac607a1b --- /dev/null +++ b/test/ck_tile/core/arch/mma/get_cmake_targets_helper.hpp @@ -0,0 +1,87 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck_tile/core/arch/arch.hpp" +#include + +namespace ck_tile::core::arch::testing { + +static CK_TILE_HOST auto getCMakeGpuTargetIds() +{ + using ck_tile::core::arch::amdgcn_target_id; +#ifdef CK_CMAKE_GPU_TARGET_IDS + constexpr uint32_t ids[] = {CK_CMAKE_GPU_TARGET_IDS}; + std::unordered_set result; + for(auto id : ids) + result.insert(static_cast(id)); + return result; +#else + return std::unordered_set{}; +#endif +} + +template +static CK_TILE_HOST bool dispatchCompilerTarget(ck_tile::core::arch::amdgcn_target_id id, + Func&& func) +{ + using namespace ck_tile::core::arch; + + // clang-format off + switch(id) + { + case amdgcn_target_id::GFX908: func(make_amdgcn_gfx9_target()); return true; + case amdgcn_target_id::GFX90A: func(make_amdgcn_gfx9_target()); return true; + case amdgcn_target_id::GFX942: func(make_amdgcn_gfx9_target()); return true; + case amdgcn_target_id::GFX950: func(make_amdgcn_gfx9_target()); return true; + case amdgcn_target_id::GFX1030: func(make_amdgcn_gfx10_3_target()); return true; + case amdgcn_target_id::GFX1031: func(make_amdgcn_gfx10_3_target()); return true; + case amdgcn_target_id::GFX1032: func(make_amdgcn_gfx10_3_target()); return true; + case amdgcn_target_id::GFX1033: func(make_amdgcn_gfx10_3_target()); return true; + case amdgcn_target_id::GFX1034: func(make_amdgcn_gfx10_3_target()); return true; + case amdgcn_target_id::GFX1035: func(make_amdgcn_gfx10_3_target()); return true; + case amdgcn_target_id::GFX1036: func(make_amdgcn_gfx10_3_target()); return true; + case amdgcn_target_id::GFX103_GENERIC: func(make_amdgcn_gfx10_3_target()); return true; + case amdgcn_target_id::GFX1100: func(make_amdgcn_gfx11_target()); return true; + case amdgcn_target_id::GFX1101: func(make_amdgcn_gfx11_target()); return true; + case amdgcn_target_id::GFX1102: func(make_amdgcn_gfx11_target()); return true; + case amdgcn_target_id::GFX1103: func(make_amdgcn_gfx11_target()); return true; + case amdgcn_target_id::GFX1150: func(make_amdgcn_gfx11_target()); return true; + case amdgcn_target_id::GFX1151: func(make_amdgcn_gfx11_target()); return true; + case amdgcn_target_id::GFX1152: func(make_amdgcn_gfx11_target()); return true; + case amdgcn_target_id::GFX1153: func(make_amdgcn_gfx11_target()); return true; + case amdgcn_target_id::GFX11_GENERIC: func(make_amdgcn_gfx11_target()); return true; + case amdgcn_target_id::GFX1200: func(make_amdgcn_gfx12_target()); return true; + case amdgcn_target_id::GFX1201: func(make_amdgcn_gfx12_target()); return true; + case amdgcn_target_id::GFX12_GENERIC: func(make_amdgcn_gfx12_target()); return true; + case amdgcn_target_id::GFX1250: func(make_amdgcn_gfx12_target()); return true; + case amdgcn_target_id::HOST: return false; + } + // clang-format on + __builtin_unreachable(); +} + +static CK_TILE_HOST constexpr int32_t getCMakeWaveSize() +{ + using ck_tile::core::arch::amdgcn_target_id; +#ifdef CK_CMAKE_GPU_TARGET_IDS + constexpr uint32_t ids[] = {CK_CMAKE_GPU_TARGET_IDS}; + constexpr index_t targets_size = sizeof(ids) / sizeof(ids[0]); + static_assert(targets_size > 0); + constexpr auto first_target_id = static_cast(ids[0]); + if constexpr(first_target_id >= amdgcn_target_id::GFX908 && + first_target_id <= amdgcn_target_id::GFX950) + { + return 64; + } + else + { + return 32; + } +#else + static_assert(false, "Configure CK_CMAKE_GPU_TARGET_IDS before calling this function."); + return 0; +#endif +} +} // namespace ck_tile::core::arch::testing diff --git a/test/ck_tile/core/arch/mma/get_wave_size_helper.hpp b/test/ck_tile/core/arch/mma/get_wave_size_helper.hpp deleted file mode 100644 index 84a3f955e5..0000000000 --- a/test/ck_tile/core/arch/mma/get_wave_size_helper.hpp +++ /dev/null @@ -1,34 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#pragma once - -#include -#include - -#include "ck_tile/core/arch/arch.hpp" -#include -#include "ck_tile/host/hip_check_error.hpp" - -namespace { - -__global__ void getWaveSizeForSelectedOp(uint32_t* waveSize) -{ - using CompilerTarget = decltype(ck_tile::core::arch::get_compiler_target()); - - if(waveSize) - *waveSize = static_cast(CompilerTarget::WAVE_SIZE_ID); -} - -static __host__ uint32_t getDeviceWaveSize() -{ - uint32_t* d_wave_size; - HIP_CHECK_ERROR(hipMalloc(&d_wave_size, sizeof(uint32_t))); - getWaveSizeForSelectedOp<<<1, 64>>>(d_wave_size); - HIP_CHECK_ERROR(hipDeviceSynchronize()); - uint32_t wave_size; - HIP_CHECK_ERROR(hipMemcpy(&wave_size, d_wave_size, sizeof(uint32_t), hipMemcpyDeviceToHost)); - return wave_size; -} - -} // namespace diff --git a/test/ck_tile/core/arch/mma/pipeline/pipeline_tests_helper.hpp b/test/ck_tile/core/arch/mma/pipeline/pipeline_tests_helper.hpp index 8460100aa9..edd5828a5b 100644 --- a/test/ck_tile/core/arch/mma/pipeline/pipeline_tests_helper.hpp +++ b/test/ck_tile/core/arch/mma/pipeline/pipeline_tests_helper.hpp @@ -10,223 +10,421 @@ #include #include "ck_tile/core/arch/arch.hpp" +#include "ck_tile/core/arch/mma/utility/tile_distribution_encoding_calculator.hpp" +#include "ck_tile/core/arch/mma/utility/tile_distribution_encoding_register_mapper.hpp" #include "ck_tile/core/numeric/type_convert.hpp" +#include "ck_tile/core/numeric/vector_type.hpp" #include "ck_tile/host/hip_check_error.hpp" +#include "ck_tile/host/kernel_launch.hpp" #include -#include "../get_wave_size_helper.hpp" - -template -struct MmaPipelineTest -{ - using AType = AType_; - using BType = BType_; - using CType = CType_; - using ScaleAType = ScaleAType_; - using ScaleBType = ScaleBType_; - static constexpr auto WaveTileM = WaveTileM_; - static constexpr auto WaveTileN = WaveTileN_; - static constexpr auto WaveTileK = WaveTileK_; - - void test_pipeline(std::function shouldSkip, - std::function kernel, - std::function getExpected, - std::function aInitializer = nullptr) - { - using namespace ck_tile; - using namespace ck_tile::core::arch; +#include "../get_cmake_targets_helper.hpp" - int devCount; - hipDevice_t dev; - HIP_CHECK_ERROR(hipGetDevice(&dev)); - HIP_CHECK_ERROR(hipGetDeviceCount(&devCount)); +namespace mma_pipeline_test { - hipDeviceProp_t devProp; - HIP_CHECK_ERROR(hipGetDeviceProperties(&devProp, dev)); +using namespace ck_tile; +using namespace ck_tile::core::arch; +using namespace ck_tile::core::arch::mma; +using namespace ck_tile::core::arch::testing; - auto currentArchId = hip_device_prop_gcn_arch_name_to_amdgcn_target_id(devProp.gcnArchName); - bool hasDevice = static_cast(devCount > 0); - int deviceWarpSize = devProp.warpSize; - - if(!hasDevice || shouldSkip(currentArchId)) +inline bool hipTargetMatchesCmakeTargets(amdgcn_target_id arch) +{ + const auto cmake_targets = getCMakeGpuTargetIds(); + if(cmake_targets.count(arch) == 0) + { + // gfx12-generic and gfx11-generic make no difference with the specialized archs. + // Some CI pipelines make use of that and configure the project with the generic + // flags besides compiling for (f.e.) gfx1201. + if(arch >= amdgcn_target_id::GFX1200 && arch <= amdgcn_target_id::GFX12_GENERIC) { - GTEST_SKIP() << "No HIP device found. Skipping test."; + return (cmake_targets.count(amdgcn_target_id::GFX12_GENERIC) > 0); } - - // WaveTile size, also the expected fragment size (MmaTile) from the selector. - // Note: Actual FragK might be slightly different due to hardware implementation, but the - // test_accum_over_k kernel will loop over the K dimension to ensure that the total K is - // correct. - static constexpr uint32_t FragM = WaveTileM; - static constexpr uint32_t FragN = WaveTileN; - static constexpr uint32_t FragK = WaveTileK; - - // The number of elements per thread - uint32_t AElements = FragM * FragK / deviceWarpSize; - uint32_t BElements = FragN * FragK / deviceWarpSize; - uint32_t CElements = FragM * FragN / deviceWarpSize; - - uint32_t ASize = AElements * sizeof(AType); - uint32_t BSize = BElements * sizeof(BType); - uint32_t CSize = CElements * sizeof(CType); - - // Initialize A (use custom initializer or default all 1's), B to all 1's, C to all 0's - std::vector h_a(AElements); - if(aInitializer) + else if(arch >= amdgcn_target_id::GFX1100 && arch <= amdgcn_target_id::GFX11_GENERIC) { - for(size_t i = 0; i < AElements; ++i) - h_a[i] = aInitializer(i); + return (cmake_targets.count(amdgcn_target_id::GFX11_GENERIC) > 0); } - else + } + return true; +} +template +void reference_matmul(std::vector& C, + const std::vector& A, + const std::vector& B, + uint32_t M, + uint32_t N, + uint32_t K) +{ + for(uint32_t m = 0; m < M; ++m) + { + for(uint32_t n = 0; n < N; ++n) { - std::fill(h_a.begin(), h_a.end(), type_convert(1)); + float acc = 0.0f; + for(uint32_t k = 0; k < K; ++k) + { + acc += type_convert(A[m * K + k]) * type_convert(B[k * N + n]); + } + C[m * N + n] = static_cast(acc); } - std::vector h_b(BElements, type_convert(1)); - std::vector h_c(CElements, type_convert(0)); - std::vector h_out(CElements, type_convert(0)); + } +} - AType* d_a; - BType* d_b; - CType* d_c; - CType* d_out; +template +T deterministic_value(uint32_t row, uint32_t col, uint32_t minor_dim) +{ + float v = static_cast((row * minor_dim + col) % 7 + 1) * 0.25f; + return type_convert(v); +} + +// Apply 2:4 sparsity pattern to A matrix in-place (for sparse pipeline tests). +// Every group of 4 consecutive K elements keeps slots 0 and 2, zeros slots 1 and 3. +template +void apply_sparse_pattern(std::vector& A, uint32_t M, uint32_t K) +{ + for(uint32_t m = 0; m < M; ++m) + { + for(uint32_t k = 0; k < K; k += 4) + { + // Keep slots 0, 2. Zero out slots 1, 3. + if(k + 1 < K) + A[m * K + k + 1] = static_cast(0); + if(k + 3 < K) + A[m * K + k + 3] = static_cast(0); + } + } +} + +// Fill per-lane A fragments from logical A[M][K] matrix. +// For dense pipelines: AVecType = InternalAVecT[FragsM][FragsK] +// For sparse pipelines: AVecType = ExternalAFragVecT[FragsM][FragsK] (uncompressed) +template +void fill_a_fragments(typename Pipeline::AVecType* a_per_lane, + const std::vector& A_matrix, + uint32_t K, + uint32_t waveSize) +{ + using MmaOp = typename Pipeline::MmaOp; + using ARegMap = TileDistrEncRegMap::AWarpDstrEncoding>; + using AFragScalar = typename vector_traits::scalar_type; - HIP_CHECK_ERROR(hipMalloc(&d_a, ASize)); - HIP_CHECK_ERROR(hipMalloc(&d_b, BSize)); - HIP_CHECK_ERROR(hipMalloc(&d_c, CSize)); - HIP_CHECK_ERROR(hipMalloc(&d_out, CSize)); + constexpr uint32_t FragM = Pipeline::FragM; + constexpr uint32_t FragK = Pipeline::FragK; + constexpr uint32_t FragsM = Pipeline::FragsM; + constexpr uint32_t FragsK = Pipeline::FragsK; - // Copy inputs to device - HIP_CHECK_ERROR(hipMemcpy(d_a, h_a.data(), ASize, hipMemcpyHostToDevice)); - HIP_CHECK_ERROR(hipMemcpy(d_b, h_b.data(), BSize, hipMemcpyHostToDevice)); - HIP_CHECK_ERROR(hipMemcpy(d_c, h_c.data(), CSize, hipMemcpyHostToDevice)); + constexpr uint32_t kCompressionRatio = MmaOp::kCompressionRatio; - const auto wave_size = getDeviceWaveSize(); - kernel(wave_size, d_a, d_b, d_c, d_out); - HIP_CHECK_ERROR(hipDeviceSynchronize()); + // The A register map maps (lane, vec_idx) -> (m_within_frag, k_within_frag) + // For sparse: k_within_frag is in the compressed K domain (K / kCompressionRatio) + constexpr index_t a_vec_size = ARegMap::num_vector_items; + constexpr index_t external_a_frag_vec_size = a_vec_size * kCompressionRatio; - HIP_CHECK_ERROR(hipMemcpy(h_out.data(), d_out, CSize, hipMemcpyDeviceToHost)); + for(uint32_t lane = 0; lane < waveSize; ++lane) + { + auto* lane_a = reinterpret_cast(&a_per_lane[lane]); - // Verify output against expected value for all elements - for(size_t i = 0; i < CElements; ++i) + for(uint32_t bm = 0; bm < FragsM; ++bm) { - EXPECT_NEAR(h_out[i], getExpected(FragK), 1e-3); + for(uint32_t bk = 0; bk < FragsK; ++bk) + { + uint32_t frag_offset = (bm * FragsK + bk) * external_a_frag_vec_size; + + if constexpr(kCompressionRatio > 1) + { + // Sparse: fill external (uncompressed) vector + for(index_t ev = 0; ev < external_a_frag_vec_size; ++ev) + { + index_t compressed_v = ev / kCompressionRatio; + index_t sub_pos = ev % kCompressionRatio; + + auto coords = + ARegMap::calc_matrix_indices_from_lane_vector(lane, compressed_v); + uint32_t m_local = coords[0]; + uint32_t k_compressed = coords[1]; + uint32_t k_local = k_compressed * kCompressionRatio + sub_pos; + + uint32_t m_global = bm * FragM + m_local; + uint32_t k_global = bk * FragK + k_local; + + lane_a[frag_offset + ev] = + static_cast(A_matrix[m_global * K + k_global]); + } + } + else + { + // Dense/Scale: direct mapping + for(index_t v = 0; v < a_vec_size; ++v) + { + auto coords = ARegMap::calc_matrix_indices_from_lane_vector(lane, v); + uint32_t m_local = coords[0]; + uint32_t k_local = coords[1]; + + uint32_t m_global = bm * FragM + m_local; + uint32_t k_global = bk * FragK + k_local; + + lane_a[frag_offset + v] = + static_cast(A_matrix[m_global * K + k_global]); + } + } + } } - - HIP_CHECK_ERROR(hipFree(d_a)); - HIP_CHECK_ERROR(hipFree(d_b)); - HIP_CHECK_ERROR(hipFree(d_c)); - HIP_CHECK_ERROR(hipFree(d_out)); } +} + +// Fill per-lane B fragments from logical B[K][N] matrix. +// BVecType = InternalBVecT[FragsN][FragsK] +template +void fill_b_fragments(typename Pipeline::BVecType* b_per_lane, + const std::vector& B_matrix, + uint32_t N, + uint32_t waveSize) +{ + using MmaOp = typename Pipeline::MmaOp; + using BRegMap = TileDistrEncRegMap::BWarpDstrEncoding>; + using BFragScalar = typename vector_traits::scalar_type; - void - test_pipeline(std::function shouldSkip, - std::function kernel, - std::function getExpected, - std::function aInitializer = nullptr) - { - using namespace ck_tile; - using namespace ck_tile::core::arch; - - int devCount; - hipDevice_t dev; - HIP_CHECK_ERROR(hipGetDevice(&dev)); - HIP_CHECK_ERROR(hipGetDeviceCount(&devCount)); + constexpr uint32_t FragN = Pipeline::FragN; + constexpr uint32_t FragK = Pipeline::FragK; + constexpr uint32_t FragsN = Pipeline::FragsN; + constexpr uint32_t FragsK = Pipeline::FragsK; - hipDeviceProp_t devProp; - HIP_CHECK_ERROR(hipGetDeviceProperties(&devProp, dev)); + constexpr index_t b_vec_size = BRegMap::num_vector_items; - auto currentArchId = hip_device_prop_gcn_arch_name_to_amdgcn_target_id(devProp.gcnArchName); - bool hasDevice = static_cast(devCount > 0); - int deviceWarpSize = devProp.warpSize; + for(uint32_t lane = 0; lane < waveSize; ++lane) + { + auto* lane_b = reinterpret_cast(&b_per_lane[lane]); - if(!hasDevice || shouldSkip(currentArchId)) + for(uint32_t bn = 0; bn < FragsN; ++bn) { - GTEST_SKIP() << "No HIP device found. Skipping test."; + for(uint32_t bk = 0; bk < FragsK; ++bk) + { + uint32_t frag_offset = (bn * FragsK + bk) * b_vec_size; + + for(index_t v = 0; v < b_vec_size; ++v) + { + auto coords = BRegMap::calc_matrix_indices_from_lane_vector(lane, v); + uint32_t n_local = coords[0]; + uint32_t k_local = coords[1]; + + uint32_t n_global = bn * FragN + n_local; + uint32_t k_global = bk * FragK + k_local; + + // B matrix is stored as B[K][N] + lane_b[frag_offset + v] = + static_cast(B_matrix[k_global * N + n_global]); + } + } } + } +} + +// Extract C matrix from per-lane C fragments. +// CVecType = InternalCVecT[FragsM][FragsN] +template +void extract_c_matrix(const typename Pipeline::CVecType* c_per_lane, + std::vector& C_matrix, + uint32_t N, + uint32_t waveSize) +{ + using MmaOp = typename Pipeline::MmaOp; + using CRegMap = TileDistrEncRegMap::CWarpDstrEncoding>; + using CFragScalar = typename vector_traits::scalar_type; + + constexpr uint32_t FragM = Pipeline::FragM; + constexpr uint32_t FragN = Pipeline::FragN; + constexpr uint32_t FragsM = Pipeline::FragsM; + constexpr uint32_t FragsN = Pipeline::FragsN; - // WaveTile size, also the expected fragment size (MmaTile) from the selector. - // Note: Actual FragK might be slightly different due to hardware implementation, but the - // test_accum_over_k kernel will loop over the K dimension to ensure that the total K is - // correct. - static constexpr uint32_t FragM = WaveTileM; - static constexpr uint32_t FragN = WaveTileN; - static constexpr uint32_t FragK = WaveTileK; - - // The number of elements per thread - uint32_t AElements = FragM * FragK / deviceWarpSize / numeric_traits::PackedSize; - uint32_t BElements = FragN * FragK / deviceWarpSize / numeric_traits::PackedSize; - uint32_t CElements = FragM * FragN / deviceWarpSize; - - uint32_t ASize = AElements * sizeof(AType); - uint32_t BSize = BElements * sizeof(BType); - uint32_t CSize = CElements * sizeof(CType); - uint32_t ScaleASize = 1 * sizeof(ScaleAType); - uint32_t ScaleBSize = 1 * sizeof(ScaleBType); - - // Initialize A (use custom initializer or default all 1's), B to all 1's, C to all 0's - std::vector h_a(AElements); - if(aInitializer) + constexpr index_t c_vec_size = CRegMap::num_vector_items; + + for(uint32_t lane = 0; lane < waveSize; ++lane) + { + auto* lane_c = reinterpret_cast(&c_per_lane[lane]); + + for(uint32_t bm = 0; bm < FragsM; ++bm) { - for(size_t i = 0; i < AElements; ++i) - h_a[i] = aInitializer(i); + for(uint32_t bn = 0; bn < FragsN; ++bn) + { + uint32_t frag_offset = (bm * FragsN + bn) * c_vec_size; + + for(index_t v = 0; v < c_vec_size; ++v) + { + auto coords = CRegMap::calc_matrix_indices_from_lane_vector(lane, v); + uint32_t m_local = coords[0]; + uint32_t n_local = coords[1]; + + uint32_t m_global = bm * FragM + m_local; + uint32_t n_global = bn * FragN + n_local; + + C_matrix[m_global * N + n_global] = + static_cast(lane_c[frag_offset + v]); + } + } } - else + } +} + +/// Internal: runs the test with a fully resolved Pipeline type. +/// Called from run_pipeline_matrix_test after dispatching on compiler target. +template +void run_pipeline_matrix_test_impl(uint32_t M, + uint32_t N, + uint32_t K, + uint32_t waveSize, + KernelType kernel, + bool isSparse, + bool transposeExpected = false, + float referenceScale = 1.0f) +{ + std::vector A_matrix(M * K); + std::vector B_matrix(K * N); + std::vector C_expected(M * N, static_cast(0)); + std::vector C_actual(M * N, static_cast(0)); + + for(uint32_t m = 0; m < M; ++m) + for(uint32_t k = 0; k < K; ++k) + A_matrix[m * K + k] = deterministic_value(m, k, K); + + for(uint32_t k = 0; k < K; ++k) + for(uint32_t n = 0; n < N; ++n) + B_matrix[k * N + n] = deterministic_value(k, n, N); + + if(isSparse) + { + apply_sparse_pattern(A_matrix, M, K); + } + + reference_matmul(C_expected, A_matrix, B_matrix, M, N, K); + + using AVecType = typename Pipeline::AVecType; + using BVecType = typename Pipeline::BVecType; + using CVecType = typename Pipeline::CVecType; + + const size_t a_buf_size = waveSize * sizeof(AVecType); + const size_t b_buf_size = waveSize * sizeof(BVecType); + const size_t c_buf_size = waveSize * sizeof(CVecType); + + std::vector h_a(a_buf_size, 0); + std::vector h_b(b_buf_size, 0); + std::vector h_c(c_buf_size, 0); + + fill_a_fragments(reinterpret_cast(h_a.data()), A_matrix, K, waveSize); + fill_b_fragments(reinterpret_cast(h_b.data()), B_matrix, N, waveSize); + + void *d_a, *d_b, *d_c; + HIP_CHECK_ERROR(hipMalloc(&d_a, a_buf_size)); + HIP_CHECK_ERROR(hipMalloc(&d_b, b_buf_size)); + HIP_CHECK_ERROR(hipMalloc(&d_c, c_buf_size)); + + HIP_CHECK_ERROR(hipMemcpy(d_a, h_a.data(), a_buf_size, hipMemcpyHostToDevice)); + HIP_CHECK_ERROR(hipMemcpy(d_b, h_b.data(), b_buf_size, hipMemcpyHostToDevice)); + HIP_CHECK_ERROR(hipMemset(d_c, 0, c_buf_size)); + + ck_tile::launch_kernel(ck_tile::stream_config{}, + ck_tile::make_kernel(kernel, dim3(1), dim3(waveSize), 0, d_a, d_b, d_c)); + HIP_CHECK_ERROR(hipDeviceSynchronize()); + + HIP_CHECK_ERROR(hipMemcpy(h_c.data(), d_c, c_buf_size, hipMemcpyDeviceToHost)); + extract_c_matrix( + reinterpret_cast(h_c.data()), C_actual, N, waveSize); + + for(uint32_t m = 0; m < M; ++m) + { + for(uint32_t n = 0; n < N; ++n) { - std::fill(h_a.begin(), h_a.end(), type_convert(1.0f)); + // When transposeExpected is true, the kernel computes C^T via SwapAB, + // so compare actual C[m][n] against reference C[n][m]. + constexpr float relative_tolerance = 1e-2f; + constexpr float absolute_tolerance = 1e-3f; + + float expected = transposeExpected ? static_cast(C_expected[n * M + m]) + : static_cast(C_expected[m * N + n]); + expected *= referenceScale; + float actual = static_cast(C_actual[m * N + n]); + EXPECT_NEAR( + actual, expected, std::abs(expected) * relative_tolerance + absolute_tolerance) + << "Mismatch at C[" << m << "][" << n << "]"; } - std::vector h_b(BElements, type_convert(1.0f)); - std::vector h_c(CElements, type_convert(0.0f)); - std::vector h_out(CElements, type_convert(0.0f)); - // The actual scale is computed as pow(2, scale - 127), so: - // 126 -> 2^-1 and 129 -> 2^2. - ScaleAType h_scale_a = 126; - ScaleBType h_scale_b = 129; - - AType* d_a; - BType* d_b; - CType* d_c; - CType* d_out; - ScaleAType* d_scale_a; - ScaleBType* d_scale_b; - - HIP_CHECK_ERROR(hipMalloc(&d_a, ASize)); - HIP_CHECK_ERROR(hipMalloc(&d_b, BSize)); - HIP_CHECK_ERROR(hipMalloc(&d_c, CSize)); - HIP_CHECK_ERROR(hipMalloc(&d_out, CSize)); - HIP_CHECK_ERROR(hipMalloc(&d_scale_a, ScaleASize)); - HIP_CHECK_ERROR(hipMalloc(&d_scale_b, ScaleBSize)); - - // Copy inputs to device - HIP_CHECK_ERROR(hipMemcpy(d_a, h_a.data(), ASize, hipMemcpyHostToDevice)); - HIP_CHECK_ERROR(hipMemcpy(d_b, h_b.data(), BSize, hipMemcpyHostToDevice)); - HIP_CHECK_ERROR(hipMemcpy(d_c, h_c.data(), CSize, hipMemcpyHostToDevice)); - HIP_CHECK_ERROR(hipMemcpy(d_scale_a, &h_scale_a, ScaleASize, hipMemcpyHostToDevice)); - HIP_CHECK_ERROR(hipMemcpy(d_scale_b, &h_scale_b, ScaleBSize, hipMemcpyHostToDevice)); - - const auto wave_size = getDeviceWaveSize(); - kernel(wave_size, d_a, d_b, d_c, d_out, d_scale_a, d_scale_b); - HIP_CHECK_ERROR(hipDeviceSynchronize()); - - HIP_CHECK_ERROR(hipMemcpy(h_out.data(), d_out, CSize, hipMemcpyDeviceToHost)); - - // Verify output against expected value for all elements - for(size_t i = 0; i < CElements; ++i) + } + + HIP_CHECK_ERROR(hipFree(d_a)); + HIP_CHECK_ERROR(hipFree(d_b)); + HIP_CHECK_ERROR(hipFree(d_c)); +} + +/// @tparam PipelineFactory A template template that, given a CompilerTarget type, produces +/// the Pipeline type: PipelineFactory::type +/// @tparam KernelType Kernel functor struct with kBlockSize and __device__ operator() +/// @tparam AScalar Scalar type for A matrix (e.g., fp16_t) +/// @tparam BScalar Scalar type for B matrix (e.g., fp16_t) +/// @tparam CScalar Scalar type for C matrix (e.g., fp32_t) +/// @param M WaveTile M dimension +/// @param N WaveTile N dimension +/// @param K WaveTile K dimension +/// @param shouldSkip Predicate returning true if current device should skip +/// @param kernel Kernel functor instance to launch via make_kernel +/// @param isSparse Whether to apply 2:4 sparsity pattern to A +/// @param transposeExpected When true, compare against transposed reference (for +/// SwapAB/TransposeC) +/// @param referenceScale Scalar multiplier applied to the reference matmul result before +/// comparison (e.g., to account for scale-MMA scaling factors) +template