diff --git a/AGENTS.md b/AGENTS.md index 9f4ec08521..3148ff0999 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -46,7 +46,7 @@ See [docs/install_vortex.md](docs/install_vortex.md) for the full recipe. The no - **Out-of-tree.** From the repo root: ```bash mkdir -p build && cd build - ../configure --xlen=64 --tooldir=$HOME/tools # or --xlen=32 + ../configure --xlen=32 --tooldir=$HOME/tools # or --xlen=64 ./ci/toolchain_install.sh # first time only make -s ``` @@ -64,6 +64,7 @@ See [docs/bug_fixes.md](docs/bug_fixes.md) for the full rationale and examples. - **Fix root causes, not symptoms.** Diagnose before patching. - **Don't paper over upstream regressions** or mask bugs with fallback paths and suppressed warnings. +- **Don't add legacy API compatibility or fallback code** unless it is explicitly required for the change. - **If a patch is genuinely unavoidable** (e.g. blocked by an external dep), label it as a patch explicitly *in the commit message* and pair it with a follow-up to do the proper fix. --- @@ -73,7 +74,6 @@ See [docs/bug_fixes.md](docs/bug_fixes.md) for the full rationale and examples. See [docs/testing.md](docs/testing.md) and [docs/debugging.md](docs/debugging.md) for recipes. The rules: - **All test commands run from `build/`.** The generated `ci/` scripts assume it. -- **120-second timeout cap** on every test invocation. No exceptions. - **`CONFIGS` must match on both sides.** `blackbox.sh` only rebuilds the driver. If your test was compiled with `-DVX_CFG_NUM_THREADS=4`, blackbox can't fix that — rebuild the app first: ```bash make -C tests/regression/ clean @@ -82,6 +82,7 @@ See [docs/testing.md](docs/testing.md) and [docs/debugging.md](docs/debugging.md ``` - **`make tests` / `make -C tests/regression` build with *default* macros.** Use `CONFIGS` + explicit per-app rebuild for non-default configurations. - **`--rebuild=1` forces a driver rebuild** even if the hardware configuration is unchanged. Use it when iterating on the driver itself; `--rebuild=0` suppresses rebuild regardless. +- **Keep smoke/regression runs lean.** Don't enable debug or perf collection unless the run is explicitly for debug or measurement. - **RTL coverage path is `xrt`, not `rtlsim`.** When discussing or planning RTL verification, `xrt` is the canonical path — `rtlsim` bypasses the AFU surface. `rtlsim` remains useful for fast iteration on processor RTL; `xrt` is what proves the full integration. - **`ci/regression.sh` is the canonical source of tested configurations.** Use it to discover supported parameter combinations before inventing ad hoc ones. - **Perf-regression baselines (`ci/perf/baselines/*.json`) are golden data — never hand-edit them, and never "fix" a red perf gate by bumping the number.** They are regenerated only by `pytest ci -m perf_gate --update-baselines` (a human-run, reviewed step), and CI must never pass that flag. A `perf_gate` failure means real cycles moved: root-cause it, or — if the change is intended — regenerate the baseline so the diff shows the perf delta for review. Same rule as image goldens and `known_issue:`. @@ -108,6 +109,23 @@ make -C tests/opencl run-rtlsim `blackbox.sh` exposes the common knobs directly: `--clusters=`, `--cores=`, `--warps=`, `--threads=`, `--l2cache`, `--l3cache`, `--debug=`, `--perf=`. For anything not exposed as a flag, use `CONFIGS="-D..."` (all parameters take the `VX_CFG_*` prefix — e.g. `-DVX_CFG_NUM_THREADS=8`, `-DVX_CFG_EXT_TCU_ENABLE`). Baseline parameters live in `VX_config.toml` and `VX_types.toml` at the repo root — edit those only when an override is needed for *all* builds, and re-`configure` afterward. +```bash +./ci/blackbox.sh --driver=simx --app=sgemm --clusters=1 --cores=2 --warps=4 --threads=4 --l2cache +``` +OR + +```bash +CONFIGS="-DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_TYPE_TFR -DITYPE=bf16 -DOTYPE=fp32" ./ci/blackbox.sh --driver=rtlsim --app=sgemm_tcu +``` + +### Roofline Analysis + +Use `--perf=1` for detailed performance counters such as scheduler utilization, pipeline stalls, instruction mix, and memory latency. + +```bash +/usr/bin/python3 ../perf/roofline.py --app=sgemm_tcu --driver=simx --cores=1 --warps=4 --threads=8 --issue-width=2 --perf=1 --by-cycle --output=sgemm_tcu_roofline.png +``` + --- ## 5. Design & Architecture Rules diff --git a/VX_config.toml b/VX_config.toml index 7e46c2174b..b663baa627 100644 --- a/VX_config.toml +++ b/VX_config.toml @@ -253,6 +253,7 @@ VX_CFG_TCU_NVFP4_ENABLE = false VX_CFG_TCU_SPARSE_ENABLE = false VX_CFG_TCU_DSM_ENABLE = false VX_CFG_TCU_WGMMA_ENABLE = false +VX_CFG_TCU_FEDP2K = false [dxa] # Cluster-level DXA engine count, decoupled from cores/socket size. diff --git a/ci/perf/baselines/dxa.json b/ci/perf/baselines/dxa.json index b57945e94c..75abe31a8e 100644 --- a/ci/perf/baselines/dxa.json +++ b/ci/perf/baselines/dxa.json @@ -61,7 +61,7 @@ }, "dxa:perf_gate-wgmma-sp-dxa:rtlsim": { "32": { - "cycles": 113300, + "cycles": 108646, "instrs": 9344 }, "app": "sgemm_tcu_wg_sp_dxa", diff --git a/ci/perf/baselines/tensor_wg.json b/ci/perf/baselines/tensor_wg.json index ec27cd07f0..d2d3328504 100644 --- a/ci/perf/baselines/tensor_wg.json +++ b/ci/perf/baselines/tensor_wg.json @@ -1,22 +1,42 @@ { - "tensor_wg:perf_gate-wgmma-fp16:rtlsim": { + "tensor_wg:perf_gate-wgmma-fp16-ss:rtlsim": { "32": { - "cycles": 758510, - "instrs": 525536 + "cycles": 107212, + "instrs": 66640 }, "app": "sgemm_tcu_wg", - "args": "-m 128 -n 128 -k 128", - "config_hash": "020544d97c1a5a27", - "configs": "-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=32" + "args": "-m 64 -n 64 -k 64", + "config_hash": "64856bfe7975f530", + "configs": "-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=32 -DWGMMA_SS" + }, + "tensor_wg:perf_gate-wgmma-fedp2k-rs:rtlsim": { + "32": { + "cycles": 166490, + "instrs": 108512 + }, + "app": "sgemm_tcu_wg", + "args": "-m64 -n64 -k64", + "config_hash": "84f3ab0ee89d1c7e", + "configs": "-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_FEDP2K -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS" }, - "tensor_wg:perf_gate-wgmma-sparse:rtlsim": { + "tensor_wg:perf_gate-wgmma-sp-rs:rtlsim": { "32": { - "cycles": 812012, + "cycles": 778048, "instrs": 438752 }, "app": "sgemm_tcu_wg_sp", "args": "-m 128 -n 128 -k 128", - "config_hash": "b94d94645f0461c8", - "configs": "-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DWGMMA_NRC=32" + "config_hash": "6524c99e95160ff1", + "configs": "-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DWGMMA_NRC=32 -DWGMMA_RS" + }, + "tensor_wg:perf_gate-wgmma-mx-fedp2k-rs:rtlsim": { + "32": { + "cycles": 139877, + "instrs": 86304 + }, + "app": "sgemm_tcu_wg_mx", + "args": "-m64 -n64 -k64", + "config_hash": "a0d52f0bb1cc7294", + "configs": "-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DVX_CFG_TCU_FEDP2K -DVX_CFG_TCU_TYPE_DPI -DITYPE=mxfp8 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS" } } diff --git a/ci/testcases/tensor.yaml b/ci/testcases/tensor.yaml index 9a97ccee1b..5adce6ed70 100644 --- a/ci/testcases/tensor.yaml +++ b/ci/testcases/tensor.yaml @@ -152,6 +152,13 @@ tests: -DITYPE=fp16 -DOTYPE=fp32 args: -w 4 +- id: sgemm_tcu-fedp2k + via: blackbox + drivers: + - simx + app: sgemm_tcu + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_FEDP2K -DITYPE=fp16 -DOTYPE=fp32 + # SimX<->rtlsim cycle parity on the native (non-DPI) fp16 TCU datapath: same # config on both drivers, cycles must agree within tolerance (see model_parity.yaml). - id: model_parity-fp16 @@ -205,3 +212,13 @@ tests: app: sgemm2_tcu args: -m 128 -n 128 -k 64 configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DITYPE=fp16 -DOTYPE=fp32 + +- id: model_parity-fedp2k + check: model_parity + # Same pre-existing LSU timing gap as model_parity-fp16; FEDP2K changes the + # datapath width but does not change the retired instruction stream. + known_issue: "cycle-parity gap ~6%: SimX does not model the decoupled LSU pending pool (instrs match)" + via: blackbox + app: sgemm_tcu + args: -m64 -n64 -k64 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_FEDP2K -DVX_CFG_TCU_TYPE_DPI -DITYPE=fp16 -DOTYPE=fp32 diff --git a/ci/testcases/tensor_mx.yaml b/ci/testcases/tensor_mx.yaml index b85e22d8f3..c6c92cf77d 100644 --- a/ci/testcases/tensor_mx.yaml +++ b/ci/testcases/tensor_mx.yaml @@ -97,58 +97,29 @@ tests: drivers: - simx app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 + configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DTCU_MX_TLS args: "-k32" - id: sgemm_tcu_mx-simx-nvfp4-nt8 via: blackbox drivers: - simx app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DTCU_MX_TLS args: "-k64" - id: sgemm_tcu_mx-simx-nvfp4-nt16 via: blackbox drivers: - simx app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 + configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DTCU_MX_TLS args: "-k64" - id: sgemm_tcu_mx-simx-nvfp4-nt32 via: blackbox drivers: - simx app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 + configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DTCU_MX_TLS args: "-k128" -- id: sgemm_tcu_mx-simx-mxint8-nt4 - via: blackbox - drivers: - - simx - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 - args: "-k32" -- id: sgemm_tcu_mx-simx-mxint8-nt8 - via: blackbox - drivers: - - simx - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 - args: "-k32" -- id: sgemm_tcu_mx-simx-mxint8-nt16 - via: blackbox - drivers: - - simx - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 - args: "-k32" - known_issue: "mxint8 NT16 simx numerical mismatch (NT8/NT32 pass) — pending MX fix" -- id: sgemm_tcu_mx-simx-mxint8-nt32 - via: blackbox - drivers: - - simx - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 - args: "-k64" - id: sgemm_tcu_mx-rtlsim-mxfp8-nt4 via: blackbox drivers: @@ -156,7 +127,6 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI args: "-k32" - known_issue: "mxfp8 rtlsim DPI-path numerical mismatch (simx passes) — pending MX fix" - id: sgemm_tcu_mx-rtlsim-mxfp8-nt8 via: blackbox drivers: @@ -164,7 +134,6 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI args: "-k32" - known_issue: "mxfp8 rtlsim DPI-path numerical mismatch (simx passes) — pending MX fix" - id: sgemm_tcu_mx-rtlsim-mxfp8-nt16 via: blackbox drivers: @@ -172,7 +141,6 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI args: "-k32" - known_issue: "mxfp8 rtlsim DPI-path numerical mismatch (simx passes) — pending MX fix" - id: sgemm_tcu_mx-rtlsim-mxfp8-nt32 via: blackbox drivers: @@ -180,7 +148,6 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI args: "-k64" - known_issue: "mxfp8 rtlsim DPI-path numerical mismatch (simx passes) — pending MX fix" - id: sgemm_tcu_mx-rtlsim-mxbf8-nt4 via: blackbox drivers: @@ -242,57 +209,29 @@ tests: drivers: - rtlsim app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI + configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI -DTCU_MX_TLS args: "-k32" - id: sgemm_tcu_mx-rtlsim-nvfp4-nt8 via: blackbox drivers: - rtlsim app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI -DTCU_MX_TLS args: "-k64" - id: sgemm_tcu_mx-rtlsim-nvfp4-nt16 via: blackbox drivers: - rtlsim app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI + configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI -DTCU_MX_TLS args: "-k64" - id: sgemm_tcu_mx-rtlsim-nvfp4-nt32 via: blackbox drivers: - rtlsim app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI + configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI -DTCU_MX_TLS args: "-k128" -- id: sgemm_tcu_mx-rtlsim-mxint8-nt4 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 -DVX_CFG_TCU_TYPE_DPI - args: "-k32" -- id: sgemm_tcu_mx-rtlsim-mxint8-nt8 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 -DVX_CFG_TCU_TYPE_DPI - args: "-k32" -- id: sgemm_tcu_mx-rtlsim-mxint8-nt16 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 -DVX_CFG_TCU_TYPE_DPI - args: "-k32" -- id: sgemm_tcu_mx-rtlsim-mxint8-nt32 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 -DVX_CFG_TCU_TYPE_DPI - args: "-k64" - id: sgemm_tcu_mx-simx-nvfp4-nt4-tls via: blackbox drivers: diff --git a/ci/testcases/tensor_sp_mx.yaml b/ci/testcases/tensor_sp_mx.yaml index 6e6ab8dc67..2d6c20458c 100644 --- a/ci/testcases/tensor_sp_mx.yaml +++ b/ci/testcases/tensor_sp_mx.yaml @@ -6,10 +6,6 @@ defaults: - 64 tier: smoke touches: [hw/rtl/tcu, sim/simx/tcu, sw/kernel/include/vx_tensor.h] - # Seed coverage only — the sparse-MX path has no validated reference yet, so - # every case currently mismatches. Tracked xfail until the sparse-MX result is - # validated; drop this once the cases pass. - known_issue: "sparse-MX (sgemm_tcu_sp_mx) lacks a validated reference; results mismatch" tests: - id: sgemm_tcu_sp_mx-simx-mxfp8-nt4 via: blackbox @@ -25,20 +21,6 @@ tests: app: sgemm_tcu_sp_mx configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 args: "-k32" -- id: sgemm_tcu_sp_mx-simx-mxint8-nt4 - via: blackbox - drivers: - - simx - app: sgemm_tcu_sp_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 - args: "-k32" -- id: sgemm_tcu_sp_mx-simx-mxint8-nt8 - via: blackbox - drivers: - - simx - app: sgemm_tcu_sp_mx - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 - args: "-k32" - id: sgemm_tcu_sp_mx-rtlsim-mxfp8-nt4 via: blackbox drivers: @@ -53,17 +35,3 @@ tests: app: sgemm_tcu_sp_mx configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI args: "-k32" -- id: sgemm_tcu_sp_mx-rtlsim-mxint8-nt4 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_sp_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 -DVX_CFG_TCU_TYPE_DPI - args: "-k32" -- id: sgemm_tcu_sp_mx-rtlsim-mxint8-nt8 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_sp_mx - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=mxint8 -DOTYPE=int32 -DVX_CFG_TCU_TYPE_DPI - args: "-k32" diff --git a/ci/testcases/tensor_wg.yaml b/ci/testcases/tensor_wg.yaml index 69dfd24fdd..073deae765 100644 --- a/ci/testcases/tensor_wg.yaml +++ b/ci/testcases/tensor_wg.yaml @@ -7,54 +7,6 @@ defaults: tier: smoke touches: [hw/rtl/tcu, sim/simx/tcu, sw/kernel/include/vx_tensor.h] tests: -- id: sgemm_tcu-1 - via: blackbox - drivers: - - simx - app: sgemm_tcu - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -- id: sgemm_tcu-2 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -- id: sgemm_tcu_sp-1 - via: blackbox - drivers: - - simx - app: sgemm_tcu_sp - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -- id: sgemm_tcu_sp-2 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_sp - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -- id: sgemm_tcu_sp-3 - via: blackbox - drivers: - - simx - app: sgemm_tcu_sp - configs: -DVX_CFG_NUM_THREADS=2 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=int8 -DOTYPE=int32 -- id: sgemm_tcu_sp-4 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_sp - configs: -DVX_CFG_NUM_THREADS=2 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_TYPE_DPI -DVX_CFG_TCU_INT8_ENABLE -DITYPE=int8 -DOTYPE=int32 -- id: sgemm_tcu_sp-5 - via: blackbox - drivers: - - simx - app: sgemm_tcu_sp - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -- id: sgemm_tcu_sp-6 - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_sp - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_TYPE_DPI -DITYPE=fp16 -DOTYPE=fp32 - id: sgemm_tcu_wg-1 via: blackbox drivers: @@ -104,6 +56,36 @@ tests: app: sgemm_tcu_wg configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_TF32_ENABLE -DITYPE=tf32 -DOTYPE=fp32 -DWGMMA_NRC=16 +- id: sgemm_tcu_wg-8 + via: blackbox + drivers: + - simx + app: sgemm_tcu_wg + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_SS +- id: sgemm_tcu_wg-9 + via: blackbox + drivers: + - rtlsim + app: sgemm_tcu_wg + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_SS +- id: sgemm_tcu_wg-10 + via: blackbox + drivers: + - simx + app: sgemm_tcu_wg + configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 + +- id: sgemm_tcu_wg-fedp2k + via: blackbox + drivers: + - simx + app: sgemm_tcu_wg + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_FEDP2K -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS + - id: sgemm_tcu_wg_sp-1 via: blackbox drivers: @@ -146,54 +128,68 @@ tests: app: sgemm_tcu_wg_sp configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DWGMMA_NRC=8 -- id: sgemm_tcu_wg-8 +- id: sgemm_tcu_wg_sp-7 via: blackbox drivers: - simx - app: sgemm_tcu_wg + app: sgemm_tcu_wg_sp configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE - -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_SS -- id: sgemm_tcu_wg-9 + -DVX_CFG_TCU_SPARSE_ENABLE -DWGMMA_NRC=8 -DWGMMA_SS +- id: sgemm_tcu_wg_sp-8 via: blackbox drivers: - rtlsim - app: sgemm_tcu_wg + app: sgemm_tcu_wg_sp configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE - -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_SS -- id: sgemm_tcu_wg_sp-7 + -DVX_CFG_TCU_SPARSE_ENABLE -DWGMMA_NRC=8 -DWGMMA_SS + +- id: sgemm_tcu_wg_mx-baseline via: blackbox drivers: - simx - app: sgemm_tcu_wg_sp + app: sgemm_tcu_wg_mx configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE - -DVX_CFG_TCU_SPARSE_ENABLE -DWGMMA_NRC=8 -DWGMMA_SS -- id: sgemm_tcu_wg_sp-8 + -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS + +- id: sgemm_tcu_wg_sp-fedp2k-rs via: blackbox drivers: - - rtlsim + - simx app: sgemm_tcu_wg_sp configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE - -DVX_CFG_TCU_SPARSE_ENABLE -DWGMMA_NRC=8 -DWGMMA_SS -- id: sgemm_tcu-3 + -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_FEDP2K -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS + +- id: sgemm_tcu_wg_sp-fedp2k-ss via: blackbox drivers: - simx - app: sgemm_tcu - configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -- id: sgemm_tcu_sp-7 + app: sgemm_tcu_wg_sp + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_FEDP2K -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_SS + +- id: sgemm_tcu_wg_mx-fedp2k-rs via: blackbox drivers: - simx - app: sgemm_tcu_sp - configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -- id: sgemm_tcu_wg-10 + app: sgemm_tcu_wg_mx + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DVX_CFG_TCU_FEDP2K -DITYPE=mxfp8 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS + +- id: sgemm_tcu_wg_mx-fedp2k-ss via: blackbox drivers: - simx - app: sgemm_tcu_wg - configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE - -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 + app: sgemm_tcu_wg_mx + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DVX_CFG_TCU_FEDP2K -DITYPE=mxfp8 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_SS +- id: sgemm_tcu_wg_mx-fedp2k-nvfp4 + via: blackbox + drivers: + - simx + app: sgemm_tcu_wg_mx + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DVX_CFG_TCU_FEDP2K -DITYPE=nvfp4 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS -DTCU_MX_TLS # SimX<->rtlsim cycle parity on the WGMMA int8 datapath: same config on both # drivers, cycles must agree within tolerance (see model_parity.yaml). - id: model_parity-wgmma-int8 @@ -203,18 +199,70 @@ tests: configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_INT8_ENABLE -DITYPE=int8 -DOTYPE=int32 -DWGMMA_NRC=32 -# Dense WGMMA fp16 datapath, cycles-vs-baseline (same K=128 depth as the int8 -# parity twin above). -- id: perf_gate-wgmma-fp16 +# Dense WGMMA fp16 SS datapath, cycles-vs-baseline. +- id: perf_gate-wgmma-fp16-ss check: perf_gate via: blackbox app: sgemm_tcu_wg - args: -m 128 -n 128 -k 128 - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=32 + args: -m 64 -n 64 -k 64 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=32 -DWGMMA_SS + +- id: perf_gate-wgmma-fedp2k-rs + check: perf_gate + via: blackbox + app: sgemm_tcu_wg + args: -m64 -n64 -k64 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_FEDP2K -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS + +- id: model_parity-wgmma-fedp2k-rs + check: model_parity + known_issue: "cycle-parity gap ~8%: SimX WGMMA setup/tile-buffer timing is optimistic (instrs match)" + via: blackbox + app: sgemm_tcu_wg + args: -m64 -n64 -k64 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_FEDP2K -DVX_CFG_TCU_TYPE_DPI -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS + # Cooperative sparse WGMMA (default fp16 operands, 2:4 structured sparsity). -- id: perf_gate-wgmma-sparse +- id: perf_gate-wgmma-sp-rs check: perf_gate via: blackbox app: sgemm_tcu_wg_sp args: -m 128 -n 128 -k 128 - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DWGMMA_NRC=32 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DWGMMA_NRC=32 -DWGMMA_RS + +- id: model_parity-wgmma-mx-baseline + check: model_parity + via: blackbox + app: sgemm_tcu_wg_mx + args: -m64 -n64 -k64 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DVX_CFG_TCU_TYPE_DPI -DITYPE=mxfp8 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS + +- id: model_parity-wgmma-sp-fedp2k-ss + check: model_parity + # SimX plans a complete sparse WGMMA tile-buffer transaction while RTL + # refills the per-block A and shared B buffers stepwise. + known_issue: "cycle-parity gap ~11%: SimX sparse WGMMA tile-buffer timing is conservative (instrs match)" + via: blackbox + app: sgemm_tcu_wg_sp + args: -m64 -n64 -k64 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_FEDP2K -DVX_CFG_TCU_TYPE_DPI -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_SS + +- id: model_parity-wgmma-mx-fedp2k-rs + check: model_parity + via: blackbox + app: sgemm_tcu_wg_mx + args: -m64 -n64 -k64 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DVX_CFG_TCU_FEDP2K -DVX_CFG_TCU_TYPE_DPI -DITYPE=mxfp8 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS + +- id: perf_gate-wgmma-mx-fedp2k-rs + check: perf_gate + via: blackbox + app: sgemm_tcu_wg_mx + args: -m64 -n64 -k64 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_ISSUE_WIDTH=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_WGMMA_ENABLE + -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP8_ENABLE -DVX_CFG_TCU_FEDP2K -DVX_CFG_TCU_TYPE_DPI -DITYPE=mxfp8 -DOTYPE=fp32 -DWGMMA_NRC=8 -DWGMMA_RS diff --git a/docs/designs/tensor_core_wgmma_engine.md b/docs/designs/tensor_core_wgmma_engine.md index 429f0ffadb..19e6d15bf7 100644 --- a/docs/designs/tensor_core_wgmma_engine.md +++ b/docs/designs/tensor_core_wgmma_engine.md @@ -54,8 +54,9 @@ expands the macro into per-block micro-ops; each block executes a FEDP ## 2. Data types, opcodes, and configuration **Formats** ([`VX_tcu_pkg.sv:27-39`](../../hw/rtl/tcu/VX_tcu_pkg.sv#L27), -[`sw/common/tensor_cfg.h:25-66`](../../sw/common/tensor_cfg.h#L25)): -fp32, fp16, bf16, fp8 (e4m3), bf8 (e5m2), tf32, and integer i32/i8/u8/i4/u4. +[`sw/common/tensor_cfg.h`](../../sw/common/tensor_cfg.h)): +fp32, fp16, bf16, fp8 (e4m3), bf8 (e5m2), tf32, integer +i32/i8/u8/i4/u4, and block-scaled mxfp8/mxbf8/mxfp4/nvfp4. Per-format enables: `VX_CFG_TCU_TF32_ENABLE`, `_BF16_`, `_FP8_`, `_INT_` ([`VX_config.toml:238-241`](../../VX_config.toml#L238)). @@ -79,8 +80,9 @@ Sparsity is now a distinct opcode (`*_SP`) rather than a runtime flag. `VX_CFG_TCU_TYPE` selects the FEDP backend (`DPI`/`DSP`/`BHF`/`TFR`; default `TFR` for ASIC, `DSP` for synthesis, `DPI` when DPI is enabled); `VX_CFG_NUM_TCU_LANES = NUM_THREADS`; `VX_CFG_NUM_TCU_BLOCKS = ISSUE_WIDTH`; -`VX_CFG_TCU_SPARSE_ENABLE`; `VX_CFG_TCU_WGMMA_ENABLE`. - +`VX_CFG_TCU_SPARSE_ENABLE`; `VX_CFG_TCU_MX_ENABLE`; `VX_CFG_TCU_WGMMA_ENABLE`; +`VX_CFG_TCU_FEDP2K` doubles the dense WGMMA FEDP width and halves its +K-step count --- ## 3. RTL module inventory @@ -141,6 +143,37 @@ selected at first-uop from `desc[31:16]` [`VX_tcu_bbuf.sv:286-601`](../../hw/rtl/tcu/VX_tcu_bbuf.sv#L286)). The `Q+1` LMEM masters arbitrate to one port through `VX_mem_arb`. +**FEDP2K geometry and register operands.** The ordinary dense WGMMA FEDP +consumes `tcK` 32-bit words per A row and executes two K steps. With +`VX_CFG_TCU_FEDP2K`, `fedpK = 2*tcK` and dense WGMMA executes one K step. +For RS instructions, `rs1` carries the lower `tcK` A words and `rs2` +carries the upper `tcK` words. Sparse WGMMA remains compressed to `tcK` +words and therefore does not consume an upper-A `rs2` operand. SS obtains +both A halves from the tile buffer. + +**Descriptor transport and uop sequencing.** x10 holds the optional A +shared-memory descriptor and x11 holds the B descriptor. FEDP2K support is +a compile-time capability; within such a build, descriptor setup is selected +per decoded instruction: + +| WGMMA mode | First emitted uop | Compute `rs1` | Compute `rs2` | +|---|---|---|---| +| Dense RS + FEDP2K | Descriptor-only setup from x11 | Lower A | Upper A | +| Dense RS without FEDP2K | First compute, with x11 fused | A | x11 on first compute only | +| Sparse RS, including FEDP2K | First compute, with x11 fused | Compressed A | x11 on first compute only | +| SS, including FEDP2K | First compute, with x10/x11 fused | x10 on first compute only | x11 on first compute only | + +Thus `needs_setup = FEDP2K && RS && dense`. This is the only combination +where descriptor B and upper A compete for the same warp `rs2` vector. +The setup uop does not read x10 and has no FEDP writeback; it latches the B +descriptor and resets the tile-buffer transaction. Fused RS modes capture +x11 when the first compute uop fires. SS likewise presents x10/x11 directly +on its first compute uop, allowing tile-buffer refill to start without a +descriptor-capture bubble. The descriptors are latched when that uop fires. +Later compute uops reuse the descriptor latches. `fu_lock` marks the first +emitted uop in either sequence, while `is_first_uop` marks the first compute +uop, so lockstep and buffer allocation do not confuse setup with compute. + **Lock-step / warpgroup gate.** `VX_tcu_lockstep` enforces single-CTA occupancy of the shared B buffer: a single owner latch plus per-block `in_expansion_r` (set on the first sub-uop, cleared on the last). A block @@ -154,6 +187,14 @@ configured FEDP cell. Sparse 2:4 routes B through `VX_tcu_sp_mux` using the Accumulation walks k through the C accumulator register; results return via `VX_lane_gather` → `commit_if`. +For MX inputs, independent A and B scale arrays are preloaded by TCU_LD +into `VX_tcu_mx_meta`. The core selects scale bytes by logical row/column +and K block, including the wider FEDP2K span, and applies the format's +block size (32 elements for MXFP formats, 16 for NVFP4). Sparse addressing +accounts for the 2:4 logical-K expansion. MX integer-8 helper/datapath +variants are intentionally absent; supported MX formats use the dedicated +floating-point conversion path. + **TCU_LD path.** `INST_TCU_LD` is handled by `VX_tcu_agu`, which walks the metadata stride and issues to the **shared** `VX_lsu_scheduler` ([`hw/rtl/core/VX_lsu_scheduler.sv`](../../hw/rtl/core/VX_lsu_scheduler.sv), @@ -176,6 +217,10 @@ Operand load goes through channels (`load_lmem_word`); there is **no** `plan_a/plan_b`, `ready_a/ready_b`, `read_a/read_b`. The RTL and SimX share the same k-outer iteration order, per-block-A + shared-B structure, and lock-step deadlock contract, enabling SimX↔RTL cycle parity work. +The uop generator uses the same `FEDP2K && RS && dense` setup predicate as +RTL, including fused descriptor reads for SS and sparse RS. SimX also +mirrors the widened dense A operand, persistent descriptors, and separate +MX A/B metadata SRAMs and scale indexing. Perf CSRs: `TBUF_STALLS`, `TBUF_CACHE_HITS`, `LMEM_READS` ([`VX_types.toml:570-577`](../../VX_types.toml#L570)). diff --git a/hw/rtl/core/VX_scoreboard.sv b/hw/rtl/core/VX_scoreboard.sv index d9433482a0..88221c6e38 100644 --- a/hw/rtl/core/VX_scoreboard.sv +++ b/hw/rtl/core/VX_scoreboard.sv @@ -42,6 +42,7 @@ module VX_scoreboard import VX_gpu_pkg::*; #( // credit covers ops still in operand collection (not yet at the queue). wire [NUM_EX_UNITS-1:0] fu_issue; wire [NUM_EX_UNITS-1:0] fu_goingfull; + reg [NUM_EX_UNITS-1:0] fu_locked, fu_locked_n; // going-full (not full): a 1-slot guard band keeps outstanding <= queue // depth despite the registered suppress lag, so an issued op never stalls @@ -182,10 +183,11 @@ module VX_scoreboard import VX_gpu_pkg::*; #( wire [EX_BITS-1:0] ex_sel = ibuffer_fire ? ibuffer_if[w].data.ex_type : staging_if[w].data.ex_type; reg operands_ready_r; - // Readiness folds data hazards and FU-congestion into one flop; FU-lock - // is enforced downstream by masking the arbiter requests. + wire fu_lock_sel = ibuffer_fire ? ibuffer_if[w].data.fu_lock : staging_if[w].data.fu_lock; wire data_ready = ~((|regs_busy) || xregs_busy); - wire operands_ready_n = data_ready && ~fu_goingfull[ex_sel]; + wire operands_ready_n = data_ready + && ~fu_goingfull[ex_sel] + && ~(fu_locked_n[ex_sel] && fu_lock_sel); always @(posedge clk) begin if (reset) begin @@ -240,15 +242,8 @@ module VX_scoreboard import VX_gpu_pkg::*; #( wire [PER_ISSUE_WARPS-1:0][OUT_DATAW-1:0] arb_data_in; wire [PER_ISSUE_WARPS-1:0] arb_ready_in; - // FU lock: a sequence must not interleave with another warp at the same FU. - // fu_locked ('1 = open, one-hot = locked) gates arb_valid_in so only the lock - // holder is requested while it holds the lock. - reg [PER_ISSUE_WARPS-1:0] fu_locked; - for (genvar w = 0; w < PER_ISSUE_WARPS; ++w) begin : g_arb_data_in - // operands_ready carries data-hazard + FU-congestion; fu_locked adds the - // FU-lock gate so only the lock holder is requested during a sequence. - assign arb_valid_in[w] = staging_if[w].valid && operands_ready[w] && fu_locked[w]; + assign arb_valid_in[w] = staging_if[w].valid && operands_ready[w]; assign arb_data_in[w] = { staging_if[w].data.uuid, @@ -325,23 +320,28 @@ module VX_scoreboard import VX_gpu_pkg::*; #( assign fu_issue[e] = issue_fire && (issue_ex == EX_BITS'(e)); end - // Lock to the granted warp on acquire, hold across its sequence, reopen on - // release. arb_onehot is the granted warp, registered into fu_locked. wire issue_fu_lock = staging_fu_lock_vec[arb_index]; wire issue_fu_unlock = staging_fu_unlock_vec[arb_index]; - always @(posedge clk) begin - if (reset) begin - fu_locked <= '1; - end else if (issue_fire) begin + always @(*) begin + fu_locked_n = fu_locked; + if (issue_fire) begin if (issue_fu_lock && ~issue_fu_unlock) begin - fu_locked <= arb_onehot; - end else if (issue_fu_unlock) begin - fu_locked <= '1; + fu_locked_n[issue_ex] = 1'b1; + end else if (~issue_fu_lock && issue_fu_unlock) begin + fu_locked_n[issue_ex] = 1'b0; end end end + always @(posedge clk) begin + if (reset) begin + fu_locked <= '0; + end else begin + fu_locked <= fu_locked_n; + end + end + VX_elastic_buffer #( .DATAW (LOG_NUM_REQS + OUT_DATAW), .SIZE (`TO_OUT_BUF_SIZE(OUT_BUF)), diff --git a/hw/rtl/tcu/VX_tcu_abuf.sv b/hw/rtl/tcu/VX_tcu_abuf.sv index fd8cd63c3e..9a896b7a73 100644 --- a/hw/rtl/tcu/VX_tcu_abuf.sv +++ b/hw/rtl/tcu/VX_tcu_abuf.sv @@ -62,11 +62,13 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Execute-side observation (req_valid is already gated to WGMMA) input wire [NW_WIDTH-1:0] req_wid, input wire req_valid, + input wire req_setup, input wire [3:0] req_step_m, input wire [3:0] req_step_n, input wire [3:0] req_step_k, input wire [`VX_CFG_XLEN-1:0] req_desc_a, input wire req_a_is_smem, + input wire req_is_sparse, input wire [UUID_WIDTH-1:0] req_uuid, // LMEM bank-parallel read port @@ -74,7 +76,7 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Outputs output wire abuf_ready, - output wire [TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] abuf_rs1_data + output wire [TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] abuf_rs1_data ); `UNUSED_SPARAM (INSTANCE_ID) `UNUSED_VAR (req_wid) @@ -85,10 +87,14 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam BANK_SEL_BITS = $clog2(NUM_BANKS); localparam WORD_SIZE_LOG2 = $clog2(`VX_CFG_XLEN / 8); - localparam A_BLOCK_WORDS = TCU_TC_M * TCU_TC_K; + localparam A_BLOCK_WORDS = TCU_TC_M * TCU_WG_FEDP_K; + localparam A_BLOCK_WORDS_SP = TCU_TC_M * TCU_TC_K; localparam A_BLOCK_BANK_ROWS = (A_BLOCK_WORDS + NUM_BANKS - 1) / NUM_BANKS; + localparam A_BLOCK_BANK_ROWS_SP = (A_BLOCK_WORDS_SP + NUM_BANKS - 1) / NUM_BANKS; localparam BLOCK_WORDS_PADDED = A_BLOCK_BANK_ROWS * NUM_BANKS; + localparam BLOCK_WORDS_PADDED_SP = A_BLOCK_BANK_ROWS_SP * NUM_BANKS; localparam A_STRIPE_BANK_ROWS = TCU_WG_M_STEPS * A_BLOCK_BANK_ROWS; + localparam A_STRIPE_BANK_ROWS_SP = TCU_WG_M_STEPS * A_BLOCK_BANK_ROWS_SP; localparam A_STRIPE_WORDS = A_STRIPE_BANK_ROWS * NUM_BANKS; // XLEN ratio: each physical LMEM bank-row carries XLEN_RATIO logical // 32-bit bank-rows side-by-side. Smem layout is XLEN-independent, so the @@ -97,6 +103,7 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam BANK_ROW_WORDS = NUM_BANKS * XLEN_RATIO; // 32-bit words per LMEM bank-row localparam BANK_ROW_WORDS_LOG2 = $clog2(BANK_ROW_WORDS); localparam A_STRIPE_LMEM_ROWS = (A_STRIPE_WORDS + BANK_ROW_WORDS - 1) / BANK_ROW_WORDS; + localparam A_STRIPE_LMEM_ROWS_SP = (A_STRIPE_BANK_ROWS_SP * NUM_BANKS + BANK_ROW_WORDS - 1) / BANK_ROW_WORDS; // Row-major path issues one LMEM read per logical row (M_STEPS*TC_M rows). localparam A_TOTAL_ROWS = TCU_WG_M_STEPS * TCU_TC_M; localparam FETCH_CTR_W_BM = `CLOG2(A_STRIPE_LMEM_ROWS + 1); @@ -107,10 +114,8 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // per-row 32-bit-word stride. Cap at 14 bits (= 64KB smem / 4B). localparam LDM_W = 14; - // Canonical-config invariant: one A-block fits one (32-bit-equivalent) - // bank-row (TC_M*TC_K == NUM_BANKS). Non-canonical configs requiring - // A_SUB_BLOCKS packing in the output mux are not supported. - `STATIC_ASSERT (A_BLOCK_BANK_ROWS == 1, ("VX_tcu_abuf assumes one A-block per bank-row")) + // Canonical configs fit an A block in one bank-row normally and two + // bank-rows with dense FEDP2K. // ----------------------------------------------------------------------- // Resident slot state @@ -125,29 +130,28 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // alongside fetch_base_r and slot_step_k_r. logic slot_row_major_r; logic [LDM_W-1:0] slot_ldm_words_r; + logic slot_is_sparse_r; wire [`UP(K_STEPS_W)-1:0] req_step_k_trunc = `UP(K_STEPS_W)'(req_step_k); if (4 > K_STEPS_W) begin : g_step_k_upper_unused `UNUSED_VAR (req_step_k[3:`UP(K_STEPS_W)]) end - // The uop expander only reads rs1 (desc_a) on uop 0 of a WGMMA expansion. - // On non-first uops, req_desc_a is garbage and cannot participate in - // the residency check. - wire is_first_uop = (req_step_m == '0) && (req_step_n == '0) && (req_step_k == '0); + // The WGMMA wrapper supplies the setup-latched descriptor on compute uops. + wire is_first_compute_uop = (req_step_m == '0) && (req_step_n == '0) && (req_step_k == '0); `UNUSED_VAR (req_step_n) // only used in is_first_uop computation - // Force a refetch on every WGMMA's first uop. A cooperative-load + // Force a refetch on every WGMMA's first compute uop. A cooperative-load // pattern (K-tile loop rewrites A_warp_smem in place, issues a fresh // WGMMA with an unchanged descriptor) would incorrectly hit the cached // stripe and serve stale A data if residency were checked by desc_a alone. - // refetched_for_first_uop_r gates stripe_resident: cleared until the - // current WGMMA's first_uop fetch completes, then set; cleared again on - // the first non-first uop so the next WGMMA always triggers a fresh fetch. - reg refetched_for_first_uop_r; + // refetched_for_first_compute_r gates stripe_resident: cleared until the + // current WGMMA's first compute fetch completes, then set; cleared again + // on later compute uops so the next WGMMA always triggers a fresh fetch. + reg refetched_for_first_compute_r; wire stripe_resident = slot_valid_r && (slot_step_k_r == req_step_k_trunc) - && (!is_first_uop || refetched_for_first_uop_r); + && (!is_first_compute_uop || refetched_for_first_compute_r); // RS mode (a_from_smem=0): A from registers, abuf bypassed → always ready. wire need_smem = req_valid && req_a_is_smem; @@ -180,14 +184,12 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam STRIPE_STRIDE_BANK_ROWS = A_STRIPE_LMEM_ROWS; - // Use latched desc_a base on non-first uops (req_desc_a is garbage there - // because the uop expander gates the rs1 register read on uop 0 only). - // Without this, k-stripe-transition refills mid-WGMMA would compute the - // wrong fetch_base. + // Compute uops use the setup-latched descriptor. Without this, + // k-stripe-transition refills mid-WGMMA would compute the wrong fetch_base. wire [BANK_ADDR_WIDTH-1:0] effective_desc_a_row_base = - is_first_uop ? desc_a_row_base : slot_desc_a_row_base_r; + is_first_compute_uop ? desc_a_row_base : slot_desc_a_row_base_r; wire [LDM_W-1:0] effective_ldm_words = - is_first_uop ? desc_a_ldm_words : slot_ldm_words_r; + is_first_compute_uop ? desc_a_ldm_words : slot_ldm_words_r; wire effective_row_major = (effective_ldm_words != '0); // Block-major stripe base (unchanged): one fetch covers M_STEPS A-blocks. @@ -220,7 +222,8 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // row-major: A_TOTAL_ROWS per-row reads (one per logical row) wire [FETCH_CTR_W-1:0] target_fetches = slot_row_major_r ? FETCH_CTR_W'(A_TOTAL_ROWS) - : FETCH_CTR_W'(A_STRIPE_LMEM_ROWS); + : (slot_is_sparse_r ? FETCH_CTR_W'(A_STRIPE_LMEM_ROWS_SP) + : FETCH_CTR_W'(A_STRIPE_LMEM_ROWS)); wire all_requested = (req_ctr_r >= target_fetches); wire can_issue = in_fetch && !all_requested @@ -231,9 +234,10 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Per-row 32-bit-word offset (relative to fetch_base_r * BANK_ROW_WORDS). // Width = LDM (row stride) + counter + a bit for step_k*TC_K headroom. localparam ROW_OFF_W = LDM_W + FETCH_CTR_W + 4; + wire [3:0] row_k_words = slot_is_sparse_r ? 4'(TCU_TC_K) : 4'(TCU_WG_FEDP_K); wire [ROW_OFF_W-1:0] row_word_off_req = ROW_OFF_W'(req_ctr_r) * ROW_OFF_W'(slot_ldm_words_r) - + ROW_OFF_W'(slot_step_k_r) * ROW_OFF_W'(TCU_TC_K); + + ROW_OFF_W'(slot_step_k_r) * ROW_OFF_W'(row_k_words); wire [BANK_ADDR_WIDTH-1:0] row_lmem_addr = fetch_base_r + BANK_ADDR_WIDTH'(row_word_off_req >> BANK_ROW_WORDS_LOG2); @@ -264,63 +268,71 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( fetch_base_r <= '0; slot_row_major_r <= 1'b0; slot_ldm_words_r <= '0; - refetched_for_first_uop_r <= 1'b0; + slot_is_sparse_r <= 1'b0; + refetched_for_first_compute_r <= 1'b0; end else begin - // Inflight tracker (single outstanding request at a time) - if (tcu_lmem_if.rsp_valid) - req_inflight_r <= 1'b0; - if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) - req_inflight_r <= 1'b1; - - // Latch desc_a's row base + ldm_words on first uop of every WGMMA - // so non-first uops (k-stripe-transition refills) can reuse them - // without needing the gated req_desc_a bus. - if (req_valid && is_first_uop) begin - slot_desc_a_row_base_r <= desc_a_row_base; - slot_ldm_words_r <= desc_a_ldm_words; - slot_row_major_r <= (desc_a_ldm_words != '0); - end + if (req_setup) begin + fsm_state_r <= S_IDLE; + req_ctr_r <= '0; + rsp_ctr_r <= '0; + req_inflight_r <= 1'b0; + slot_valid_r <= 1'b0; + slot_fetching_r <= 1'b0; + refetched_for_first_compute_r <= 1'b0; + end else begin + // Inflight tracker (single outstanding request at a time) + if (tcu_lmem_if.rsp_valid) + req_inflight_r <= 1'b0; + if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) + req_inflight_r <= 1'b1; + + // Latch descriptor fields on the first compute uop. + if (req_valid && is_first_compute_uop) begin + slot_desc_a_row_base_r <= desc_a_row_base; + slot_ldm_words_r <= desc_a_ldm_words; + slot_row_major_r <= (desc_a_ldm_words != '0); + end - // refetched_for_first_uop_r: set when fetch completes for - // this WGMMA's first uop; cleared on the next non-first uop. - if (last_rsp && is_first_uop) - refetched_for_first_uop_r <= 1'b1; - else if (fire && !is_first_uop) - refetched_for_first_uop_r <= 1'b0; - - case (fsm_state_r) - S_IDLE: begin - if (alloc_en) begin - fsm_state_r <= S_FETCH; - slot_fetching_r <= 1'b1; - slot_valid_r <= 1'b0; - slot_step_k_r <= req_step_k_trunc; - // Row-major base is the A_warp start bank-row (no - // step_k offset; per-row arithmetic derives the - // exact LMEM addr). Block-major base is the stripe - // origin (step_k already factored in). - fetch_base_r <= effective_row_major - ? effective_desc_a_row_base - : stripe_base; - req_ctr_r <= '0; - rsp_ctr_r <= '0; - req_inflight_r <= 1'b0; + if (last_rsp && is_first_compute_uop) + refetched_for_first_compute_r <= 1'b1; + else if (fire && !is_first_compute_uop) + refetched_for_first_compute_r <= 1'b0; + + case (fsm_state_r) + S_IDLE: begin + if (alloc_en) begin + fsm_state_r <= S_FETCH; + slot_fetching_r <= 1'b1; + slot_valid_r <= 1'b0; + slot_step_k_r <= req_step_k_trunc; + slot_is_sparse_r <= req_is_sparse; + // Row-major base is the A_warp start bank-row (no + // step_k offset; per-row arithmetic derives the + // exact LMEM addr). Block-major base is the stripe + // origin (step_k already factored in). + fetch_base_r <= effective_row_major + ? effective_desc_a_row_base + : stripe_base; + req_ctr_r <= '0; + rsp_ctr_r <= '0; + req_inflight_r <= 1'b0; + end end - end - S_FETCH: begin - if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) - req_ctr_r <= req_ctr_r + FETCH_CTR_W'(1); - if (last_rsp) begin - fsm_state_r <= S_IDLE; - slot_fetching_r <= 1'b0; - slot_valid_r <= 1'b1; - req_inflight_r <= 1'b0; - end else if (tcu_lmem_if.rsp_valid) begin - rsp_ctr_r <= rsp_ctr_r + FETCH_CTR_W'(1); + S_FETCH: begin + if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) + req_ctr_r <= req_ctr_r + FETCH_CTR_W'(1); + if (last_rsp) begin + fsm_state_r <= S_IDLE; + slot_fetching_r <= 1'b0; + slot_valid_r <= 1'b1; + req_inflight_r <= 1'b0; + end else if (tcu_lmem_if.rsp_valid) begin + rsp_ctr_r <= rsp_ctr_r + FETCH_CTR_W'(1); + end end - end - default: fsm_state_r <= S_IDLE; - endcase + default: fsm_state_r <= S_IDLE; + endcase + end end end @@ -337,7 +349,7 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // within the LMEM response. wire [ROW_OFF_W-1:0] row_word_off_rsp = ROW_OFF_W'(rsp_ctr_r) * ROW_OFF_W'(slot_ldm_words_r) - + ROW_OFF_W'(slot_step_k_r) * ROW_OFF_W'(TCU_TC_K); + + ROW_OFF_W'(slot_step_k_r) * ROW_OFF_W'(row_k_words); wire [BANK_ROW_WORDS_LOG2:0] row_lane_rsp = (BANK_ROW_WORDS_LOG2+1)'( row_word_off_rsp & ROW_OFF_W'(BANK_ROW_WORDS - 1)); @@ -346,13 +358,13 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( storage_wren = '0; if (in_fetch && tcu_lmem_if.rsp_valid) begin if (slot_row_major_r) begin - // Row-major: write TC_K words of row rsp_ctr_r into - // storage[rsp_ctr_r * TC_K .. + TC_K). Source words start + // Row-major: write k_words of row rsp_ctr_r into + // storage[rsp_ctr_r * k_words .. + k_words). Source words start // at row_lane_rsp inside the LMEM response. - for (int k = 0; k < TCU_TC_K; ++k) begin - automatic int dst = int'(rsp_ctr_r) * TCU_TC_K + k; + for (int k = 0; k < TCU_WG_FEDP_K; ++k) begin + automatic int dst = int'(rsp_ctr_r) * int'(row_k_words) + k; automatic int src = int'(row_lane_rsp) + k; - if (dst < A_STRIPE_WORDS && src < BANK_ROW_WORDS) begin + if (k < int'(row_k_words) && dst < A_STRIPE_WORDS && src < BANK_ROW_WORDS) begin storage_wren[dst] = 1'b1; storage_wdata[dst * 32 +: 32] = tcu_lmem_if.rsp_data.data[src * 32 +: 32]; @@ -395,17 +407,27 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ); // ----------------------------------------------------------------------- - // Output: select A-block based on step_m, pass through as 32-bit words. - // Storage holds blocks at [m_blk * BLOCK_WORDS_PADDED .. +BLOCK_WORDS_PADDED). + // Output: select A-block based on step_m. Sparse data is stored compactly, + // then expanded into FEDP_K row stride with the unavailable tail lanes zero. // ----------------------------------------------------------------------- - logic [TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] rs1_mux; + logic [TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] rs1_mux; always_comb begin rs1_mux = '0; - for (int lane = 0; lane < TCU_BLOCK_CAP; ++lane) begin - if (lane < int'(A_BLOCK_WORDS)) begin - int src_idx; - src_idx = int'(req_step_m) * BLOCK_WORDS_PADDED + lane; + for (int lane = 0; lane < TCU_WG_A_DATA_SIZE; ++lane) begin + int src_idx; + if (req_is_sparse) begin + automatic int row = lane / int'(TCU_WG_FEDP_K); + automatic int k = lane % int'(TCU_WG_FEDP_K); + if (row < int'(TCU_TC_M) && k < int'(TCU_TC_K)) begin + src_idx = int'(req_step_m) * int'(BLOCK_WORDS_PADDED_SP) + + row * int'(TCU_TC_K) + + k; + if (src_idx < int'(A_STRIPE_WORDS)) + rs1_mux[lane] = `VX_CFG_XLEN'(storage_rdata[src_idx]); + end + end else if (lane < int'(A_BLOCK_WORDS)) begin + src_idx = int'(req_step_m) * int'(BLOCK_WORDS_PADDED) + lane; if (src_idx < int'(A_STRIPE_WORDS)) rs1_mux[lane] = `VX_CFG_XLEN'(storage_rdata[src_idx]); end diff --git a/hw/rtl/tcu/VX_tcu_bbuf.sv b/hw/rtl/tcu/VX_tcu_bbuf.sv index 61e0514bc1..4e9317385f 100644 --- a/hw/rtl/tcu/VX_tcu_bbuf.sv +++ b/hw/rtl/tcu/VX_tcu_bbuf.sv @@ -16,24 +16,25 @@ `ifdef VX_CFG_TCU_WGMMA_ENABLE // -// TB-shared B buffer (1 bank-row storage per slot). +// TB-shared B buffer (block-major SMEM, 1 bank-row storage). // -// Single instance per VX_tcu_unit. Holds the bank-row(s) of B that contain -// the current (step_k, step_n) block. All Q tcu_cores read from the same -// buffer (structural fan-out, not arbitrated). +// Single instance per VX_tcu_unit. Holds the bank-row of B that contains +// the current (step_k, step_n) block. All Q tcu_cores read from the +// same buffer (structural fan-out, not arbitrated). // -// Two B layouts, selected at runtime by req_is_sparse: +// For canonical configs where TC_K * TC_N < NUM_BANKS, one bank-row holds +// B_SUB_BLOCKS = NUM_BANKS / (TC_K * TC_N) consecutive (k,n) blocks. +// Refill key = {desc_b, bank_row_index} where +// bank_row_index = (step_k * N_STEPS + step_n) >> LG_B_SUB_BLOCKS. // -// Dense (block-major): one logical 32-bit bank-row holds B_SUB_BLOCKS -// consecutive (k,n) blocks. The bank-row is stored verbatim into slot A; -// tcu_core's b_off picks the block within at execute time. +// The bus to tcu_core carries the whole bank-row; tcu_core's b_off +// (= step_n & (B_SUB_BLOCKS-1) << LG_B_BS) selects within. // -// Sparse (flat candidate-pair): each (step_k, step_n) block occupies two -// contiguous physical bank-rows laid out in FEDP candidate-pair order -// (matches vx_tensor.h b_sp_flat_idx). The two rows are stored verbatim -// into slots A and B; a fixed read permutation (constant wiring) presents -// them to the FEDP as rs2[k_idx*TC_N*2 + n_in*2 + cand]. No transpose -// crossbar — store-and-read are both straight wires. +// Block-major within-block layout: +// B_smem[(k*N_STEPS+n) * BLOCK_WORDS + j*(TC_K*i_ratio) + k_in_elem] +// Each 32-bit word packs i_ratio K-elements at one (j, k_word) cell. +// This matches tcu_core's `b_col[k] = rs2_data[b_off + j*TC_K + k]` indexing, +// so bbuf is word pass-through (no format-aware extraction needed here). // module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( @@ -52,6 +53,7 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // TB-level uop observation (req_valid is already gated to WGMMA at wrapper) input wire req_valid, + input wire req_setup, input wire req_is_first_uop, input wire req_is_sparse, input wire [3:0] req_step_m, @@ -76,7 +78,8 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam BANK_SEL_BITS = $clog2(NUM_BANKS); localparam WORD_SIZE_LOG2 = $clog2(`VX_CFG_XLEN / 8); - localparam B_BLOCK_WORDS = TCU_TC_K * TCU_TC_N; + localparam B_BLOCK_WORDS = TCU_WG_FEDP_K * TCU_TC_N; + localparam B_BLOCK_WORDS_SP = TCU_TC_K * TCU_TC_N; localparam B_BUF_WORDS = NUM_BANKS; // storage per slot // (= 1 logical 32-bit bank-row) localparam LG_B_SUB_BLOCKS = $clog2(TCU_WG_B_SUB_BLOCKS); @@ -85,11 +88,33 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam XLEN_RATIO = `VX_CFG_XLEN / 32; localparam LG_XLEN_RATIO = (XLEN_RATIO > 1) ? $clog2(XLEN_RATIO) : 0; + localparam LOGICAL_ROWS_PER_BLK = TCU_WG_B_BLOCK_SIZE_SP / NUM_BANKS; + `STATIC_ASSERT (LOGICAL_ROWS_PER_BLK == 1 || LOGICAL_ROWS_PER_BLK == 2, + ("flat sparse-B supports 1 or 2 logical bank-rows per block")) + localparam bit SPARSE_TWO_SLOT = (LOGICAL_ROWS_PER_BLK == 2); + localparam bit SPARSE_TWO_FETCH = SPARSE_TWO_SLOT && (XLEN_RATIO == 1); + // Canonical-config invariant: 1 logical (32-bit-equivalent) bank-row // holds B_SUB_BLOCKS blocks (the smem layout is XLEN-independent). `STATIC_ASSERT (B_BLOCK_WORDS * TCU_WG_B_SUB_BLOCKS == NUM_BANKS, ("VX_tcu_bbuf assumes one bank-row per B_SUB_BLOCKS blocks")) + // K-major (= row-major SMEM where K runs along contiguous bytes; the + // WGMMA SS-descriptor's canonical layout) fetch path. Engaged + // when desc_b's stride field (bits [31:16]) is non-zero. Performs + // TCU_TC_N per-N-row LMEM reads per (step_k, step_n) WGMMA uop, writing + // TCU_WG_FEDP_K 32-bit words per row into storage at the b_off offset + // tcu_core's `b_off + j*tcK + k` indexing will read. Mirrors + // VX_tcu_abuf.sv's row-major fetch for A. + localparam LDM_W = 14; + localparam BANK_ROW_WORDS = NUM_BANKS * XLEN_RATIO; + localparam BANK_ROW_WORDS_LOG2= $clog2(BANK_ROW_WORDS); + localparam LG_B_BLOCK_WORDS = $clog2(B_BLOCK_WORDS); + localparam KM_CTR_W = $clog2(TCU_TC_N + 1); + // 32-bit-word offset width for K-major address arithmetic. + // step_n × TCU_TC_N + j has ≤ 4+4 bits; × ldm_words adds LDM_W. + localparam KM_OFF_W = LDM_W + 4 + 4; + // ----------------------------------------------------------------------- // Block-index compute (variable N_STEPS via cd_nregs). // K_STEPS=2 always; N_STEPS=4/8/16 for cd_nregs=0/1/2 (NRC=8/16/32). @@ -103,13 +128,18 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( default: block_index = {req_step_k[0], req_step_n[3:0]}; // N_STEPS=16 endcase end - // K_STEPS=2 always → only req_step_k[0] selects the block; upper bits unused. - `UNUSED_VAR (req_step_k[3:1]) + if (4 > 1) begin : g_step_k_upper_unused + `UNUSED_VAR (req_step_k[3:1]) + end - // Dense LMEM bank-row offset: one logical 32-bit bank-row holds - // B_SUB_BLOCKS dense blocks; XLEN_RATIO of them fit in one physical - // bank-row, so the LMEM addr advances every (B_SUB_BLOCKS * XLEN_RATIO) - // blocks. + // LMEM bank-row offset. + // + // Dense: one logical 32-bit bank-row holds B_SUB_BLOCKS dense + // blocks; XLEN_RATIO of them fit in one physical bank-row. + // So advance the LMEM addr every + // (B_SUB_BLOCKS * XLEN_RATIO) blocks. + // Sparse blocks use the flat producer layout and can span one or two + // logical 32-bit bank rows. localparam TOTAL_SHIFT = LG_B_SUB_BLOCKS + LG_XLEN_RATIO; wire [4:0] dense_offset = (TOTAL_SHIFT == 0) ? block_index @@ -123,81 +153,81 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ? '0 : SUB_HALF_W'(({27'b0, block_index} >> LG_B_SUB_BLOCKS) & ((1 << LG_XLEN_RATIO) - 1)); + localparam SP_LOGROW_W = 6; + wire [SP_LOGROW_W-1:0] sp_logrow_a = + SP_LOGROW_W'(block_index) * SP_LOGROW_W'(LOGICAL_ROWS_PER_BLK); + wire [SP_LOGROW_W-1:0] sp_logrow_b = sp_logrow_a + SP_LOGROW_W'(1); + wire [SUB_HALF_W-1:0] sp_sub_a = (LG_XLEN_RATIO == 0) ? '0 + : SUB_HALF_W'(sp_logrow_a & SP_LOGROW_W'((1 << LG_XLEN_RATIO) - 1)); + wire [SUB_HALF_W-1:0] sp_sub_b = (LG_XLEN_RATIO == 0) ? '0 + : SUB_HALF_W'(sp_logrow_b & SP_LOGROW_W'((1 << LG_XLEN_RATIO) - 1)); + // ----------------------------------------------------------------------- - // Descriptor base address + // Address compute (block-major) // ----------------------------------------------------------------------- localparam DESC_ADDR_W = BANK_ADDR_WIDTH + BANK_SEL_BITS; wire [DESC_ADDR_W-1:0] desc_b_word_base = DESC_ADDR_W'(req_desc_b[15:0] >> WORD_SIZE_LOG2); wire [BANK_ADDR_WIDTH-1:0] desc_b_row_base = desc_b_word_base[BANK_SEL_BITS +: BANK_ADDR_WIDTH]; - `UNUSED_VAR (req_desc_b[`VX_CFG_XLEN-1:16]) + // desc_b's upper 16 bits encode the per-row byte stride (WGMMA + // SS-descriptor `ldm`). Non-zero stride selects the K-major fetch path. + wire [LDM_W-1:0] desc_b_ldm_words = LDM_W'(req_desc_b[31:16] >> 2); + if (`VX_CFG_XLEN > 32) begin : g_desc_b_upper_unused + `UNUSED_VAR (req_desc_b[`VX_CFG_XLEN-1:32]) + end if (BANK_SEL_BITS > 0) begin : g_addr_lsb_unused `UNUSED_VAR (desc_b_word_base[BANK_SEL_BITS-1:0]) end // ----------------------------------------------------------------------- // Resident slots - // slot A: dense bank-row, or sparse flat bank-row 0 - // slot B: sparse flat bank-row 1 (sparse only) + // slot A: holds k_blk=0 bank-row (dense uses this exclusively) + // slot B: holds k_blk=1 bank-row (sparse only) // ----------------------------------------------------------------------- logic slot_a_valid_r; logic [BANK_ADDR_WIDTH-1:0] slot_a_addr_r; logic [BANK_ADDR_WIDTH-1:0] slot_desc_b_row_base_r; logic slot_fetching_r; - logic [SUB_HALF_W-1:0] slot_a_sub_half_r; // dense XLEN>32 half select + // Dense within-physical-bank-row half (XLEN>32 only). + logic [SUB_HALF_W-1:0] slot_a_sub_half_r; + // Second slot for sparse blocks spanning two logical rows. logic slot_b_valid_r; logic [BANK_ADDR_WIDTH-1:0] slot_b_addr_r; - logic [SUB_HALF_W-1:0] slot_b_sub_half_r; // sparse XLEN>32 second-row half + logic [SUB_HALF_W-1:0] slot_b_sub_half_r; + // Mode the slot pair was filled under (sparse vs dense). On mode + // transition for the same warpgroup we must refill. logic slot_is_sparse_r; - - // req_desc_b is only valid on the first uop of a WGMMA expansion; latch - // desc_b on first uop and use the latched base for subsequent uops. + // K-major mode + per-WGMMA latched fields. Descriptor fields are latched + // on setup; per-compute step fields are latched when a fetch is allocated. + logic slot_row_major_r; + logic [LDM_W-1:0] slot_ldm_words_r; + logic [3:0] slot_step_k_r; + logic [3:0] slot_step_n_r; + // K-major multi-fetch counters (count up to TCU_TC_N requests / responses + // per WGMMA uop, one per N-row of the (step_k, step_n) block). + logic [KM_CTR_W-1:0] km_req_ctr_r; + logic [KM_CTR_W-1:0] km_rsp_ctr_r; + + // The WGMMA wrapper supplies the setup-latched descriptor on compute uops. `UNUSED_VAR (req_step_m) - wire is_first_uop = req_is_first_uop; wire [BANK_ADDR_WIDTH-1:0] effective_desc_b_row_base = - is_first_uop ? desc_b_row_base : slot_desc_b_row_base_r; - - // ----------------------------------------------------------------------- - // Fetch addresses - // ----------------------------------------------------------------------- + req_is_first_uop ? desc_b_row_base : slot_desc_b_row_base_r; + // K-major slot fields (slot_row_major_r / slot_ldm_words_r / + // slot_step_k_r / slot_step_n_r) are latched at alloc_en — see the + // always_ff below. The K-major addressing arithmetic reads them + // directly. desc_b_ldm_words on the bus is used only for setup. + // Per-mode fetch addresses. wire [BANK_ADDR_WIDTH-1:0] fetch_addr_dense = effective_desc_b_row_base + BANK_ADDR_WIDTH'(dense_offset); - - // Flat sparse-B: each (step_k, step_n) block occupies LOGICAL_ROWS_PER_BLK - // consecutive logical 32-bit bank-rows in FEDP candidate-pair order (matches - // vx_tensor.h b_sp_flat_idx). This count is XLEN-independent — the smem - // layout is identical for any XLEN. Logical rows are packed XLEN_RATIO per - // physical LMEM bank-row, so logical row L lives at physical row - // (L >> LG_XLEN_RATIO), sub-half (L & (XLEN_RATIO-1)) — exactly as the dense - // path maps its logical rows. slot A holds the block's logical row 0; slot B - // (two-row blocks) holds logical row 1. - localparam LOGICAL_ROWS_PER_BLK = TCU_WG_B_BLOCK_SIZE_SP / NUM_BANKS; - `STATIC_ASSERT (LOGICAL_ROWS_PER_BLK == 1 || LOGICAL_ROWS_PER_BLK == 2, - ("flat sparse-B supports 1 or 2 logical bank-rows per block")) - // slot B feeds the read mux whenever a block spans two logical rows. - localparam bit SPARSE_TWO_SLOT = (LOGICAL_ROWS_PER_BLK == 2); - // A second LMEM fetch is needed only when those two logical rows fall in - // different physical rows — i.e. XLEN_RATIO==1 (XLEN=32). At XLEN>=64 the two - // rows share one physical row, filled as two sub-halves of a single fetch. - localparam bit SPARSE_TWO_FETCH = SPARSE_TWO_SLOT && (XLEN_RATIO == 1); - - // Logical-row index of each slot's data → physical row + within-row sub-half. - localparam SP_LOGROW_W = 6; // block_index (5b) * LOGICAL_ROWS_PER_BLK (<=2) - wire [SP_LOGROW_W-1:0] sp_logrow_a = SP_LOGROW_W'(block_index) * SP_LOGROW_W'(LOGICAL_ROWS_PER_BLK); - wire [SP_LOGROW_W-1:0] sp_logrow_b = sp_logrow_a + SP_LOGROW_W'(1); - wire [SUB_HALF_W-1:0] sp_sub_a = (LG_XLEN_RATIO == 0) ? '0 - : SUB_HALF_W'(sp_logrow_a & SP_LOGROW_W'((1 << LG_XLEN_RATIO) - 1)); - wire [SUB_HALF_W-1:0] sp_sub_b = (LG_XLEN_RATIO == 0) ? '0 - : SUB_HALF_W'(sp_logrow_b & SP_LOGROW_W'((1 << LG_XLEN_RATIO) - 1)); - wire [BANK_ADDR_WIDTH-1:0] fetch_addr_a_flat = + wire [BANK_ADDR_WIDTH-1:0] fetch_addr_a_sparse = effective_desc_b_row_base + BANK_ADDR_WIDTH'(sp_logrow_a >> LG_XLEN_RATIO); - wire [BANK_ADDR_WIDTH-1:0] fetch_addr_b_flat = + wire [BANK_ADDR_WIDTH-1:0] fetch_addr_b_sparse = effective_desc_b_row_base + BANK_ADDR_WIDTH'(sp_logrow_b >> LG_XLEN_RATIO); wire [BANK_ADDR_WIDTH-1:0] fetch_addr_a = - req_is_sparse ? fetch_addr_a_flat : fetch_addr_dense; + req_is_sparse ? fetch_addr_a_sparse : fetch_addr_dense; wire bank_row_resident_dense = slot_a_valid_r && !slot_is_sparse_r @@ -206,22 +236,76 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( wire bank_row_resident_sparse = slot_a_valid_r && slot_is_sparse_r - && (slot_a_addr_r == fetch_addr_a_flat) + && (slot_a_addr_r == fetch_addr_a_sparse) && (slot_a_sub_half_r == sp_sub_a) - && (!SPARSE_TWO_SLOT || (slot_b_valid_r && (slot_b_addr_r == fetch_addr_b_flat) + && (!SPARSE_TWO_SLOT || (slot_b_valid_r + && (slot_b_addr_r == fetch_addr_b_sparse) && (slot_b_sub_half_r == sp_sub_b))); - wire bank_row_resident = req_is_sparse ? bank_row_resident_sparse - : bank_row_resident_dense; + // K-major residency (dense + sparse): same (step_k, step_n) already in + // the slot pair. For sparse, both slots must be filled because sparse B + // is split across the two dense candidate streams. + wire bank_row_resident_kmajor = + slot_a_valid_r && slot_row_major_r + && (!slot_is_sparse_r || slot_b_valid_r) + && (slot_step_k_r == req_step_k) + && (slot_step_n_r == req_step_n); + + // Block-major residency is the existing dense/sparse branch; K-major + // overrides it when the first compute descriptor or latched slot mode + // selects K-major. + wire req_wants_kmajor = + req_is_first_uop ? (desc_b_ldm_words != '0) : slot_row_major_r; + + wire bank_row_resident = req_wants_kmajor + ? bank_row_resident_kmajor + : (req_is_sparse ? bank_row_resident_sparse : bank_row_resident_dense); wire need_fetch = req_valid && !bank_row_resident; wire alloc_en = need_fetch && !slot_fetching_r; assign bbuf_ready = !req_valid || bank_row_resident; + // ----------------------------------------------------------------------- + // K-major address generation + // ----------------------------------------------------------------------- + // For row r (= km_req_ctr_r) of the current (step_k, step_n) block: + // word_off = (step_n × TCU_TC_N + r) × ldm_words + step_k × k_words + // bank_row = base + (word_off >> log2(BANK_ROW_WORDS)) + // lane = word_off & (BANK_ROW_WORDS - 1) + // k_words is fedpK for dense and tcK for sparse + + wire [3:0] km_k_words = slot_is_sparse_r ? 4'(TCU_TC_K) : 4'(TCU_WG_FEDP_K); + + wire [KM_OFF_W-1:0] km_word_off_req = + KM_OFF_W'(slot_step_n_r) * KM_OFF_W'(TCU_TC_N) * KM_OFF_W'(slot_ldm_words_r) + + KM_OFF_W'(km_req_ctr_r) * KM_OFF_W'(slot_ldm_words_r) + + KM_OFF_W'(slot_step_k_r) * KM_OFF_W'(km_k_words); + wire [BANK_ADDR_WIDTH-1:0] km_lmem_addr = + slot_desc_b_row_base_r + BANK_ADDR_WIDTH'(km_word_off_req >> BANK_ROW_WORDS_LOG2); + + wire [KM_OFF_W-1:0] km_word_off_rsp = + KM_OFF_W'(slot_step_n_r) * KM_OFF_W'(TCU_TC_N) * KM_OFF_W'(slot_ldm_words_r) + + KM_OFF_W'(km_rsp_ctr_r) * KM_OFF_W'(slot_ldm_words_r) + + KM_OFF_W'(slot_step_k_r) * KM_OFF_W'(km_k_words); + wire [BANK_ROW_WORDS_LOG2:0] km_lane_rsp = (BANK_ROW_WORDS_LOG2+1)'( + km_word_off_rsp & KM_OFF_W'(BANK_ROW_WORDS - 1)); + + // Storage offset where this K-major block lands (matches tcu_core's + // b_off = (step_n & (B_SUB_BLOCKS-1)) << LG_B_BLOCK_WORDS so the + // FEDP's `rs2[b_off + j*tcK + k]` reads what we wrote). + localparam BUF_OFF_W = $clog2(B_BUF_WORDS); + wire [BUF_OFF_W:0] km_b_off; + if (LG_B_SUB_BLOCKS > 0) begin : g_km_b_off + assign km_b_off = (BUF_OFF_W+1)'(slot_step_n_r[LG_B_SUB_BLOCKS-1:0]) + << LG_B_BLOCK_WORDS; + end else begin : g_km_b_off_zero + assign km_b_off = '0; + end + // ----------------------------------------------------------------------- // Fetch FSM - // S_IDLE → S_FETCH_A → S_IDLE (dense) - // S_IDLE → S_FETCH_A → S_FETCH_B → S_IDLE (sparse) + // S_IDLE → S_FETCH_A → S_IDLE (dense) + // S_IDLE → S_FETCH_A → S_FETCH_B → S_IDLE (sparse) // ----------------------------------------------------------------------- typedef enum logic [1:0] { @@ -236,12 +320,24 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( wire in_fetch = in_fetch_a || in_fetch_b; logic req_inflight_r; - // Single outstanding request per fetch state (req_inflight_r gates re-fire). - wire can_issue = in_fetch && !req_inflight_r; - wire last_rsp = in_fetch && tcu_lmem_if.rsp_valid; - + // K-major: counter-gated multi-fire (TC_N requests per uop, single + // outstanding). Block-major: original single-fire-per-state behavior + // (km_req_ctr_r is reset on alloc; goes 0 → 1 for one block-major fire). + wire km_more_to_request = + slot_row_major_r ? (km_req_ctr_r < KM_CTR_W'(TCU_TC_N)) + : (km_req_ctr_r == '0); + wire can_issue = in_fetch && !req_inflight_r && km_more_to_request; + wire km_final_rsp = slot_row_major_r + && tcu_lmem_if.rsp_valid + && (km_rsp_ctr_r == KM_CTR_W'(TCU_TC_N - 1)); + wire bm_final_rsp = !slot_row_major_r && tcu_lmem_if.rsp_valid; + wire last_rsp = in_fetch && (km_final_rsp || bm_final_rsp); + + // Issue address: K-major path uses km_lmem_addr (re-derived per + // km_req_ctr_r); block-major path uses the original slot_a/b addr. wire [BANK_ADDR_WIDTH-1:0] active_lmem_addr = - in_fetch_b ? slot_b_addr_r : slot_a_addr_r; + slot_row_major_r ? km_lmem_addr + : (in_fetch_b ? slot_b_addr_r : slot_a_addr_r); assign tcu_lmem_if.req_valid = can_issue; assign tcu_lmem_if.req_data.rw = 1'b0; @@ -267,88 +363,119 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( slot_a_sub_half_r <= '0; slot_b_sub_half_r <= '0; slot_is_sparse_r <= 1'b0; + slot_row_major_r <= 1'b0; + slot_ldm_words_r <= '0; + slot_step_k_r <= '0; + slot_step_n_r <= '0; + km_req_ctr_r <= '0; + km_rsp_ctr_r <= '0; end else begin - if (tcu_lmem_if.rsp_valid) - req_inflight_r <= 1'b0; - if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) - req_inflight_r <= 1'b1; - - // Latch desc_b base on first uop so non-first uops re-derive - // fetch_addr without the gated req_desc_b bus. - if (req_valid && is_first_uop) - slot_desc_b_row_base_r <= desc_b_row_base; - - case (fsm_state_r) - S_IDLE: begin - if (alloc_en) begin - fsm_state_r <= S_FETCH_A; - slot_fetching_r <= 1'b1; - slot_a_valid_r <= 1'b0; - slot_b_valid_r <= 1'b0; - slot_a_addr_r <= fetch_addr_a; - slot_a_sub_half_r <= req_is_sparse ? sp_sub_a : dense_sub_half; - // slot_b is read only in sparse two-row mode. - slot_b_addr_r <= fetch_addr_b_flat; - slot_b_sub_half_r <= sp_sub_b; - slot_is_sparse_r <= req_is_sparse; - req_inflight_r <= 1'b0; - end + if (req_setup) begin + slot_a_valid_r <= 1'b0; + slot_b_valid_r <= 1'b0; + slot_fetching_r <= 1'b0; + fsm_state_r <= S_IDLE; + req_inflight_r <= 1'b0; + km_req_ctr_r <= '0; + km_rsp_ctr_r <= '0; + end else begin + if (tcu_lmem_if.rsp_valid) + req_inflight_r <= 1'b0; + if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) + req_inflight_r <= 1'b1; + + // Latch descriptor fields on the first compute uop. + if (req_valid && req_is_first_uop) begin + slot_desc_b_row_base_r <= desc_b_row_base; + slot_ldm_words_r <= desc_b_ldm_words; + slot_row_major_r <= (desc_b_ldm_words != '0); end - S_FETCH_A: begin - if (last_rsp) begin - slot_a_valid_r <= 1'b1; - req_inflight_r <= 1'b0; - // Two-row sparse at XLEN=32 needs a second physical fetch - // into slot B. At XLEN>=64 the second logical row is the - // other sub-half of this same response (slot B is filled - // alongside slot A below), so the block is complete now. - // One-row sparse and dense are complete after slot A. - if (slot_is_sparse_r && SPARSE_TWO_FETCH) begin - fsm_state_r <= S_FETCH_B; - end else begin - if (slot_is_sparse_r && SPARSE_TWO_SLOT) - slot_b_valid_r <= 1'b1; + + // K-major req/rsp counters advance independently of FSM state + // (single-outstanding still enforced via req_inflight_r). + if (slot_row_major_r && tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) + km_req_ctr_r <= km_req_ctr_r + KM_CTR_W'(1); + if (slot_row_major_r && tcu_lmem_if.rsp_valid && !km_final_rsp) + km_rsp_ctr_r <= km_rsp_ctr_r + KM_CTR_W'(1); + + case (fsm_state_r) + S_IDLE: begin + if (alloc_en) begin + fsm_state_r <= S_FETCH_A; + slot_fetching_r <= 1'b1; + slot_a_valid_r <= 1'b0; + slot_b_valid_r <= 1'b0; + slot_a_addr_r <= fetch_addr_a; + slot_a_sub_half_r <= req_is_sparse ? sp_sub_a : dense_sub_half; + slot_b_addr_r <= fetch_addr_b_sparse; + slot_b_sub_half_r <= sp_sub_b; + slot_is_sparse_r <= req_is_sparse; + // K-major per-compute slot state for this + // (step_k, step_n) block. + slot_step_k_r <= req_step_k; + slot_step_n_r <= req_step_n; + // slot_row_major_r is latched from the setup + // descriptor and reused for all compute refills. + km_req_ctr_r <= '0; + km_rsp_ctr_r <= '0; + req_inflight_r <= 1'b0; + end + end + S_FETCH_A: begin + if (last_rsp) begin + slot_a_valid_r <= 1'b1; + req_inflight_r <= 1'b0; + if (slot_is_sparse_r && slot_row_major_r) begin + slot_b_valid_r <= 1'b1; + fsm_state_r <= S_IDLE; + slot_fetching_r <= 1'b0; + end else if (slot_is_sparse_r && SPARSE_TWO_FETCH) begin + fsm_state_r <= S_FETCH_B; + end else begin + if (slot_is_sparse_r && SPARSE_TWO_SLOT) + slot_b_valid_r <= 1'b1; + fsm_state_r <= S_IDLE; + slot_fetching_r <= 1'b0; + end + end + end + S_FETCH_B: begin + if (last_rsp) begin fsm_state_r <= S_IDLE; slot_fetching_r <= 1'b0; + slot_b_valid_r <= 1'b1; + req_inflight_r <= 1'b0; end end - end - S_FETCH_B: begin - if (last_rsp) begin - fsm_state_r <= S_IDLE; - slot_fetching_r <= 1'b0; - slot_b_valid_r <= 1'b1; - req_inflight_r <= 1'b0; - end - end - default: fsm_state_r <= S_IDLE; - endcase + default: fsm_state_r <= S_IDLE; + endcase + end end end // ----------------------------------------------------------------------- - // Storage (LUTRAM): two slots × NUM_BANKS 32-bit words, written verbatim. - // Dense: slot A holds one logical 32-bit bank-row (picked by sub_half - // at XLEN>32, the lower NUM_BANKS otherwise). - // Sparse: slot A holds flat bank-row 0, slot B holds flat bank-row 1. - // tcu_core's b_off (dense) and the flat read permutation (sparse) select - // the operands at execute time. + // Storage (LUTRAM): two slots × NUM_BANKS 32-bit words. + // slot A is always written on FETCH_A response (dense + sparse). + // slot B is written on FETCH_B response (sparse only). + // + // Dense write: copies one logical 32-bit bank-row (NUM_BANKS words) + // from the physical response — picked by sub_half at + // XLEN>32, the lower NUM_BANKS otherwise. tcu_core's + // b_off then picks within those words at execute time. + // + // Sparse write: extracts the compressed candidate lanes from one or two + // dense B blocks and stores them in the compact rs2 order. // ----------------------------------------------------------------------- logic [B_BUF_WORDS*32-1:0] storage_a_wdata, storage_b_wdata; logic [B_BUF_WORDS-1:0] storage_a_wren, storage_b_wren; - // 32-bit-word offset of each slot's logical bank-row within the physical - // LMEM response (NUM_BANKS * XLEN_RATIO words). Dense and sparse both index - // by their resolved sub-half (= 0 for dense/sparse one-row and at XLEN=32). + // Per-slot logical-row offsets within a physical LMEM response. localparam OFF_W = $clog2(NUM_BANKS * XLEN_RATIO) + 1; wire [OFF_W-1:0] a_off_words = OFF_W'(slot_a_sub_half_r) * OFF_W'(NUM_BANKS); wire [OFF_W-1:0] b_off_words = OFF_W'(slot_b_sub_half_r) * OFF_W'(NUM_BANKS); - - // At XLEN>=64 a two-row sparse block's second logical row is the other - // sub-half of slot A's response, so slot B is filled during the slot-A fetch - // (no separate S_FETCH_B). At XLEN=32 it is a distinct fetch (in_fetch_b). - wire sparse_b_from_a = SPARSE_TWO_SLOT && !SPARSE_TWO_FETCH && slot_is_sparse_r; + wire sparse_b_from_a = SPARSE_TWO_SLOT && !SPARSE_TWO_FETCH + && slot_is_sparse_r && !slot_row_major_r; always_comb begin storage_a_wdata = '0; @@ -356,22 +483,69 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( storage_b_wdata = '0; storage_b_wren = '0; if (tcu_lmem_if.rsp_valid) begin - for (int b = 0; b < B_BUF_WORDS; ++b) begin - if (in_fetch_a) begin - storage_a_wren[b] = 1'b1; - storage_a_wdata[b * 32 +: 32] = - tcu_lmem_if.rsp_data.data[(int'(a_off_words) + b) * 32 +: 32]; - if (sparse_b_from_a) begin + if (slot_row_major_r && slot_is_sparse_r) begin + // K-major sparse: per N-row r (= km_rsp_ctr_r = column j) the + // response carries this column's K dimension z-major with the + // two sparse candidates adjacent — word (z*2 + cand) holds the + // FEDP's bword{cand} for K-step z. + // The FEDP reads rs2[k_idx*(TC_N*2) + j*2 + cand] with k_idx=z; + // slot_a||slot_b splits at B_BLOCK_WORDS_SP (= TC_K*TC_N). + // Drive each (z,cand) word to its exact flat target and route by slot. + for (int z = 0; z < TCU_TC_K; ++z) begin + for (int c = 0; c < 2; ++c) begin + automatic int src = int'(km_lane_rsp) + z * 2 + c; + automatic int tgt = z * (TCU_TC_N * 2) + int'(km_rsp_ctr_r) * 2 + c; + if (src < (NUM_BANKS * XLEN_RATIO) && in_fetch_a) begin + if (tgt < int'(B_BLOCK_WORDS_SP)) begin + if (tgt < B_BUF_WORDS) begin + storage_a_wren[tgt] = 1'b1; + storage_a_wdata[tgt * 32 +: 32] = + tcu_lmem_if.rsp_data.data[src * 32 +: 32]; + end + end else begin + automatic int tgt_b = tgt - int'(B_BLOCK_WORDS_SP); + if (tgt_b < B_BUF_WORDS) begin + storage_b_wren[tgt_b] = 1'b1; + storage_b_wdata[tgt_b * 32 +: 32] = + tcu_lmem_if.rsp_data.data[src * 32 +: 32]; + end + end + end + end + end + end else if (slot_row_major_r) begin + // K-major dense: write fedpK words for this row (km_rsp_ctr_r) + // into storage[km_b_off + km_rsp_ctr_r * fedpK .. + fedpK), + // sourced from the response at lane km_lane_rsp. + for (int k = 0; k < TCU_WG_FEDP_K; ++k) begin + automatic int dst = int'(km_b_off) + + int'(km_rsp_ctr_r) * TCU_WG_FEDP_K + + k; + automatic int src = int'(km_lane_rsp) + k; + if (dst < B_BUF_WORDS && src < (NUM_BANKS * XLEN_RATIO) && in_fetch_a) begin + storage_a_wren[dst] = 1'b1; + storage_a_wdata[dst * 32 +: 32] = + tcu_lmem_if.rsp_data.data[src * 32 +: 32]; + end + end + end else begin + for (int b = 0; b < B_BUF_WORDS; ++b) begin + if (in_fetch_a) begin + storage_a_wren[b] = 1'b1; + storage_a_wdata[b * 32 +: 32] = + tcu_lmem_if.rsp_data.data[(int'(a_off_words) + b) * 32 +: 32]; + if (sparse_b_from_a) begin + storage_b_wren[b] = 1'b1; + storage_b_wdata[b * 32 +: 32] = + tcu_lmem_if.rsp_data.data[(int'(b_off_words) + b) * 32 +: 32]; + end + end + if (in_fetch_b) begin storage_b_wren[b] = 1'b1; storage_b_wdata[b * 32 +: 32] = tcu_lmem_if.rsp_data.data[(int'(b_off_words) + b) * 32 +: 32]; end end - if (in_fetch_b) begin - storage_b_wren[b] = 1'b1; - storage_b_wdata[b * 32 +: 32] = - tcu_lmem_if.rsp_data.data[(int'(b_off_words) + b) * 32 +: 32]; - end end end end @@ -408,7 +582,9 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( .clk (clk), .reset (reset), .read (1'b1), - .write ((in_fetch_b || (sparse_b_from_a && in_fetch_a)) && tcu_lmem_if.rsp_valid), + .write ((in_fetch_b && tcu_lmem_if.rsp_valid) + || (in_fetch_a && tcu_lmem_if.rsp_valid && slot_is_sparse_r + && (slot_row_major_r || sparse_b_from_a))), .wren (storage_b_wren), .waddr (1'b0), .wdata (storage_b_wdata), @@ -417,12 +593,11 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ); // ----------------------------------------------------------------------- - // Output mux (constant wiring). - // Dense: rs2[0..NUM_BANKS-1] = storage_A. tcu_core's b_off picks within. - // Sparse: flat read permutation — storage holds the block N-inner - // (word = kw_in*tcN + n_in); the FEDP wants - // rs2[k_idx*tcN*2 + n_in*2 + cand] with kw_in = k_idx*2 + cand. - // slot_a||slot_b split at B_BLOCK_WORDS. + // Output mux. + // Dense: rs2[0..NUM_BANKS-1] = storage_A (legacy). tcu_core's b_off + // picks within at execute time. + // Sparse K-major is written directly in FEDP order. Block-major keeps + // the flat producer layout and applies a fixed read permutation. // ----------------------------------------------------------------------- logic [TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0] rs2_mux; @@ -430,18 +605,26 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( rs2_mux = '0; for (int lane = 0; lane < TCU_WG_RS2_WIDTH; ++lane) begin if (slot_is_sparse_r) begin - automatic int unsigned k_idx_l = lane / (TCU_TC_N * 2); - automatic int unsigned rem_l = lane % (TCU_TC_N * 2); - automatic int unsigned n_in_l = rem_l / 2; - automatic int unsigned cand_l = rem_l % 2; - automatic int unsigned w_l = (k_idx_l * 2 + cand_l) * TCU_TC_N + n_in_l; - // Slot A holds the first B_BUF_WORDS words; the remainder (two-row - // only) lives in slot B. Index masked to stay in range when the - // block fits one row (slot B then unused and synth-pruned). - if (w_l < int'(B_BUF_WORDS)) begin - rs2_mux[lane] = `VX_CFG_XLEN'(storage_a_rdata[w_l]); - end else if (w_l < int'(2 * B_BLOCK_WORDS)) begin - rs2_mux[lane] = `VX_CFG_XLEN'(storage_b_rdata[(w_l - int'(B_BUF_WORDS)) & (B_BUF_WORDS-1)]); + if (slot_row_major_r) begin + if (lane < int'(B_BLOCK_WORDS_SP)) begin + rs2_mux[lane] = `VX_CFG_XLEN'(storage_a_rdata[lane]); + end else if (lane < int'(2 * B_BLOCK_WORDS_SP)) begin + rs2_mux[lane] = `VX_CFG_XLEN'( + storage_b_rdata[lane - int'(B_BLOCK_WORDS_SP)]); + end + end else begin + automatic int unsigned k_idx_l = lane / (TCU_TC_N * 2); + automatic int unsigned rem_l = lane % (TCU_TC_N * 2); + automatic int unsigned n_in_l = rem_l / 2; + automatic int unsigned cand_l = rem_l % 2; + automatic int unsigned w_l = + (k_idx_l * 2 + cand_l) * TCU_TC_N + n_in_l; + if (w_l < int'(B_BUF_WORDS)) begin + rs2_mux[lane] = `VX_CFG_XLEN'(storage_a_rdata[w_l]); + end else if (w_l < int'(2 * B_BLOCK_WORDS)) begin + rs2_mux[lane] = `VX_CFG_XLEN'( + storage_b_rdata[(w_l - int'(B_BUF_WORDS)) & (B_BUF_WORDS-1)]); + end end end else if (lane < int'(B_BUF_WORDS)) begin rs2_mux[lane] = `VX_CFG_XLEN'(storage_a_rdata[lane]); @@ -491,7 +674,7 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( if (alloc_en) `TRACE(3, ("%t: %s bbuf: alloc desc_b=0x%0h sparse=%0d step_k=%0d step_n=%0d addr_a=0x%0h addr_b=0x%0h sub_half=%0d\n", $time, INSTANCE_ID, req_desc_b, req_is_sparse, req_step_k, req_step_n, - fetch_addr_a, fetch_addr_b_flat, dense_sub_half)) + fetch_addr_a, fetch_addr_b_sparse, dense_sub_half)) if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) `TRACE(3, ("%t: %s bbuf: rd_req addr=0x%0h\n", $time, INSTANCE_ID, tcu_lmem_if.req_data.addr)) diff --git a/hw/rtl/tcu/VX_tcu_core.sv b/hw/rtl/tcu/VX_tcu_core.sv index fd2378ef48..5c1eaafafd 100644 --- a/hw/rtl/tcu/VX_tcu_core.sv +++ b/hw/rtl/tcu/VX_tcu_core.sv @@ -22,7 +22,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( input wire reset, `ifdef VX_CFG_TCU_WGMMA_ENABLE - input wire [TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, + input wire [TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, input wire [TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0] tbuf_rs2_data, input wire tbuf_ready, `endif @@ -43,30 +43,32 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ); `UNUSED_SPARAM (INSTANCE_ID); + localparam FEDP_K = TCU_WG_FEDP_K; + `ifdef VX_CFG_TCU_TYPE_DSP localparam FCVT_LATENCY = 1; localparam FMUL_LATENCY = 8; localparam FADD_LATENCY = 11; - localparam FACC_LATENCY = $clog2(2 * TCU_TC_K + 1) * FADD_LATENCY; - localparam FEDP_LATENCY = FCVT_LATENCY + FMUL_LATENCY + FACC_LATENCY; + localparam FACC_LATENCY = $clog2(2 * FEDP_K) * FADD_LATENCY; + localparam FEDP_LATENCY = FCVT_LATENCY + FMUL_LATENCY + FACC_LATENCY + FADD_LATENCY; `elsif VX_CFG_TCU_TYPE_BHF localparam FMUL_LATENCY = 2; localparam FADD_LATENCY = 2; localparam FRND_LATENCY = 1; - localparam FACC_LATENCY = $clog2(2 * TCU_TC_K + 1) * (FADD_LATENCY + FRND_LATENCY); - localparam FEDP_LATENCY = (FMUL_LATENCY + FRND_LATENCY) + 1 + FACC_LATENCY; + localparam FACC_LATENCY = $clog2(2 * FEDP_K) * (FADD_LATENCY + FRND_LATENCY); + localparam FEDP_LATENCY = (FMUL_LATENCY + FRND_LATENCY) + 1 + FACC_LATENCY + (FADD_LATENCY + FRND_LATENCY); `elsif VX_CFG_TCU_TYPE_FPNEW localparam FMUL_LATENCY = 6; localparam FMUX_LATENCY = 1; localparam FADD_LATENCY = 7; - localparam FACC_LATENCY = $clog2(2 * TCU_TC_K) * FADD_LATENCY; + localparam FACC_LATENCY = $clog2(2 * FEDP_K) * FADD_LATENCY; localparam FEDP_LATENCY = FMUL_LATENCY + FMUX_LATENCY + FACC_LATENCY + FADD_LATENCY; `elsif VX_CFG_TCU_TYPE_DPI localparam FMUL_LATENCY = 2; localparam FACC_LATENCY = 2; localparam FEDP_LATENCY = FMUL_LATENCY + FACC_LATENCY; `else // VX_CFG_TCU_TYPE_TFR - localparam FMUL_LATENCY = 1; // USE_DSP swaps Wallace->DSP at the same latency + localparam FMUL_LATENCY = 1; localparam FALN_LATENCY = 1; localparam FACC_LATENCY = 1; localparam FRND_LATENCY = 1; @@ -79,6 +81,10 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam LG_A_BS = $clog2(TCU_A_BLOCK_SIZE); localparam LG_B_BS = $clog2(TCU_B_BLOCK_SIZE); localparam OFF_W = $clog2(TCU_BLOCK_CAP); +`ifdef VX_CFG_TCU_WGMMA_ENABLE + localparam LG_WG_B_BS = $clog2(TCU_WG_B_BLOCK_SIZE); + localparam WG_B_OFF_W = $clog2(TCU_WG_RS2_WIDTH); +`endif `ifdef VX_CFG_TCU_SPARSE_ENABLE localparam LG_B_BS_SP = $clog2(TCU_B_BLOCK_SIZE_SP); @@ -90,27 +96,17 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `endif `ifdef VX_CFG_TCU_MX_ENABLE - wire is_wmma = (execute_if.data.op_type == INST_TCU_WMMA) - `ifdef VX_CFG_TCU_SPARSE_ENABLE - || (execute_if.data.op_type == INST_TCU_WMMA_SP) - `endif - ; `ifdef VX_CFG_TCU_SPARSE_ENABLE wire mx_is_sparse = is_sparse; `else wire mx_is_sparse = 1'b0; `endif localparam FEDP_SF = TCU_MX_MAX_SF; - localparam FEDP_SCALE_W = 2 * FEDP_SF * 8; // pipelined sf_a + sf_b scale bytes `else localparam FEDP_SF = 1; - localparam FEDP_SCALE_W = 0; + `UNUSED_PARAM (FEDP_SF) `endif - // FEDP pipeline-register payload: {c_val, fmt_s, fmt_d, b_col, a_row} plus - // the MX scale bytes when enabled. - localparam FEDP_PIPE_W = 32 + 5 + 5 + 2 * TCU_TC_K * 32 + FEDP_SCALE_W; - // ----------------------------------------------------------------------- // WGMMA / WMMA abstraction layer // ----------------------------------------------------------------------- @@ -118,7 +114,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // common interface. Downstream code uses only these wires and never // references tbuf_* or is_wgmma directly. - wire [TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] rs1_data; + wire [TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] rs1_data; `ifdef VX_CFG_TCU_WGMMA_ENABLE wire [TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0] rs2_data; `else @@ -138,15 +134,25 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // RF-side rs2_data is NUM_THREADS lanes wide; the WGMMA bbuf can be // wider (TCU_WG_RS2_WIDTH lanes). Pad/truncate to the wgmma width on // the false branch so both arms match TCU_WG_RS2_WIDTH * XLEN bits. + localparam WG_RS1_BITS = TCU_WG_A_DATA_SIZE * `VX_CFG_XLEN; localparam WG_RS2_BITS = TCU_WG_RS2_WIDTH * `VX_CFG_XLEN; + wire [WG_RS1_BITS-1:0] rs1_data_rf = WG_RS1_BITS'(execute_if.data.rs1_data); wire [WG_RS2_BITS-1:0] rs2_data_rf = WG_RS2_BITS'(execute_if.data.rs2_data); - assign rs1_data = (is_wgmma && wg_a_smem) ? tbuf_rs1_data : execute_if.data.rs1_data; + assign rs1_data = (is_wgmma && wg_a_smem) ? tbuf_rs1_data : rs1_data_rf; assign rs2_data = is_wgmma ? tbuf_rs2_data : rs2_data_rf; + `ifdef VX_CFG_TCU_SPARSE_ENABLE + // Sparse metadata lives in VX_tcu_sp_meta SRAM, preloaded via TCU_LD. + wire [TCU_MAX_META_BLOCK_WIDTH-1:0] vld_meta_block = wmma_sp_meta; + `endif + assign exe_ready_extra = ~is_wgmma || tbuf_ready; `else assign rs1_data = execute_if.data.rs1_data; assign rs2_data = execute_if.data.rs2_data; + `ifdef VX_CFG_TCU_SPARSE_ENABLE + wire [TCU_MAX_META_BLOCK_WIDTH-1:0] vld_meta_block = wmma_sp_meta; + `endif assign exe_ready_extra = 1'b1; `endif @@ -158,18 +164,74 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( wire [4:0] fmt_d = execute_if.data.op_args.tcu.fmt_d; wire execute_fire = execute_if.valid && execute_if.ready; +`ifdef VX_CFG_TCU_WGMMA_ENABLE + wire is_wgmma_setup = is_wgmma && !execute_if.data.header.wb; +`else + wire is_wgmma_setup = 1'b0; +`endif + wire setup_enqueue = execute_fire && is_wgmma_setup; + wire fedp_enqueue = execute_fire && !is_wgmma_setup; + + // ----------------------------------------------------------------------- + // Sparse metadata: VX_tcu_sp_meta (for WMMA_SP) + optional tile-buffer mux + // ----------------------------------------------------------------------- + + tcu_header_t mdata_queue_in; + always_comb begin + mdata_queue_in = execute_if.data.header; + end `UNUSED_VAR ({step_m, step_n, step_k, fmt_s, fmt_d, execute_if.data}); +`ifdef VX_TCU_LD_TRACE +`ifdef VX_CFG_TCU_SPARSE_ENABLE + // META_RD trace: logs vld_meta_block at FEDP consume time. + // Format: META_RD,wid,step_m,step_k,wg_bank,word_lo32 + wire trc_is_sp = (execute_if.data.op_type == INST_OP_BITS'(INST_TCU_WMMA_SP)) + `ifdef VX_CFG_TCU_WGMMA_ENABLE + || (execute_if.data.op_type == INST_OP_BITS'(INST_TCU_WGMMA_SP)) + `endif + ; + wire [3:0] trc_wg_bank = ((TCU_K_STEPS > 2) ? (step_m << 1) : step_m) | step_k; + always @(posedge clk) begin + if (execute_fire && trc_is_sp) begin + $write("META_RD,%0d,%0d,%0d,%0d,0x%08h\n", + execute_if.data.header.wid, step_m, step_k, trc_wg_bank, + vld_meta_block[31:0]); + end + end +`endif +`endif + // ----------------------------------------------------------------------- // Pipeline control // ----------------------------------------------------------------------- wire mdata_queue_full; - wire result_fire = result_if.valid && result_if.ready; wire fedp_enable, fedp_done; + reg setup_valid_r; + tcu_header_t setup_header_r; + tcu_header_t mdata_queue_out; + + wire setup_result_fire = setup_valid_r && result_if.ready; + wire fedp_result_fire = fedp_done && result_if.ready && !setup_valid_r; + + always @(posedge clk) begin + if (reset) begin + setup_valid_r <= 1'b0; + end else begin + if (setup_result_fire) begin + setup_valid_r <= 1'b0; + end + if (setup_enqueue) begin + setup_valid_r <= 1'b1; + setup_header_r <= execute_if.data.header; + end + end + end + reg [PIPE_LATENCY-1:0] fedp_delay_pipe; always @(posedge clk) begin if (reset) begin @@ -178,18 +240,21 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( if (fedp_enable) begin fedp_delay_pipe <= fedp_delay_pipe >> 1; end - if (execute_fire) begin + if (fedp_enqueue) begin fedp_delay_pipe[PIPE_LATENCY-1] <= 1; end end end assign fedp_done = fedp_delay_pipe[0]; - assign result_if.valid = fedp_done; - assign fedp_enable = ~fedp_done || result_if.ready; - assign execute_if.ready = ~mdata_queue_full && fedp_enable && exe_ready_extra; + assign result_if.valid = setup_valid_r || fedp_done; + assign fedp_enable = ~fedp_done || fedp_result_fire; + assign execute_if.ready = is_wgmma_setup + ? ((~setup_valid_r || result_if.ready) && exe_ready_extra) + : (~mdata_queue_full && fedp_enable && exe_ready_extra); + + wire mdata_push = fedp_enqueue; - // Header FIFO: carries the result header alongside the FEDP pipeline. VX_fifo_queue #( .DATAW ($bits(tcu_header_t)), .DEPTH (MDATA_QUEUE_DEPTH), @@ -197,10 +262,10 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ) mdata_queue ( .clk (clk), .reset (reset), - .push (execute_fire), - .pop (result_fire), - .data_in(execute_if.data.header), - .data_out(result_if.data.header), + .push (mdata_push), + .pop (fedp_result_fire), + .data_in(mdata_queue_in), + .data_out(mdata_queue_out), `UNUSED_PIN(empty), `UNUSED_PIN(alm_empty), .full (mdata_queue_full), @@ -208,22 +273,26 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `UNUSED_PIN(size) ); + assign result_if.data.header = setup_valid_r ? setup_header_r : mdata_queue_out; + // ----------------------------------------------------------------------- // Operand offset computation // ----------------------------------------------------------------------- wire [OFF_W-1:0] a_off = (OFF_W'(step_m) & OFF_W'(TCU_A_SUB_BLOCKS-1)) << LG_A_BS; + wire [OFF_W-1:0] b_off_wm; `ifdef VX_CFG_TCU_SPARSE_ENABLE - wire [OFF_W-1:0] b_off = is_sparse + assign b_off_wm = is_sparse ? (OFF_W'(step_n) & OFF_W'(TCU_B_SUB_BLOCKS_SP-1)) << LG_B_BS_SP : (OFF_W'(step_n) & OFF_W'(TCU_B_SUB_BLOCKS-1)) << LG_B_BS; `else - wire [OFF_W-1:0] b_off = (OFF_W'(step_n) & OFF_W'(TCU_B_SUB_BLOCKS-1)) << LG_B_BS; + assign b_off_wm = (OFF_W'(step_n) & OFF_W'(TCU_B_SUB_BLOCKS-1)) << LG_B_BS; +`endif +`ifdef VX_CFG_TCU_WGMMA_ENABLE + wire [WG_B_OFF_W-1:0] b_off_wg = + (WG_B_OFF_W'(step_n) & WG_B_OFF_W'(TCU_WG_B_SUB_BLOCKS-1)) << LG_WG_B_BS; `endif - // ----------------------------------------------------------------------- - // Unified metadata SRAM: sparse lane masks (WMMA_SP) and MX scale factors - // ----------------------------------------------------------------------- `ifdef VX_CFG_TCU_SPARSE_ENABLE wire [TCU_MAX_META_BLOCK_WIDTH-1:0] wmma_sp_meta; `endif @@ -232,10 +301,6 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( wire [TCU_BLOCK_CAP-1:0][31:0] mx_meta_b; `endif - // The metadata SRAM exists only when a metadata-consuming mode (sparse or - // MX) is enabled; TCU_META_ENABLE = (MX or SPARSE). Its module body - // is `ifdef`-guarded on the same symbol, so the instantiation must match — - // otherwise plain TCU (no MX/SPARSE) references an empty module. `ifdef TCU_META_ENABLE VX_tcu_meta #( .INSTANCE_ID (INSTANCE_ID) @@ -260,21 +325,73 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `endif `ifdef VX_CFG_TCU_MX_ENABLE - // MX scale-factor decode (read-side complement to VX_tcu_meta's MX region). + + localparam MX_MAX_MN = TCU_TILE_M > TCU_TILE_N ? TCU_TILE_M : TCU_TILE_N; + localparam MX_IDX_W = $clog2(MX_MAX_MN); + localparam MX_TILE_K_MAX = `MAX(TCU_TILE_K, TCU_WG_K_STEPS * TCU_WG_FEDP_K); + localparam MX_K_IDX_W = `LOG2UP(MX_TILE_K_MAX * TCU_MAX_ELT_RATIO); + localparam MX_SCALE_BLOCKS_MAX = mx_scale_blocks_k_words(TCU_NVFP4_ID, MX_TILE_K_MAX); + localparam MX_SCALE_IDX_W = $clog2(MX_MAX_MN * MX_SCALE_BLOCKS_MAX); + + function automatic [7:0] mx_scale_at( + input logic [TCU_BLOCK_CAP-1:0][31:0] meta, + input logic [4:0] fmt, + input logic [MX_SCALE_IDX_W-1:0] scale_blocks_k, + input logic [MX_IDX_W-1:0] mn_idx, + input logic [MX_K_IDX_W-1:0] k_base_idx + ); + logic [MX_SCALE_IDX_W-1:0] scale_k; + logic [MX_SCALE_IDX_W-1:0] scale_idx; + logic [`LOG2UP(TCU_BLOCK_CAP)-1:0] word_idx; + logic [1:0] byte_idx; + begin + scale_k = MX_SCALE_IDX_W'(k_base_idx / mx_scale_block_size(fmt)); + scale_idx = MX_SCALE_IDX_W'(mn_idx) * scale_blocks_k + scale_k; + word_idx = `LOG2UP(TCU_BLOCK_CAP)'(scale_idx >> 2); + byte_idx = scale_idx[1:0]; + mx_scale_at = meta[word_idx][byte_idx * 8 +: 8]; + end + endfunction + wire [TCU_TC_M-1:0][FEDP_SF-1:0][7:0] mx_sf_a; wire [TCU_TC_N-1:0][FEDP_SF-1:0][7:0] mx_sf_b; - VX_tcu_mx_scale mx_scale ( - .meta_a (mx_meta_a), - .meta_b (mx_meta_b), - .step_m (step_m), - .step_n (step_n), - .step_k (step_k), - .fmt_s (fmt_s), - .is_wmma (is_wmma), - .is_sparse(mx_is_sparse), - .sf_a (mx_sf_a), - .sf_b (mx_sf_b) - ); + wire [3:0] mx_elems_per_word = 4'(32 / tcu_fmt_width(fmt_s)); + wire [MX_SCALE_IDX_W-1:0] mx_scale_blocks_k_eff = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? MX_SCALE_IDX_W'(mx_scale_blocks_k_words(fmt_s, TCU_WG_K_STEPS * TCU_WG_FEDP_K)) + : MX_SCALE_IDX_W'(mx_scale_blocks_k_words(fmt_s, TCU_TILE_K)); + `else + MX_SCALE_IDX_W'(mx_scale_blocks_k_words(fmt_s, TCU_TILE_K)); + `endif + wire [MX_K_IDX_W:0] mx_uop_k_words = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? (MX_K_IDX_W+1)'(TCU_WG_FEDP_K) : (MX_K_IDX_W+1)'(TCU_TC_K); + `else + (MX_K_IDX_W+1)'(TCU_TC_K); + `endif + wire [MX_K_IDX_W:0] mx_uop_k_elems = (MX_K_IDX_W+1)'( + mx_uop_k_words * (MX_K_IDX_W+1)'(mx_elems_per_word) + * (MX_K_IDX_W+1)'(mx_is_sparse ? 2 : 1)); + wire [MX_K_IDX_W:0] mx_fedp_k_elems = (MX_K_IDX_W+1)'( + (MX_K_IDX_W+1)'(FEDP_K) * (MX_K_IDX_W+1)'(mx_elems_per_word) + * (MX_K_IDX_W+1)'(mx_is_sparse ? 2 : 1)); + wire [MX_K_IDX_W-1:0] mx_k_base_idx = MX_K_IDX_W'(step_k * mx_uop_k_elems); + + for (genvar i = 0; i < TCU_TC_M; ++i) begin : g_mx_sf_a_i + wire [MX_IDX_W-1:0] mx_a_idx = MX_IDX_W'(step_m) * MX_IDX_W'(TCU_TC_M) + MX_IDX_W'(i); + for (genvar s = 0; s < FEDP_SF; ++s) begin : g_s + wire [MX_K_IDX_W-1:0] mx_k_idx = mx_k_base_idx + MX_K_IDX_W'((s * mx_fedp_k_elems) / FEDP_SF); + assign mx_sf_a[i][s] = mx_scale_at(mx_meta_a, fmt_s, mx_scale_blocks_k_eff, mx_a_idx, mx_k_idx); + end + end + + for (genvar j = 0; j < TCU_TC_N; ++j) begin : g_mx_sf_b_j + wire [MX_IDX_W-1:0] mx_b_idx = MX_IDX_W'(step_n) * MX_IDX_W'(TCU_TC_N) + MX_IDX_W'(j); + for (genvar s = 0; s < FEDP_SF; ++s) begin : g_s + wire [MX_K_IDX_W-1:0] mx_k_idx = mx_k_base_idx + MX_K_IDX_W'((s * mx_fedp_k_elems) / FEDP_SF); + assign mx_sf_b[j][s] = mx_scale_at(mx_meta_b, fmt_s, mx_scale_blocks_k_eff, mx_b_idx, mx_k_idx); + end + end `endif // ----------------------------------------------------------------------- @@ -286,18 +403,80 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( for (genvar i = 0; i < TCU_TC_M; ++i) begin : g_i for (genvar j = 0; j < TCU_TC_N; ++j) begin : g_j `ifdef VX_CFG_TCU_SPARSE_ENABLE - wire [TCU_TC_K-1:0][31:0] a_row, b_col, b_col_dense, b_col_sparse, b_col_1, b_col_2; + wire [FEDP_K-1:0][31:0] a_row, b_col, b_col_dense; + wire [TCU_TC_K-1:0][31:0] b_col_sparse, b_col_1, b_col_2; `else - wire [TCU_TC_K-1:0][31:0] a_row, b_col; + wire [FEDP_K-1:0][31:0] a_row, b_col; `endif `ifdef VX_CFG_TCU_MX_ENABLE wire [FEDP_SF-1:0][7:0] sf_a = mx_sf_a[i]; wire [FEDP_SF-1:0][7:0] sf_b = mx_sf_b[j]; `endif - for (genvar k_idx = 0; k_idx < TCU_TC_K; ++k_idx) begin : g_slice_assign - assign a_row[k_idx] = 32'(rs1_data[a_off + i * TCU_TC_K + k_idx]); - `ifdef VX_CFG_TCU_SPARSE_ENABLE - assign b_col_dense[k_idx] = 32'(rs2_data[b_off + j * TCU_TC_K + k_idx]); + for (genvar k_idx = 0; k_idx < FEDP_K; ++k_idx) begin : g_slice_assign + `ifdef VX_CFG_TCU_WGMMA_ENABLE + localparam int WG_B_IDX = j * TCU_WG_FEDP_K + k_idx; + `endif + if (k_idx < TCU_TC_K) begin : g_lo + `ifdef VX_CFG_TCU_WGMMA_ENABLE + wire [31:0] a_wgmma_smem = 32'(rs1_data[i * TCU_WG_FEDP_K + k_idx]); + wire [31:0] a_wgmma_reg = 32'(execute_if.data.rs1_data[i * TCU_TC_K + k_idx]); + assign a_row[k_idx] = is_wgmma + ? (wg_a_smem ? a_wgmma_smem : a_wgmma_reg) + : 32'(execute_if.data.rs1_data[a_off + i * TCU_TC_K + k_idx]); + `else + assign a_row[k_idx] = 32'(rs1_data[a_off + i * TCU_TC_K + k_idx]); + `endif + `ifdef VX_CFG_TCU_SPARSE_ENABLE + assign b_col_dense[k_idx] = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? 32'(rs2_data[int'(b_off_wg) + WG_B_IDX]) : + `endif + 32'(rs2_data[b_off_wm + j * TCU_TC_K + k_idx]); + `else + assign b_col[k_idx] = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? 32'(rs2_data[int'(b_off_wg) + WG_B_IDX]) : + `endif + 32'(rs2_data[b_off_wm + j * TCU_TC_K + k_idx]); + `endif + end else begin : g_hi + `ifdef VX_CFG_TCU_WGMMA_ENABLE + wire [31:0] a_wgmma_smem = 32'(rs1_data[i * TCU_WG_FEDP_K + k_idx]); + wire [31:0] a_wgmma_reg = + `ifdef VX_CFG_TCU_FEDP2K + `ifdef VX_CFG_TCU_SPARSE_ENABLE + is_sparse ? 32'b0 : + `endif + 32'(execute_if.data.rs2_data[i * TCU_TC_K + (k_idx - TCU_TC_K)]); + `else + 32'b0; + `endif + assign a_row[k_idx] = (is_wgmma + `ifdef VX_CFG_TCU_SPARSE_ENABLE + && !is_sparse + `endif + ) ? (wg_a_smem ? a_wgmma_smem : a_wgmma_reg) : 32'b0; + `else + assign a_row[k_idx] = 32'b0; + `endif + `ifdef VX_CFG_TCU_SPARSE_ENABLE + assign b_col_dense[k_idx] = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + (is_wgmma && !is_sparse) ? 32'(rs2_data[int'(b_off_wg) + WG_B_IDX]) : + `endif + 32'b0; + `else + assign b_col[k_idx] = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? 32'(rs2_data[int'(b_off_wg) + WG_B_IDX]) : + `endif + 32'b0; + `endif + end + end + + `ifdef VX_CFG_TCU_SPARSE_ENABLE + for (genvar k_idx = 0; k_idx < TCU_TC_K; ++k_idx) begin : g_sparse_slice_assign localparam J_SP = SYM_SPARSE ? (j % (TCU_TC_N / 2)) : j; // rs2_data sparse-pair layout differs by op: // WGMMA_SP: source is tbuf (shared mem), K-major → @@ -306,20 +485,18 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // idx = J_SP*(TC_K*2) + k_idx*2 + lane // The two layouts are incompatible; separate formulas are required. `ifdef VX_CFG_TCU_WGMMA_ENABLE - wire [31:0] b_col_1_wg = 32'(rs2_data[b_off + k_idx * TCU_TC_N * 2 + J_SP * 2]); - wire [31:0] b_col_2_wg = 32'(rs2_data[b_off + k_idx * TCU_TC_N * 2 + J_SP * 2 + 1]); - wire [31:0] b_col_1_wm = 32'(rs2_data[b_off + J_SP * TCU_TC_K * 2 + k_idx * 2]); - wire [31:0] b_col_2_wm = 32'(rs2_data[b_off + J_SP * TCU_TC_K * 2 + k_idx * 2 + 1]); + wire [31:0] b_col_1_wg = 32'(rs2_data[k_idx * TCU_TC_N * 2 + J_SP * 2]); + wire [31:0] b_col_2_wg = 32'(rs2_data[k_idx * TCU_TC_N * 2 + J_SP * 2 + 1]); + wire [31:0] b_col_1_wm = 32'(rs2_data[b_off_wm + J_SP * TCU_TC_K * 2 + k_idx * 2]); + wire [31:0] b_col_2_wm = 32'(rs2_data[b_off_wm + J_SP * TCU_TC_K * 2 + k_idx * 2 + 1]); assign b_col_1[k_idx] = is_wgmma ? b_col_1_wg : b_col_1_wm; assign b_col_2[k_idx] = is_wgmma ? b_col_2_wg : b_col_2_wm; `else - assign b_col_1[k_idx] = 32'(rs2_data[b_off + J_SP * TCU_TC_K * 2 + k_idx * 2]); - assign b_col_2[k_idx] = 32'(rs2_data[b_off + J_SP * TCU_TC_K * 2 + k_idx * 2 + 1]); - `endif - `else - assign b_col[k_idx] = 32'(rs2_data[b_off + j * TCU_TC_K + k_idx]); + assign b_col_1[k_idx] = 32'(rs2_data[b_off_wm + J_SP * TCU_TC_K * 2 + k_idx * 2]); + assign b_col_2[k_idx] = 32'(rs2_data[b_off_wm + J_SP * TCU_TC_K * 2 + k_idx * 2 + 1]); `endif end + `endif wire [31:0] c_val = 32'(execute_if.data.rs3_data[i * TCU_TC_N + j]); @@ -331,27 +508,54 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( .fmt_s (fmt_s), .b_col_in1 (b_col_1), .b_col_in2 (b_col_2), - .vld_mask (wmma_sp_meta), + .vld_mask (vld_meta_block), .b_col_out (b_col_sparse) ); - assign b_col = is_sparse ? b_col_sparse : b_col_dense; + for (genvar k_idx = 0; k_idx < FEDP_K; ++k_idx) begin : g_sparse_b_select + if (k_idx < TCU_TC_K) begin : g_lo + assign b_col[k_idx] = is_sparse ? b_col_sparse[k_idx] : b_col_dense[k_idx]; + end else begin : g_hi + assign b_col[k_idx] = is_sparse ? 32'b0 : b_col_dense[k_idx]; + end + end + + `ifdef VX_TCU_LD_TRACE + // GATHER trace: GATHER,wid,step_m,step_n,i,k,bword0,bword1,lo,hi,gathered + // One line per (i, j, k_idx); emitted only for sparse ops. + always @(posedge clk) begin + if (execute_fire && is_sparse) begin + for (int kk = 0; kk < TCU_TC_K; ++kk) begin + $write("GATHER,%0d,%0d,%0d,%0d,%0d,0x%08h,0x%08h,?,?,0x%08h\n", + execute_if.data.header.wid, step_m, step_n, + i, j*TCU_TC_K + kk, + b_col_1[kk], b_col_2[kk], b_col_sparse[kk]); + end + end + end + `endif `endif - // Dual-side sparse lane mask. + // Dual-side sparse lane mask `ifdef VX_CFG_TCU_TYPE_TFR wire [TCU_MAX_INPUTS-1:0] vld_mask_r; `ifdef VX_CFG_TCU_DSM_ENABLE + wire [TCU_MAX_INPUTS-1:0] vld_mask; VX_tcu_dsm #( - .N (TCU_TC_K), - .OUT_REG (1) + .N (TCU_TC_K) ) dual_sparse_mask ( - .clk (clk), - .reset (reset), - .enable (fedp_enable), .fmt_s (fmt_s), .a_row (a_row), .b_col (b_col), - .vld_mask (vld_mask_r) + .vld_mask (vld_mask) + ); + VX_pipe_register #( + .DATAW (TCU_MAX_INPUTS) + ) pipe_vld_mask ( + .clk (clk), + .reset (reset), + .enable (fedp_enable), + .data_in (vld_mask), + .data_out (vld_mask_r) ); `else assign vld_mask_r = '1; @@ -359,35 +563,39 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `endif wire [4:0] fmt_s_r, fmt_d_r; - wire [TCU_TC_K-1:0][31:0] a_row_r, b_col_r; + wire [FEDP_K-1:0][31:0] a_row_r, b_col_r; `ifdef VX_CFG_TCU_MX_ENABLE wire [FEDP_SF-1:0][7:0] sf_a_r, sf_b_r; `endif wire [31:0] c_val_r; - wire [FEDP_PIPE_W-1:0] fedp_pipe_in, fedp_pipe_out; `ifdef VX_CFG_TCU_MX_ENABLE - assign fedp_pipe_in = {c_val, sf_b, sf_a, fmt_s, fmt_d, b_col, a_row}; - assign {c_val_r, sf_b_r, sf_a_r, fmt_s_r, fmt_d_r, b_col_r, a_row_r} = fedp_pipe_out; + VX_pipe_register #( + .DATAW (32 + 5 + 5 + FEDP_K * 32 + FEDP_K * 32 + 2 * FEDP_SF * 8) + ) pipe_fedp ( + .clk (clk), + .reset (reset), + .enable (fedp_enable), + .data_in ({c_val, sf_b, sf_a, fmt_s, fmt_d, b_col, a_row}), + .data_out ({c_val_r, sf_b_r, sf_a_r, fmt_s_r, fmt_d_r, b_col_r, a_row_r}) + ); `else - assign fedp_pipe_in = {c_val, fmt_s, fmt_d, b_col, a_row}; - assign {c_val_r, fmt_s_r, fmt_d_r, b_col_r, a_row_r} = fedp_pipe_out; - `endif VX_pipe_register #( - .DATAW (FEDP_PIPE_W) + .DATAW (32 + 5 + 5 + FEDP_K * 32 + FEDP_K * 32) ) pipe_fedp ( .clk (clk), .reset (reset), .enable (fedp_enable), - .data_in (fedp_pipe_in), - .data_out (fedp_pipe_out) + .data_in ({c_val, fmt_s, fmt_d, b_col, a_row}), + .data_out ({c_val_r, fmt_s_r, fmt_d_r, b_col_r, a_row_r}) ); + `endif `ifdef VX_CFG_TCU_TYPE_DPI VX_tcu_fedp_dpi #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K), + .N (FEDP_K), .SF (FEDP_SF) ) fedp ( .clk (clk), @@ -395,8 +603,8 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( .enable(fedp_enable), .fmt_s (fmt_s_r), .fmt_d (fmt_d_r), - .a_row (a_row_r), - .b_col (b_col_r), + .a_row(a_row_r), + .b_col(b_col_r), `ifdef VX_CFG_TCU_MX_ENABLE .sf_a (sf_a_r), .sf_b (sf_b_r), @@ -408,15 +616,15 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_tcu_fedp_bhf #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K) + .N (FEDP_K) ) fedp ( .clk (clk), .reset (reset), .enable(fedp_enable), .fmt_s (fmt_s_r), .fmt_d (fmt_d_r), - .a_row (a_row_r), - .b_col (b_col_r), + .a_row(a_row_r), + .b_col(b_col_r), .c_val (c_val_r), .d_val (d_val[i][j]) ); @@ -424,7 +632,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_tcu_fedp_fpnew #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K) + .N (FEDP_K) ) fedp ( .clk (clk), .reset (reset), @@ -440,14 +648,13 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_tcu_fedp_tfr #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K), - .SF (FEDP_SF), - .USE_DSP (`VX_CFG_TCU_USE_DSP) + .N (FEDP_K), + .SF (FEDP_SF) ) fedp ( .clk (clk), .reset (reset), - .enable(fedp_enable), .vld_mask(vld_mask_r), + .enable(fedp_enable), .fmt_s (fmt_s_r), .fmt_d (fmt_d_r), .a_row (a_row_r), @@ -463,15 +670,15 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_tcu_fedp_dsp #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K) + .N (FEDP_K) ) fedp ( .clk (clk), .reset (reset), .enable(fedp_enable), .fmt_s (fmt_s_r), .fmt_d (fmt_d_r), - .a_row (a_row_r), - .b_col (b_col_r), + .a_row(a_row_r), + .b_col(b_col_r), .c_val (c_val_r), .d_val (d_val[i][j]) ); @@ -488,9 +695,9 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( always @(posedge clk) begin if (execute_if.valid && execute_if.ready) begin `TRACE(3, ("%t: %s FEDP-enq: wid=%0d, cta_id=%0d, i=%0d, j=%0d, m=%0d, n=%0d, a_row=", $time, INSTANCE_ID, execute_if.data.header.wid, execute_if.data.header.cta_id, i, j, step_m, step_n)) - `TRACE_ARRAY1D(2, "0x%0h", a_row, TCU_TC_K) + `TRACE_ARRAY1D(2, "0x%0h", a_row, FEDP_K) `TRACE(3, (", b_col=")); - `TRACE_ARRAY1D(2, "0x%0h", b_col, TCU_TC_K) + `TRACE_ARRAY1D(2, "0x%0h", b_col, FEDP_K) `TRACE(3, (", c_val=0x%0h (#%0d)\n", c_val, execute_if.data.header.uuid)); end if (result_if.valid && result_if.ready) begin diff --git a/hw/rtl/tcu/VX_tcu_dsm.sv b/hw/rtl/tcu/VX_tcu_dsm.sv index 5ce27e3aac..a8f88a88b7 100644 --- a/hw/rtl/tcu/VX_tcu_dsm.sv +++ b/hw/rtl/tcu/VX_tcu_dsm.sv @@ -110,13 +110,6 @@ module VX_tcu_dsm import VX_tcu_pkg::*; #( vld_mask_per_k[k][e * 2 +: 2] = {2{a_i8_nz[e] && b_i8_nz[e]}}; end end - `ifdef VX_CFG_TCU_MX_ENABLE - TCU_MXI8_ID: begin - for (int e = 0; e < 4; ++e) begin - vld_mask_per_k[k][e * 2 +: 2] = {2{a_i8_nz[e] && b_i8_nz[e]}}; - end - end - `endif `endif `ifdef VX_CFG_TCU_INT4_ENABLE TCU_I4_ID, diff --git a/hw/rtl/tcu/VX_tcu_lockstep.sv b/hw/rtl/tcu/VX_tcu_lockstep.sv index 42cf9d9ed8..ff3d438762 100644 --- a/hw/rtl/tcu/VX_tcu_lockstep.sv +++ b/hw/rtl/tcu/VX_tcu_lockstep.sv @@ -23,12 +23,10 @@ // drain before a new CTA enters. Same-CTA across blocks (the production // case for a warpgroup at the same uop) remains free. // -// Single-owner design: by construction only one CTA can occupy the bbuf -// at a time, so a single tcu_owner_r / tcu_owned_r pair is sufficient. -// Multiple blocks firing leader uops on the same cycle for the same CTA -// (the production warpgroup case) agree on the cta, so the priority- -// encoded pick agrees with all of them. No same-cycle multi-CTA-firing -// race exists in current configs. +// Per-block ownership matches the SimX model: a block records its CTA when +// its first compute uop enters WGMMA, and releases on its last compute uop. +// A block is deferred only while another block is mid-WGMMA for a different +// CTA. // // Contract: the consumer must AND `cta_conflict` into the request // validity presented to bbuf and into any downstream ready that gates @@ -56,29 +54,24 @@ module VX_tcu_lockstep import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // ----------------------------------------------------------------------- // State // ----------------------------------------------------------------------- - reg tcu_owned_r; - reg [NCTA_WIDTH-1:0] tcu_owner_r; - // Per-block "WGMMA expansion in progress" — set on first sub-uop fire, - // cleared on last sub-uop fire. Persists across LMEM-stall gaps so the - // gate keeps a different CTA's WGMMA from sneaking in mid-expansion - // and corrupting the shared bbuf descriptor latch. + // Per-block "WGMMA expansion in progress" — first compute uop joins the + // CTA-owned region; the last compute uop releases it. This persists + // across LMEM-stall gaps so a different CTA cannot enter mid-expansion + // and corrupt the shared bbuf state. reg [BLOCK_SIZE-1:0] in_expansion_r; + reg [BLOCK_SIZE-1:0][NCTA_WIDTH-1:0] cta_owner_r; // ----------------------------------------------------------------------- // Owner update. // - // Leader-fire detection: any block firing its first sub-uop this cycle. + // Leader-fire detection: any block joining a WGMMA expansion this cycle. // The owner cta is picked from the lowest-indexed leader-firing block. // Under same-CTA warpgroup fire (production path), every leader-firing // block carries the same cta_id, so the priority pick agrees with all. // ----------------------------------------------------------------------- wire [BLOCK_SIZE-1:0] leader_fire_b = exec_fire_b & is_first_uop_b & is_wgmma_b; - wire any_leader_fire = |leader_fire_b; - wire any_in_expansion_next; - // in_expansion_next reflects this-cycle updates (set on leader fire, - // clear on last-uop fire). Released the cycle after the last block - // exits expansion. + // clear on last-uop fire). wire [BLOCK_SIZE-1:0] in_expansion_next; for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_inexp_next wire is_wgmma_fire = exec_fire_b[bi] && is_wgmma_b[bi]; @@ -86,33 +79,16 @@ module VX_tcu_lockstep import VX_gpu_pkg::*, VX_tcu_pkg::*; #( (in_expansion_r[bi] || (is_wgmma_fire && is_first_uop_b[bi])) && !(is_wgmma_fire && is_last_uop_b[bi]); end - assign any_in_expansion_next = |in_expansion_next; - - // first_fire_cta: priority-encoded pick of the lowest leader-firing block. - logic [NCTA_WIDTH-1:0] first_fire_cta; - always_comb begin - first_fire_cta = '0; - for (int b = BLOCK_SIZE-1; b >= 0; b--) begin - if (leader_fire_b[b]) first_fire_cta = new_cta_b[b]; - end - end always @(posedge clk) begin if (reset) begin - tcu_owned_r <= 1'b0; - tcu_owner_r <= '0; in_expansion_r <= '0; + cta_owner_r <= '0; end else begin in_expansion_r <= in_expansion_next; - if (!tcu_owned_r) begin - if (any_leader_fire) begin - tcu_owned_r <= 1'b1; - tcu_owner_r <= first_fire_cta; - end - end else begin - // Release once no block is still mid-expansion. - if (!any_in_expansion_next) begin - tcu_owned_r <= 1'b0; + for (int bi = 0; bi < BLOCK_SIZE; ++bi) begin + if (leader_fire_b[bi]) begin + cta_owner_r[bi] <= new_cta_b[bi]; end end end @@ -122,9 +98,16 @@ module VX_tcu_lockstep import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Per-block conflict — flat combinational, no propagation chain. // ----------------------------------------------------------------------- for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_conflict - assign cta_conflict[bi] = is_wgmma_b[bi] - && tcu_owned_r - && (tcu_owner_r != new_cta_b[bi]); + logic cta_conflict_r; + always_comb begin + cta_conflict_r = 1'b0; + for (int k = 0; k < BLOCK_SIZE; ++k) begin + if (k != bi && in_expansion_r[k] && (cta_owner_r[k] != new_cta_b[bi])) begin + cta_conflict_r = 1'b1; + end + end + end + assign cta_conflict[bi] = is_wgmma_b[bi] && cta_conflict_r; end // ----------------------------------------------------------------------- diff --git a/hw/rtl/tcu/VX_tcu_pkg.sv b/hw/rtl/tcu/VX_tcu_pkg.sv index 541c88f036..cfe5d9c6fc 100644 --- a/hw/rtl/tcu/VX_tcu_pkg.sv +++ b/hw/rtl/tcu/VX_tcu_pkg.sv @@ -41,8 +41,6 @@ package VX_tcu_pkg; localparam TCU_U8_ID = 18; localparam TCU_I4_ID = 19; localparam TCU_U4_ID = 20; - // MX formats have 2nd-MSB set - localparam TCU_MXI8_ID = 24; localparam TCU_FMT_WIDTH= 5; // Set configuration parameters @@ -90,23 +88,29 @@ package VX_tcu_pkg; // WGMMA per-warp tile dimensions (NRA=4 fixed, NRC=NR variable). // Derived from block geometry: xtileM = 2*tcM, xtileK = 2*tcK. - // m_steps = k_steps = 2 always. + // Dense FEDP width doubles under VX_CFG_TCU_FEDP2K localparam TCU_WG_TILE_M = 2 * TCU_TC_M; localparam TCU_WG_TILE_K = 2 * TCU_TC_K; +`ifdef VX_CFG_TCU_FEDP2K + localparam TCU_WG_FEDP_K = 2 * TCU_TC_K; +`else + localparam TCU_WG_FEDP_K = TCU_TC_K; +`endif localparam TCU_WG_TILE_N = (TCU_WG_NR * TCU_NT) / TCU_WG_TILE_M; // WG step counts: block geometry (TC_M/TC_N/TC_K) unchanged, tile is larger localparam TCU_WG_M_STEPS = TCU_WG_TILE_M / TCU_TC_M; localparam TCU_WG_N_STEPS = TCU_WG_TILE_N / TCU_TC_N; - localparam TCU_WG_K_STEPS = TCU_WG_TILE_K / TCU_TC_K; + localparam TCU_WG_K_STEPS = TCU_WG_TILE_K / TCU_WG_FEDP_K; localparam TCU_WG_UOPS = TCU_WG_M_STEPS * TCU_WG_N_STEPS * TCU_WG_K_STEPS; // WG A/B micro-tiling (block geometry is shared with non-WG) localparam TCU_WG_A_BLOCK_SIZE = TCU_TC_M * TCU_TC_K; + localparam TCU_WG_A_DATA_SIZE = TCU_TC_M * TCU_WG_FEDP_K; localparam TCU_WG_A_SUB_BLOCKS = TCU_BLOCK_CAP / TCU_WG_A_BLOCK_SIZE; - localparam TCU_WG_B_BLOCK_SIZE = TCU_TC_K * TCU_TC_N; + localparam TCU_WG_B_BLOCK_SIZE = TCU_WG_FEDP_K * TCU_TC_N; localparam TCU_WG_B_SUB_BLOCKS = TCU_BLOCK_CAP / TCU_WG_B_BLOCK_SIZE; // Symmetric sparse flag (NT=4, NT=16: block_em == block_en) @@ -121,9 +125,11 @@ package VX_tcu_pkg; localparam TCU_B_SUB_BLOCKS_SP = TCU_BLOCK_CAP / TCU_B_BLOCK_SIZE_SP; // WGMMA_SP always needs the full candidate lane set, regardless of SYM_SPARSE. localparam TCU_WG_B_BLOCK_SIZE_SP = TCU_TC_K * TCU_TC_N * 2; - // Width of the tbuf_rs2_data port: wider only when SPARSE is enabled (WGMMA_SP path). - // Without SPARSE, only TCU_BLOCK_CAP lanes are ever consumed, so keep the port narrow. - localparam TCU_WG_RS2_WIDTH = `VX_CFG_TCU_SPARSE_ENABLED ? TCU_WG_B_BLOCK_SIZE_SP : TCU_BLOCK_CAP; + localparam TCU_WG_RS2_WIDTH_DENSE = + (TCU_WG_B_BLOCK_SIZE > TCU_BLOCK_CAP) ? TCU_WG_B_BLOCK_SIZE : TCU_BLOCK_CAP; + localparam TCU_WG_RS2_WIDTH = `VX_CFG_TCU_SPARSE_ENABLED + ? ((TCU_WG_B_BLOCK_SIZE_SP > TCU_WG_RS2_WIDTH_DENSE) ? TCU_WG_B_BLOCK_SIZE_SP : TCU_WG_RS2_WIDTH_DENSE) + : TCU_WG_RS2_WIDTH_DENSE; localparam TCU_MIN_FMT_WIDTH = 4; //int4 localparam TCU_MAX_ELT_RATIO = 32 / TCU_MIN_FMT_WIDTH; @@ -171,7 +177,7 @@ package VX_tcu_pkg; input int unsigned block_elems ); automatic int unsigned sparse_ratio = `VX_CFG_TCU_SPARSE_ENABLED ? 2 : 1; - automatic int unsigned fedp_elems = TCU_TC_K * (32 / data_bits) * sparse_ratio; + automatic int unsigned fedp_elems = TCU_WG_FEDP_K * (32 / data_bits) * sparse_ratio; return (fedp_elems + block_elems - 1) / block_elems; endfunction @@ -185,9 +191,6 @@ package VX_tcu_pkg; `endif `ifdef VX_CFG_TCU_NVFP4_ENABLE max_sf = `MAX(max_sf, mx_fedp_sf_count(4, 16)); - `endif - `ifdef VX_CFG_TCU_INT8_ENABLE - max_sf = `MAX(max_sf, mx_fedp_sf_count(8, 32)); `endif return max_sf; endfunction @@ -262,8 +265,7 @@ package VX_tcu_pkg; TCU_I8_ID, TCU_U8_ID, TCU_MXFP8_ID, - TCU_MXBF8_ID, - TCU_MXI8_ID: + TCU_MXBF8_ID: return 8; TCU_FP32_ID, TCU_I32_ID, @@ -288,7 +290,7 @@ package VX_tcu_pkg; function automatic logic tcu_fmt_is_mx(input logic [TCU_FMT_WIDTH-1:0] fmt); case (fmt) - TCU_MXFP8_ID, TCU_MXBF8_ID, TCU_MXFP4_ID, TCU_NVFP4_ID, TCU_MXI8_ID: + TCU_MXFP8_ID, TCU_MXBF8_ID, TCU_MXFP4_ID, TCU_NVFP4_ID: return 1'b1; default: return 1'b0; @@ -297,19 +299,26 @@ package VX_tcu_pkg; function automatic int unsigned mx_scale_block_size(input logic [TCU_FMT_WIDTH-1:0] fmt); case (fmt) - TCU_MXFP8_ID, TCU_MXBF8_ID, TCU_MXFP4_ID, TCU_MXI8_ID: return 32; + TCU_MXFP8_ID, TCU_MXBF8_ID, TCU_MXFP4_ID: return 32; TCU_NVFP4_ID: return 16; default: return 1; endcase endfunction - function automatic int unsigned mx_scale_blocks_k(input logic [TCU_FMT_WIDTH-1:0] fmt); + function automatic int unsigned mx_scale_blocks_k_words( + input logic [TCU_FMT_WIDTH-1:0] fmt, + input int unsigned tile_k_words + ); automatic int unsigned data_bits = tcu_fmt_width(fmt); automatic int unsigned block_elems = mx_scale_block_size(fmt); - automatic int unsigned tile_elems = (data_bits != 0) ? TCU_TILE_K * (32 / data_bits) : 0; + automatic int unsigned tile_elems = (data_bits != 0) ? tile_k_words * (32 / data_bits) : 0; return (tile_elems + block_elems - 1) / block_elems; endfunction + function automatic int unsigned mx_scale_blocks_k(input logic [TCU_FMT_WIDTH-1:0] fmt); + return mx_scale_blocks_k_words(fmt, TCU_TILE_K); + endfunction + function automatic logic [4:0] meta_num_cols(input logic [TCU_FMT_WIDTH-1:0] fmt); automatic int hw = tcu_fmt_width(fmt) / 2; return 5'((TCU_BLOCK_CAP + hw - 1) / hw); @@ -343,7 +352,6 @@ package VX_tcu_pkg; TCU_U8_ID: `TRACE(level, ("u8")) TCU_I4_ID: `TRACE(level, ("i4")) TCU_U4_ID: `TRACE(level, ("u4")) - TCU_MXI8_ID: `TRACE(level, ("mxint8")) default: `TRACE(level, ("?")) endcase endtask @@ -424,6 +432,7 @@ package VX_tcu_pkg; // step_m/step_n/step_k at every consumer. logic is_first_uop; logic is_last_uop; + logic setup_fire; `ifdef VX_CFG_TCU_SPARSE_ENABLE logic is_sparse; `endif diff --git a/hw/rtl/tcu/VX_tcu_tbuf.sv b/hw/rtl/tcu/VX_tcu_tbuf.sv index bff52c6e02..6ef22b48da 100644 --- a/hw/rtl/tcu/VX_tcu_tbuf.sv +++ b/hw/rtl/tcu/VX_tcu_tbuf.sv @@ -60,7 +60,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_mem_bus_if.master tcu_lmem_if, // Per-block operand outputs (rs2 is broadcast — bbuf is shared) - output wire [BLOCK_SIZE-1:0][TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, + output wire [BLOCK_SIZE-1:0][TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, output wire [BLOCK_SIZE-1:0][TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0]tbuf_rs2_data, output wire [BLOCK_SIZE-1:0] tbuf_ready ); @@ -84,7 +84,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Per-block abufs // ----------------------------------------------------------------------- - wire [BLOCK_SIZE-1:0][TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] abuf_rs1_data_w; + wire [BLOCK_SIZE-1:0][TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] abuf_rs1_data_w; wire [BLOCK_SIZE-1:0] abuf_ready_w; `ifdef PERF_ENABLE @@ -106,11 +106,17 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `endif .req_wid (req[b].wid), .req_valid (req[b].valid), + .req_setup (req[b].setup_fire), .req_step_m (req[b].step_m), .req_step_n (req[b].step_n), .req_step_k (req[b].step_k), .req_desc_a (req[b].desc_a), .req_a_is_smem (req[b].a_is_smem), + `ifdef VX_CFG_TCU_SPARSE_ENABLE + .req_is_sparse (req[b].is_sparse), + `else + .req_is_sparse (1'b0), + `endif .req_uuid (req[b].uuid), .tcu_lmem_if (lmem_masters[b]), .abuf_ready (abuf_ready_w[b]), @@ -136,6 +142,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `endif wire bbuf_req_valid; + wire bbuf_req_setup; wire bbuf_req_is_first_uop; wire bbuf_req_is_sparse; wire [3:0] bbuf_req_step_m; @@ -148,6 +155,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( if (BLOCK_SIZE == 1) begin : g_bbuf_inputs_n1 assign bbuf_req_uuid = req[0].uuid; assign bbuf_req_valid = req[0].valid; + assign bbuf_req_setup = req[0].setup_fire; assign bbuf_req_is_first_uop = req[0].is_first_uop; `ifdef VX_CFG_TCU_SPARSE_ENABLE assign bbuf_req_is_sparse = req[0].is_sparse; @@ -161,8 +169,10 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( assign bbuf_req_desc_b = req[0].desc_b; end else begin : g_bbuf_inputs_pe wire [BLOCK_SIZE-1:0] req_valid_v; + wire [BLOCK_SIZE-1:0] req_setup_v; for (genvar b = 0; b < BLOCK_SIZE; ++b) begin : g_req_valid_v assign req_valid_v[b] = req[b].valid; + assign req_setup_v[b] = req[b].setup_fire; end wire [BLOCK_SIZE-1:0] bbuf_sel_oh; @@ -213,6 +223,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( assign bbuf_req_uuid = sel_uuid; assign bbuf_req_valid = bbuf_sel_valid; + assign bbuf_req_setup = |req_setup_v; assign bbuf_req_is_first_uop = sel_is_first_uop; assign bbuf_req_is_sparse = sel_is_sparse; assign bbuf_req_step_m = sel_step_m; @@ -235,6 +246,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( .lmem_reads (bbuf_lmem_reads_w), `endif .req_valid (bbuf_req_valid), + .req_setup (bbuf_req_setup), .req_is_first_uop (bbuf_req_is_first_uop), .req_is_sparse (bbuf_req_is_sparse), .req_step_m (bbuf_req_step_m), diff --git a/hw/rtl/tcu/VX_tcu_unit.sv b/hw/rtl/tcu/VX_tcu_unit.sv index cbde0ff2d0..e16d572a83 100644 --- a/hw/rtl/tcu/VX_tcu_unit.sv +++ b/hw/rtl/tcu/VX_tcu_unit.sv @@ -151,7 +151,7 @@ module VX_tcu_unit import VX_gpu_pkg::*, VX_tcu_pkg::*; #( assign exec_data_w[bi] = core_execute_if[bi].data; end - wire [BLOCK_SIZE-1:0][TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data; + wire [BLOCK_SIZE-1:0][TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data; wire [BLOCK_SIZE-1:0][TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0] tbuf_rs2_data; wire [BLOCK_SIZE-1:0] tbuf_ready_eff; diff --git a/hw/rtl/tcu/VX_tcu_uops.sv b/hw/rtl/tcu/VX_tcu_uops.sv index 16599d3aad..974866a235 100644 --- a/hw/rtl/tcu/VX_tcu_uops.sv +++ b/hw/rtl/tcu/VX_tcu_uops.sv @@ -49,11 +49,22 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( localparam IDX_BITS = LG_N + LG_M + LG_K; `ifdef VX_CFG_TCU_WGMMA_ENABLE `ifdef VX_CFG_TCU_SPARSE_ENABLE - localparam MAX_WG_UOPS_SP = TCU_WG_UOPS / 2; - localparam CTR_W_MAX_WMMA = MAX_UOPS > TCU_WG_UOPS ? MAX_UOPS : TCU_WG_UOPS; + localparam TCU_WG_K_STEPS_SP = (TCU_WG_K_STEPS > 1) ? (TCU_WG_K_STEPS / 2) : 1; + localparam MAX_WG_UOPS_SP = TCU_WG_M_STEPS * TCU_WG_N_STEPS * TCU_WG_K_STEPS_SP; + `ifdef VX_CFG_TCU_FEDP2K + localparam MAX_WG_UOPS_DENSE = TCU_WG_UOPS + 1; + `else + localparam MAX_WG_UOPS_DENSE = TCU_WG_UOPS; + `endif + localparam CTR_W_MAX_WMMA = MAX_UOPS > MAX_WG_UOPS_DENSE ? MAX_UOPS : MAX_WG_UOPS_DENSE; localparam CTR_W_BASE = $clog2(CTR_W_MAX_WMMA > MAX_WG_UOPS_SP ? CTR_W_MAX_WMMA : MAX_WG_UOPS_SP); `else - localparam CTR_W_BASE = $clog2(MAX_UOPS > TCU_WG_UOPS ? MAX_UOPS : TCU_WG_UOPS); + `ifdef VX_CFG_TCU_FEDP2K + localparam MAX_WG_UOPS_DENSE = TCU_WG_UOPS + 1; + `else + localparam MAX_WG_UOPS_DENSE = TCU_WG_UOPS; + `endif + localparam CTR_W_BASE = $clog2(MAX_UOPS > MAX_WG_UOPS_DENSE ? MAX_UOPS : MAX_WG_UOPS_DENSE); `endif `else localparam CTR_W_BASE = $clog2(MAX_UOPS); @@ -78,11 +89,24 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( || (ibuf_in.op_type == INST_TCU_WGMMA_SP) `endif ; +`ifdef VX_CFG_TCU_SPARSE_ENABLE + wire wg_is_sparse = (ibuf_in.op_type == INST_TCU_WGMMA_SP); +`endif wire wg_a_from_smem = ibuf_in.op_args.tcu.a_from_smem; + wire wg_needs_setup = + `ifdef VX_CFG_TCU_FEDP2K + !wg_a_from_smem + `ifdef VX_CFG_TCU_SPARSE_ENABLE + && !wg_is_sparse + `endif + `else + 1'b0 + `endif + ; // Variable NRC based on cd_nregs: 0→8, 1→16, 2→32 // Loop order: m (inner) → n → k (outer) [K-outer] - // m_steps=2 and k_steps=2 are fixed; n varies (middle). + // m_steps=2; k_steps is 2 normally and 1 with FEDP2K. // K-outer lets independent (m,n) tiles overlap FEDP latency. localparam LG_M_WG = $clog2(TCU_WG_M_STEPS); // 1 localparam LG_K_WG = $clog2(TCU_WG_K_STEPS); // 1 @@ -93,19 +117,36 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( localparam WG_MN_NR16 = 16; // nrc=16: n_steps=8, mn=16 localparam WG_MN_NR32 = 32; // nrc=32: n_steps=16, mn=32 - // Pre-computed total uop counts for each cd_nregs value (dense) + // Pre-computed compute uop counts for each cd_nregs value (dense) localparam WG_UOPS_NR8 = WG_MN_NR8 * TCU_WG_K_STEPS; localparam WG_UOPS_NR16 = WG_MN_NR16 * TCU_WG_K_STEPS; localparam WG_UOPS_NR32 = WG_MN_NR32 * TCU_WG_K_STEPS; +`ifdef VX_CFG_TCU_SPARSE_ENABLE + localparam WG_UOPS_SP_NR8 = WG_MN_NR8 * TCU_WG_K_STEPS_SP; + localparam WG_UOPS_SP_NR16 = WG_MN_NR16 * TCU_WG_K_STEPS_SP; + localparam WG_UOPS_SP_NR32 = WG_MN_NR32 * TCU_WG_K_STEPS_SP; +`endif // Mux uop count based on cd_nregs: 0→8, 1→16, 2→32 reg [UOP_CTR_W-1:0] wg_uop_cnt; + reg [UOP_CTR_W-1:0] wg_compute_uop_cnt; always_comb begin case (ibuf_in.op_args.tcu.cd_nregs) - 2'd0: wg_uop_cnt = UOP_CTR_W'(WG_UOPS_NR8); - 2'd1: wg_uop_cnt = UOP_CTR_W'(WG_UOPS_NR16); - default: wg_uop_cnt = UOP_CTR_W'(WG_UOPS_NR32); + 2'd0: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_NR8); + 2'd1: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_NR16); + default: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_NR32); endcase + `ifdef VX_CFG_TCU_SPARSE_ENABLE + if (wg_is_sparse) begin + case (ibuf_in.op_args.tcu.cd_nregs) + 2'd0: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_SP_NR8); + 2'd1: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_SP_NR16); + default: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_SP_NR32); + endcase + wg_uop_cnt = wg_compute_uop_cnt; + end else + `endif + wg_uop_cnt = wg_compute_uop_cnt + UOP_CTR_W'(wg_needs_setup); end // K-outer index extraction: @@ -113,13 +154,14 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( // Sparse: ctr = n * m_steps + m (k always 0) // Since n_steps varies by cd_nregs, k-index bit position shifts. // m is always bit 0 (m_steps=2). n and k extracted via mux. - wire [`UP(CTR_W)-1:0] wg_idx_ctr = ctr; + wire is_wg_setup_uop = wg_needs_setup && (ctr == '0); + wire [`UP(CTR_W)-1:0] wg_idx_ctr = wg_needs_setup + ? (is_wg_setup_uop ? '0 : (ctr - `UP(CTR_W)'(1))) + : ctr; + wire is_wg_first_compute_uop = !is_wg_setup_uop && (wg_idx_ctr == '0); wire [`UP(LG_M_WG)-1:0] wg_m_index = wg_idx_ctr[0 +: `UP(LG_M_WG)]; reg [`UP(LG_K_WG)-1:0] wg_k_index; reg [`UP(LG_N_WG_MAX)-1:0] wg_n_index; -`ifdef VX_CFG_TCU_SPARSE_ENABLE - wire wg_is_sparse = (ibuf_in.op_type == INST_TCU_WGMMA_SP); -`endif always_comb begin `ifdef VX_CFG_TCU_SPARSE_ENABLE if (wg_is_sparse) begin @@ -154,10 +196,9 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( localparam [4:0] wg_ra_base = TCU_WG_RA; // Register offsets for from-reg mode - // Dense A: rs1_off = m * k_steps + k (NRA=4 regs at ra_base: f24..f27) + // Dense A: rs1_off = m * 2 + k (NRA=4 regs at ra_base: f24..f27) // Sparse A: rs1_off = m (NRA compressed to 2 regs f24,f25; f26,f27 hold metadata) - localparam LG_WG_A_SB = $clog2(`UP(TCU_WG_A_SUB_BLOCKS)); - wire [`UP(CTR_W)-1:0] wg_rs1_reg_off_dense = ((`UP(CTR_W)'(wg_m_index) >> LG_WG_A_SB) << `UP(LG_K_WG)) | `UP(CTR_W)'(wg_k_index); + wire [`UP(CTR_W)-1:0] wg_rs1_reg_off_dense = (`UP(CTR_W)'(wg_m_index) << 1) | `UP(CTR_W)'(wg_k_index); `ifdef VX_CFG_TCU_SPARSE_ENABLE wire [`UP(CTR_W)-1:0] wg_rs1_reg_off = (wg_is_sparse && !wg_a_from_smem) ? `UP(CTR_W)'(wg_m_index) @@ -179,15 +220,11 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( ; `endif - // Sparse WGMMA (RS) expands to wg_uop_cnt/2 MMA uops; WMMA_SP expands - // to TCU_UOPS (sym) or TCU_UOPS/2 (asym). + // Dense FEDP2K RS adds one descriptor setup uop. Other WGMMA modes fuse + // descriptor transport into their first compute uop. assign uop_count = `ifdef VX_CFG_TCU_WGMMA_ENABLE - is_wgmma ? ( - `ifdef VX_CFG_TCU_SPARSE_ENABLE - wg_is_sparse ? (wg_uop_cnt >> 1) : - `endif - wg_uop_cnt) : + is_wgmma ? wg_uop_cnt : `endif `ifdef VX_CFG_TCU_SPARSE_ENABLE is_sparse ? (SYM_SPARSE ? UOP_CTR_W'(TCU_UOPS) : UOP_CTR_W'(TCU_UOPS / 2)) : @@ -312,25 +349,36 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( `endif `ifdef VX_CFG_TCU_WGMMA_ENABLE if (is_wgmma) begin - ibuf_r.op_args.tcu.step_m = 4'(wg_m_index); - ibuf_r.op_args.tcu.step_n = 4'(wg_n_index); - ibuf_r.op_args.tcu.step_k = 4'(wg_k_index); - ibuf_r.wb = 1'b1; + ibuf_r.op_args.tcu.step_m = is_wg_setup_uop ? 4'd0 : 4'(wg_m_index); + ibuf_r.op_args.tcu.step_n = is_wg_setup_uop ? 4'd0 : 4'(wg_n_index); + ibuf_r.op_args.tcu.step_k = is_wg_setup_uop ? 4'd0 : 4'(wg_k_index); + ibuf_r.wb = !is_wg_setup_uop; ibuf_r.rd = make_reg_num(REG_TYPE_F, TCU_WG_RC + wg_rs3_off); ibuf_r.rs3 = make_reg_num(REG_TYPE_F, TCU_WG_RC + wg_rs3_off); - ibuf_r.used_rs[2] = 1'b1; - // Smem descriptors are invariant across the whole WGMMA expansion, - // so only fetch them on the first MMA uop. - if (wg_a_from_smem) begin + ibuf_r.used_rs[2] = !is_wg_setup_uop; + if (is_wg_setup_uop) begin + ibuf_r.used_rs[0] = 1'b0; + end else if (is_wg_first_compute_uop && wg_a_from_smem) begin + ibuf_r.rs1 = make_reg_num(REG_TYPE_I, 5'd10); + ibuf_r.used_rs[0] = 1'b1; + end else if (wg_a_from_smem) begin ibuf_r.rs1 = make_reg_num(REG_TYPE_I, 5'd10); - ibuf_r.used_rs[0] = (wg_idx_ctr == '0); + ibuf_r.used_rs[0] = 1'b0; end else begin ibuf_r.rs1 = make_reg_num(REG_TYPE_F, wg_ra_base + 5'(wg_rs1_reg_off)); ibuf_r.used_rs[0] = 1'b1; end - // B source: always smem descriptor (x11), fetched on first MMA uop + // Dense FEDP2K RS reserves rs2 for upper A and transports the B + // descriptor in its setup uop. Other modes read x11 on first compute. ibuf_r.rs2 = make_reg_num(REG_TYPE_I, 5'd11); - ibuf_r.used_rs[1] = (wg_idx_ctr == '0); + if (is_wg_setup_uop || (!wg_needs_setup && is_wg_first_compute_uop)) begin + ibuf_r.used_rs[1] = 1'b1; + end else if (wg_needs_setup) begin + ibuf_r.rs2 = make_reg_num(REG_TYPE_F, wg_ra_base + 5'(wg_rs1_reg_off) + 5'd1); + ibuf_r.used_rs[1] = 1'b1; + end else begin + ibuf_r.used_rs[1] = 1'b0; + end end else `endif begin @@ -374,14 +422,11 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( ibuf_r.fu_unlock = (uop_idx == (uop_count - UOP_CTR_W'(1))); // Expose first/last-uop predicates via op_args.tcu so downstream // consumers (bbuf, lockstep) read a single source of truth. - ibuf_r.op_args.tcu.is_first_uop = (uop_idx == '0); - ibuf_r.op_args.tcu.is_last_uop = (uop_idx == (uop_count - UOP_CTR_W'(1))); - end else - `endif - begin - ibuf_r.fu_lock = 1'b0; - ibuf_r.fu_unlock = 1'b0; + ibuf_r.op_args.tcu.is_first_uop = is_wg_first_compute_uop; + ibuf_r.op_args.tcu.is_last_uop = !is_wg_setup_uop && (uop_idx == (uop_count - UOP_CTR_W'(1))); end + `endif + // Other TCU uops retain 11 so they cannot enter a locked WGMMA sequence. end assign ibuf_out = ibuf_r; diff --git a/hw/rtl/tcu/VX_tcu_wgmma.sv b/hw/rtl/tcu/VX_tcu_wgmma.sv index 8b9d5fc9dd..2765f83431 100644 --- a/hw/rtl/tcu/VX_tcu_wgmma.sv +++ b/hw/rtl/tcu/VX_tcu_wgmma.sv @@ -46,7 +46,7 @@ module VX_tcu_wgmma import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_mem_bus_if.master tcu_lmem_if, // Outputs to tcu_core (consumed by the wrapper). - output wire [BLOCK_SIZE-1:0][TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, + output wire [BLOCK_SIZE-1:0][TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, output wire [BLOCK_SIZE-1:0][TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0] tbuf_rs2_data, output wire [BLOCK_SIZE-1:0] tbuf_ready_eff ); @@ -62,17 +62,51 @@ module VX_tcu_wgmma import VX_gpu_pkg::*, VX_tcu_pkg::*; #( wire [BLOCK_SIZE-1:0] exec_fire_b_w; wire [BLOCK_SIZE-1:0] is_first_uop_b_w; wire [BLOCK_SIZE-1:0] is_last_uop_b_w; + wire [BLOCK_SIZE-1:0] is_setup_uop_b_w; + wire [BLOCK_SIZE-1:0] needs_setup_b_w; + + logic [BLOCK_SIZE-1:0][`VX_CFG_XLEN-1:0] desc_a_r; + logic [BLOCK_SIZE-1:0][`VX_CFG_XLEN-1:0] desc_b_r; + for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_lockstep_inputs wire is_wgmma_op = (exec_data[bi].op_type == INST_TCU_WGMMA) `ifdef VX_CFG_TCU_SPARSE_ENABLE || (exec_data[bi].op_type == INST_TCU_WGMMA_SP) `endif ; - assign is_wgmma_b_w[bi] = exec_valid[bi] && is_wgmma_op; + assign is_setup_uop_b_w[bi] = is_wgmma_op && !exec_data[bi].header.wb; + assign needs_setup_b_w[bi] = + `ifdef VX_CFG_TCU_FEDP2K + is_wgmma_op && !exec_data[bi].op_args.tcu.a_from_smem + `ifdef VX_CFG_TCU_SPARSE_ENABLE + && (exec_data[bi].op_type != INST_TCU_WGMMA_SP) + `endif + `else + 1'b0 + `endif + ; + assign is_wgmma_b_w[bi] = exec_valid[bi] && is_wgmma_op && !is_setup_uop_b_w[bi]; assign new_cta_b_w[bi] = exec_data[bi].header.cta_id; assign exec_fire_b_w[bi] = exec_valid[bi] && exec_ready[bi]; - assign is_first_uop_b_w[bi] = exec_data[bi].op_args.tcu.is_first_uop; - assign is_last_uop_b_w[bi] = exec_data[bi].op_args.tcu.is_last_uop; + assign is_first_uop_b_w[bi] = is_wgmma_op && exec_data[bi].op_args.tcu.is_first_uop; + assign is_last_uop_b_w[bi] = is_wgmma_op && exec_data[bi].op_args.tcu.is_last_uop; + end + + always_ff @(posedge clk) begin + if (reset) begin + desc_a_r <= '0; + desc_b_r <= '0; + end else begin + for (int bi = 0; bi < BLOCK_SIZE; ++bi) begin + if (exec_valid[bi] && exec_ready[bi] + && (is_setup_uop_b_w[bi] + || (is_first_uop_b_w[bi] && !needs_setup_b_w[bi]))) begin + if (exec_data[bi].op_args.tcu.a_from_smem) + desc_a_r[bi] <= exec_data[bi].rs1_data[0]; + desc_b_r[bi] <= exec_data[bi].rs2_data[0]; + end + end + end end VX_tcu_lockstep #( @@ -102,18 +136,23 @@ module VX_tcu_wgmma import VX_gpu_pkg::*, VX_tcu_pkg::*; #( || (exec_data[bi].op_type == INST_TCU_WGMMA_SP) `endif ; - assign req[bi].valid = exec_valid[bi] && is_wgmma_b && !cta_conflict[bi]; + wire is_setup_uop = is_wgmma_b && !exec_data[bi].header.wb; + assign req[bi].valid = exec_valid[bi] && is_wgmma_b && !is_setup_uop + && !cta_conflict[bi]; assign req[bi].uuid = exec_data[bi].header.uuid; assign req[bi].wid = exec_data[bi].header.wid; assign req[bi].step_m = exec_data[bi].op_args.tcu.step_m; assign req[bi].step_k = exec_data[bi].op_args.tcu.step_k; assign req[bi].step_n = exec_data[bi].op_args.tcu.step_n; assign req[bi].cd_nregs = exec_data[bi].op_args.tcu.cd_nregs; - assign req[bi].desc_a = exec_data[bi].rs1_data[0]; - assign req[bi].desc_b = exec_data[bi].rs2_data[0]; + wire use_live_desc = is_first_uop_b_w[bi] + && !needs_setup_b_w[bi]; + assign req[bi].desc_a = use_live_desc ? exec_data[bi].rs1_data[0] : desc_a_r[bi]; + assign req[bi].desc_b = use_live_desc ? exec_data[bi].rs2_data[0] : desc_b_r[bi]; assign req[bi].a_is_smem = exec_data[bi].op_args.tcu.a_from_smem; assign req[bi].is_first_uop = exec_data[bi].op_args.tcu.is_first_uop; assign req[bi].is_last_uop = exec_data[bi].op_args.tcu.is_last_uop; + assign req[bi].setup_fire = exec_valid[bi] && exec_ready[bi] && is_setup_uop; `ifdef VX_CFG_TCU_SPARSE_ENABLE assign req[bi].is_sparse = (exec_data[bi].op_type == INST_TCU_WGMMA_SP); `endif @@ -148,7 +187,8 @@ module VX_tcu_wgmma import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ); for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_tbuf_eff - assign tbuf_ready_eff[bi] = tbuf_ready[bi] && !cta_conflict[bi]; + assign tbuf_ready_eff[bi] = is_setup_uop_b_w[bi] ? 1'b1 + : (tbuf_ready[bi] && !cta_conflict[bi]); end // ----------------------------------------------------------------------- diff --git a/hw/rtl/tcu/dpi/VX_tcu_fedp_dpi.sv b/hw/rtl/tcu/dpi/VX_tcu_fedp_dpi.sv index f587e21cc0..1d3b8c94da 100644 --- a/hw/rtl/tcu/dpi/VX_tcu_fedp_dpi.sv +++ b/hw/rtl/tcu/dpi/VX_tcu_fedp_dpi.sv @@ -49,6 +49,20 @@ module VX_tcu_fedp_dpi import VX_tcu_pkg::*; #( `UNUSED_VAR (fmt_d); +`ifdef VX_CFG_TCU_MX_ENABLE + function automatic [7:0] mx_scale_exp_sum( + input logic [7:0] scale_a, + input logic [7:0] scale_b + ); + logic [7:0] raw_scale_a, raw_scale_b; + begin + raw_scale_a = scale_a - 8'd127; + raw_scale_b = scale_b - 8'd127; + mx_scale_exp_sum = raw_scale_a + raw_scale_b; + end + endfunction +`endif + wire [31:0] mult_result [N]; // multiplication stage @@ -57,20 +71,7 @@ module VX_tcu_fedp_dpi import VX_tcu_pkg::*; #( reg [63:0] temp, prod; reg [4:0] fflags; `ifdef VX_CFG_TCU_MX_ENABLE - `ifdef VX_CFG_TCU_FP8_ENABLE reg [63:0] group_prod; - reg [7:0] raw_sf_a, raw_sf_b, raw_sf; - `elsif VX_CFG_TCU_FP4_ENABLE - reg [63:0] group_prod; - `endif - `ifdef VX_CFG_TCU_FP4_ENABLE - `ifdef VX_CFG_TCU_NVFP4_ENABLE - reg [63:0] a_sf, b_sf, temp_sf; - `endif - `ifdef VX_CFG_TCU_MXFP4_ENABLE - reg [7:0] raw_sf_a_f4, raw_sf_b_f4, raw_sf_f4; - `endif - `endif `endif `UNUSED_VAR({fflags, prod[63:32]}); @@ -139,11 +140,8 @@ module VX_tcu_fedp_dpi import VX_tcu_pkg::*; #( dpi_fadd(enable, int'(0), temp, group_prod, 3'b0, group_prod, fflags); end end - raw_sf_a = sf_a[s] - 8'd127; - raw_sf_b = sf_b[s] - 8'd127; - raw_sf = raw_sf_a + raw_sf_b; if (group_prod[30:0] != 0) begin - group_prod[30:23] = group_prod[30:23] + raw_sf; + group_prod[30:23] = group_prod[30:23] + mx_scale_exp_sum(sf_a[s], sf_b[s]); end dpi_fadd(enable, int'(0), group_prod, prod, 3'b0, prod, fflags); end @@ -160,11 +158,8 @@ module VX_tcu_fedp_dpi import VX_tcu_pkg::*; #( dpi_fadd(enable, int'(0), temp, group_prod, 3'b0, group_prod, fflags); end end - raw_sf_a = sf_a[s] - 8'd127; - raw_sf_b = sf_b[s] - 8'd127; - raw_sf = raw_sf_a + raw_sf_b; if (group_prod[30:0] != 0) begin - group_prod[30:23] = group_prod[30:23] + raw_sf; + group_prod[30:23] = group_prod[30:23] + mx_scale_exp_sum(sf_a[s], sf_b[s]); end dpi_fadd(enable, int'(0), group_prod, prod, 3'b0, prod, fflags); end @@ -184,11 +179,8 @@ module VX_tcu_fedp_dpi import VX_tcu_pkg::*; #( dpi_fadd(enable, int'(0), temp, group_prod, 3'b0, group_prod, fflags); end end - raw_sf_a_f4 = sf_a[s] - 8'd127; - raw_sf_b_f4 = sf_b[s] - 8'd127; - raw_sf_f4 = raw_sf_a_f4 + raw_sf_b_f4; if (group_prod[30:0] != 0) begin - group_prod[30:23] = group_prod[30:23] + raw_sf_f4; + group_prod[30:23] = group_prod[30:23] + mx_scale_exp_sum(sf_a[s], sf_b[s]); end dpi_fadd(enable, int'(0), group_prod, prod, 3'b0, prod, fflags); end @@ -207,10 +199,10 @@ module VX_tcu_fedp_dpi import VX_tcu_pkg::*; #( dpi_fadd(enable, int'(0), temp, group_prod, 3'b0, group_prod, fflags); end end - dpi_f2f(enable, int'(0), int'(4), {56'hffffffffffffff, sf_a[s]}, 3'b0, a_sf, fflags); - dpi_f2f(enable, int'(0), int'(4), {56'hffffffffffffff, sf_b[s]}, 3'b0, b_sf, fflags); - dpi_fmul(enable, int'(0), a_sf, b_sf, 3'b0, temp_sf, fflags); - dpi_fmul(enable, int'(0), group_prod, temp_sf, 3'b0, group_prod, fflags); + dpi_f2f(enable, int'(0), int'(4), {56'hffffffffffffff, sf_a[s]}, 3'b0, a_f, fflags); + dpi_f2f(enable, int'(0), int'(4), {56'hffffffffffffff, sf_b[s]}, 3'b0, b_f, fflags); + dpi_fmul(enable, int'(0), a_f, b_f, 3'b0, temp, fflags); + dpi_fmul(enable, int'(0), group_prod, temp, 3'b0, group_prod, fflags); dpi_fadd(enable, int'(0), group_prod, prod, 3'b0, prod, fflags); end end @@ -230,34 +222,6 @@ module VX_tcu_fedp_dpi import VX_tcu_pkg::*; #( prod += a_row[i][8 * j +: 8] * b_col[i][8 * j +: 8]; end end - `ifdef VX_CFG_TCU_MX_ENABLE - TCU_MXI8_ID: begin - prod = 0; - for (int j = 0; j < 4; j++) begin - integer sf_slot; - reg signed [31:0] raw_prod; - reg [31:0] abs_prod; - reg signed [31:0] scaled_prod; - reg signed [8:0] combined_sf; - reg [8:0] shift_amt; - raw_prod = $signed({{24{a_row[i][8 * j + 7]}}, a_row[i][8 * j +: 8]}) * $signed({{24{b_col[i][8 * j + 7]}}, b_col[i][8 * j +: 8]}); - sf_slot = ((i * 4 + j) * SF) / (N * 4); - combined_sf = $signed({1'b0, sf_a[sf_slot]}) + $signed({1'b0, sf_b[sf_slot]}) - 9'sd266; - if (combined_sf[8]) begin - // Negative: right shift with truncation toward zero - shift_amt = -combined_sf; - abs_prod = raw_prod[31] ? (-raw_prod) : raw_prod; - scaled_prod = abs_prod >> shift_amt; - if (raw_prod[31]) begin - scaled_prod = -scaled_prod; - end - end else begin - scaled_prod = raw_prod <<< combined_sf; - end - prod += 64'($signed(scaled_prod)); - end - end - `endif `endif `ifdef VX_CFG_TCU_INT4_ENABLE TCU_I4_ID: begin diff --git a/hw/rtl/tcu/tfr/VX_tcu_tfr_lane_mask.sv b/hw/rtl/tcu/tfr/VX_tcu_tfr_lane_mask.sv index 1388640630..4276ef295d 100644 --- a/hw/rtl/tcu/tfr/VX_tcu_tfr_lane_mask.sv +++ b/hw/rtl/tcu/tfr/VX_tcu_tfr_lane_mask.sv @@ -87,9 +87,6 @@ module VX_tcu_tfr_lane_mask import VX_tcu_pkg::*; #( `ifdef VX_CFG_TCU_INT8_ENABLE TCU_I8_ID, TCU_U8_ID: lane_mask = mask_8; - `ifdef VX_CFG_TCU_MX_ENABLE - TCU_MXI8_ID: lane_mask = mask_8; - `endif `endif `ifdef VX_CFG_TCU_INT4_ENABLE TCU_I4_ID, diff --git a/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_f8.sv b/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_f8.sv index 79e6962037..b9b32a4df5 100644 --- a/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_f8.sv +++ b/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_f8.sv @@ -190,20 +190,31 @@ module VX_tcu_tfr_mul_f8 import VX_tcu_pkg::*; wire [EXP_W-1:0] final_exp; `ifdef VX_CFG_TCU_MX_ENABLE + wire [EXP_W-1:0] bias_cpa_sum, bias_cpa_carry; wire [3*EXP_W-1:0] sf_comp = fmt_f[3] ? {EXP_W'(sf_a), EXP_W'(sf_b), -EXP_W'(254)} : (3*EXP_W)'(0); VX_csa_tree #( - .N(5), + .N(4), .W(EXP_W), .S(EXP_W) ) exp_sf_csa ( - .operands ({EXP_W'(max_pre_sum), EXP_W'(bias_sel), sf_comp}), - .sum (max_pre_sum_cpa), - .carry (bias_sel_cpa) + .operands ({EXP_W'(bias_sel), sf_comp}), + .sum (bias_cpa_sum), + .carry (bias_cpa_carry) + ); + VX_ks_adder #( + .N(EXP_W), + .BYPASS(`FORCE_BUILTIN_ADDER(EXP_W)) + ) exp_bias_add ( + .dataa(bias_cpa_sum), + .datab(bias_cpa_carry), + .cin(1'b0), + .sum(bias_sel_cpa), + `UNUSED_PIN(cout) ); `else - assign max_pre_sum_cpa = EXP_W'(max_pre_sum); assign bias_sel_cpa = EXP_W'(bias_sel); `endif + assign max_pre_sum_cpa = EXP_W'(max_pre_sum); VX_ks_adder #( .N(EXP_W), .BYPASS(`FORCE_BUILTIN_ADDER(EXP_W)) diff --git a/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_i8.sv b/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_i8.sv index 08c7509a05..acffc293b0 100644 --- a/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_i8.sv +++ b/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_i8.sv @@ -28,10 +28,6 @@ module VX_tcu_tfr_mul_i8 import VX_tcu_pkg::*; #( input wire [N-1:0][31:0] a_row, input wire [N-1:0][31:0] b_col, -`ifdef VX_CFG_TCU_MX_ENABLE - input wire [7:0] sf_a, - input wire [7:0] sf_b, -`endif output logic [TCK-1:0][24:0] result ); @@ -76,41 +72,11 @@ module VX_tcu_tfr_mul_i8 import VX_tcu_pkg::*; #( `UNUSED_PIN(cout) ); -`ifdef VX_CFG_TCU_MX_ENABLE - wire signed [8:0] combined_sf = $signed(sf_a + sf_b - 9'd266); - wire is_right_shift = combined_sf[8]; - wire shift_overflow = (combined_sf > 9'sd24) || (combined_sf < -9'sd24); - wire [4:0] shift_amount = is_right_shift ? (-combined_sf[4:0]) : combined_sf[4:0]; - - wire signed [24:0] y_mxi8_scaled [2]; - for (genvar j = 0; j < 2; ++j) begin : g_mxi8 - wire signed [24:0] raw_prod = {{8{y_prod_i8[j][16]}}, y_prod_i8[j]}; - assign y_mxi8_scaled[j] = shift_overflow ? 25'sd0 - : is_right_shift ? (raw_prod >>> shift_amount) - : (raw_prod <<< shift_amount); - end - - wire [24:0] y_mxi8_add_res; - VX_ks_adder #( - .N (25), - .BYPASS (`FORCE_BUILTIN_ADDER(25)) - ) mxi8_ksa ( - .cin (1'b0), - .dataa (y_mxi8_scaled[0]), - .datab (y_mxi8_scaled[1]), - .sum (y_mxi8_add_res), - `UNUSED_PIN(cout) - ); -`endif - // Output muxing always_comb begin case ({1'b1, fmt_i}) TCU_I8_ID: result[i] = 25'($signed(y_i8_add_res)); TCU_U8_ID: result[i] = {8'b0, y_i8_add_res}; - `ifdef VX_CFG_TCU_MX_ENABLE - TCU_MXI8_ID: result[i] = y_mxi8_add_res; - `endif default: result[i] = '0; endcase end diff --git a/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_join.sv b/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_join.sv index e0f30e2066..f191fefbfe 100644 --- a/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_join.sv +++ b/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_join.sv @@ -126,9 +126,6 @@ module VX_tcu_tfr_mul_join import VX_tcu_pkg::*; #( `ifdef VX_CFG_TCU_INT8_ENABLE TCU_I8_ID, TCU_U8_ID - `ifdef VX_CFG_TCU_MX_ENABLE - , TCU_MXI8_ID - `endif : begin sig_sel = sig_int8; exp_sel = '0; diff --git a/hw/rtl/tcu/tfr/VX_tcu_tfr_shared_mul.sv b/hw/rtl/tcu/tfr/VX_tcu_tfr_shared_mul.sv index 2ca075cd12..cfb7e58ba9 100644 --- a/hw/rtl/tcu/tfr/VX_tcu_tfr_shared_mul.sv +++ b/hw/rtl/tcu/tfr/VX_tcu_tfr_shared_mul.sv @@ -175,33 +175,22 @@ module VX_tcu_tfr_shared_mul import VX_tcu_pkg::*; #( `endif `ifdef VX_CFG_TCU_INT8_ENABLE - // I8 / U8 / MXI8 + // I8 / U8 wire [TCK-1:0][24:0] mul_int8_sig; - wire [SF-1:0][TCK-1:0][24:0] mul_int8_sig_s; - for (genvar s = 0; s < SF; ++s) begin : g_mul_i8_sf - VX_tcu_tfr_mul_i8 #( - .N(N), - .TCK(TCK), - .USE_DSP(USE_DSP) - ) mul_int8 ( - .clk (clk), - .valid_in (valid_in), - .req_id (req_id), - .vld_mask (vld_mask), - .fmt_i (fmt_s[3:0]), - .a_row (a_row), - .b_col (b_col), - `ifdef VX_CFG_TCU_MX_ENABLE - .sf_a (sf_a[s]), - .sf_b (sf_b[s]), - `endif - .result (mul_int8_sig_s[s]) - ); - end - for (genvar i = 0; i < TCK; ++i) begin : g_mul_i8_lane - localparam SF_SLOT = (i * SF) / TCK; - assign mul_int8_sig[i] = mul_int8_sig_s[SF_SLOT][i]; - end + VX_tcu_tfr_mul_i8 #( + .N(N), + .TCK(TCK), + .USE_DSP(USE_DSP) + ) mul_int8 ( + .clk (clk), + .valid_in (valid_in), + .req_id (req_id), + .vld_mask (vld_mask), + .fmt_i (fmt_s[3:0]), + .a_row (a_row), + .b_col (b_col), + .result (mul_int8_sig) + ); `endif `ifdef VX_CFG_TCU_INT4_ENABLE diff --git a/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_desc_uop.csv b/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_desc_uop.csv new file mode 100644 index 0000000000..b463e5db18 --- /dev/null +++ b/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_desc_uop.csv @@ -0,0 +1,13 @@ +model,fedp_latency,mode,threads,warps,issue_width,itype,otype,idle_pct,occupancy,occupancy_pct,simt_util,simt_util_pct,stall_fetch_pct,stall_ibuf_pct,stall_scrb_pct,stall_opds_pct,stall_alu_pct,stall_lsu_pct,stall_sfu_pct,stall_fpu_pct,stall_tcu_pct,mix_alu_pct,mix_lsu_pct,mix_sfu_pct,mix_fpu_pct,mix_tcu_pct,branches,divergent_branches,divergence_pct,ifetch_latency,load_latency,loads,stores,instrs,cycles,ipc,status +baseline,4,RS,8,8,4,fp16,fp32,32,8.0,100,8.0,100,0,0,77,2,6,1,0,0,2,84,8,3,0,5,6336,0,0,4.02,9.39,59904,45568,168032,234768,0.716,PASS +baseline,4,RS,32,8,4,fp16,fp32,43,8.0,100,32.0,100,0,0,66,1,5,4,1,0,2,83,8,3,1,5,816,0,0,4.68,22.53,31232,25088,21568,36126,0.597,PASS +baseline,4,SS,8,8,4,fp16,fp32,32,8.0,100,8.0,100,0,0,73,2,6,0,0,0,3,84,7,3,0,5,6336,0,0,4.02,9.88,43520,45568,160864,224501,0.717,PASS +baseline,4,SS,32,8,4,fp16,fp32,44,8.0,100,32.0,100,0,0,69,1,5,4,1,0,4,83,7,3,1,5,816,0,0,4.71,27.09,23040,25088,20672,35326,0.585,PASS +fedp2k,4,RS,8,8,4,fp16,fp32,31,8.0,100,8.0,100,0,0,71,3,5,2,1,0,1,79,12,4,1,4,6336,0,0,4.1,9.69,59904,45568,110944,155274,0.715,PASS +fedp2k,4,RS,32,8,4,fp16,fp32,47,8.0,100,32.0,100,0,0,56,3,4,6,1,0,1,78,12,5,1,4,816,0,0,5.02,22.5,31232,25088,14208,25786,0.551,PASS +fedp2k,4,SS,8,8,4,fp16,fp32,30,8.0,100,8.0,100,0,0,71,3,5,1,0,0,3,79,11,5,1,5,6336,0,0,4.03,9.86,43520,45568,100128,137848,0.726,PASS +fedp2k,4,SS,32,8,4,fp16,fp32,48,7.9,99,32.0,100,0,0,53,1,3,6,1,0,4,78,11,5,1,4,816,0,0,5.1,27.13,23040,25088,13104,24392,0.537,PASS +fedp2k,8,RS,8,8,4,fp16,fp32,31,8.0,100,8.0,100,0,0,73,3,5,1,1,0,2,79,12,4,1,4,6336,0,0,4.1,9.61,59904,45568,110944,155220,0.715,PASS +fedp2k,8,RS,32,8,4,fp16,fp32,47,8.0,100,32.0,100,0,0,56,3,4,6,1,0,1,78,12,5,1,4,816,0,0,5.02,22.5,31232,25088,14208,25786,0.551,PASS +fedp2k,8,SS,8,8,4,fp16,fp32,30,8.0,100,8.0,100,0,0,68,3,5,1,1,0,4,79,11,5,1,5,6336,0,0,4.03,10.05,43520,45568,100128,137867,0.726,PASS +fedp2k,8,SS,32,8,4,fp16,fp32,48,8.0,99,32.0,100,0,0,61,1,3,6,1,0,4,78,11,5,1,4,816,0,0,5.1,27.26,23040,25088,13104,24291,0.539,PASS diff --git a/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_tcu_meta_ld.csv b/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_tcu_meta_ld.csv new file mode 100644 index 0000000000..cd1cd4e526 --- /dev/null +++ b/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_tcu_meta_ld.csv @@ -0,0 +1,13 @@ +model,fedp_latency,mode,threads,warps,issue_width,itype,otype,idle_pct,occupancy,occupancy_pct,simt_util,simt_util_pct,stall_fetch_pct,stall_ibuf_pct,stall_scrb_pct,stall_opds_pct,stall_alu_pct,stall_lsu_pct,stall_sfu_pct,stall_fpu_pct,stall_tcu_pct,mix_alu_pct,mix_lsu_pct,mix_sfu_pct,mix_fpu_pct,mix_tcu_pct,branches,divergent_branches,divergence_pct,ifetch_latency,load_latency,loads,stores,instrs,cycles,ipc,status +baseline,4,RS,8,8,4,fp16,fp32,32,8.0,100,7.9,99,0,0,69,1,5,1,0,0,2,82,8,5,0,5,7936,512,6,4.02,9.32,59904,46080,172960,243040,0.712,PASS +baseline,4,RS,32,8,4,fp16,fp32,43,7.9,99,31.5,99,0,0,62,1,5,4,1,0,2,81,8,5,1,5,1024,64,6,4.67,21.85,31232,25152,22240,37409,0.595,PASS +baseline,4,SS,8,8,4,fp16,fp32,33,8.0,100,7.8,98,0,0,70,1,6,0,0,0,3,82,7,5,0,6,7936,512,6,4.02,9.94,43520,46592,167392,239417,0.699,PASS +baseline,4,SS,32,8,4,fp16,fp32,43,8.0,99,31.3,98,0,0,68,1,5,4,1,0,4,81,8,5,1,5,1024,64,6,4.69,26.6,23040,25216,21552,36545,0.59,PASS +fedp2k,4,RS,8,8,4,fp16,fp32,34,8.0,100,7.8,98,0,0,66,3,5,1,1,0,1,77,12,7,0,4,7936,512,6,4.1,9.6,59904,46080,115872,169531,0.683,PASS +fedp2k,4,RS,32,8,4,fp16,fp32,47,8.0,100,31.3,98,0,0,54,2,3,5,1,0,1,76,12,7,1,4,1024,64,6,4.98,22.05,31232,25152,14880,27394,0.543,PASS +fedp2k,4,SS,8,8,4,fp16,fp32,30,7.9,99,7.7,97,0,0,61,3,4,1,1,0,2,76,11,7,1,5,7936,512,6,4.03,9.48,43520,46592,106656,147193,0.725,PASS +fedp2k,4,SS,32,8,4,fp16,fp32,48,7.9,99,31.0,97,0,0,50,2,3,5,1,0,4,75,12,8,1,5,1024,64,6,5.04,26.94,23040,25216,13984,26141,0.535,PASS +fedp2k,8,RS,8,8,4,fp16,fp32,33,8.0,100,7.8,98,0,0,66,3,5,1,1,0,2,77,12,7,0,4,7936,512,6,4.11,9.7,59904,46080,115872,167617,0.691,PASS +fedp2k,8,RS,32,8,4,fp16,fp32,47,8.0,100,31.3,98,0,0,54,2,3,5,1,0,2,76,12,7,1,4,1024,64,6,4.98,22.02,31232,25152,14880,27285,0.545,PASS +fedp2k,8,SS,8,8,4,fp16,fp32,30,7.9,99,7.7,97,0,0,60,3,4,0,1,0,3,76,11,7,1,5,7936,512,6,4.03,9.52,43520,46592,106656,147023,0.725,PASS +fedp2k,8,SS,32,8,4,fp16,fp32,48,7.9,99,31.0,97,0,0,51,2,3,5,1,0,5,75,12,8,1,5,1024,64,6,5.04,26.75,23040,25216,13984,26256,0.533,PASS diff --git a/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp16.csv b/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp16.csv new file mode 100644 index 0000000000..9e740871fb --- /dev/null +++ b/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp16.csv @@ -0,0 +1,13 @@ +model,fedp_latency,mode,threads,warps,issue_width,itype,otype,idle_pct,occupancy,occupancy_pct,simt_util,simt_util_pct,stall_fetch_pct,stall_ibuf_pct,stall_scrb_pct,stall_opds_pct,stall_alu_pct,stall_lsu_pct,stall_sfu_pct,stall_fpu_pct,stall_tcu_pct,mix_alu_pct,mix_lsu_pct,mix_sfu_pct,mix_fpu_pct,mix_tcu_pct,branches,divergent_branches,divergence_pct,ifetch_latency,load_latency,loads,stores,instrs,cycles,ipc,status +baseline,4,RS,8,8,4,fp16,fp32,60,8.0,99,8.0,100,0,0,57,0,1,9,1,0,1,71,9,16,0,4,13056,0,0,4.23,41.85,61440,36864,138592,335658,0.413,PASS +baseline,4,SS,8,8,4,fp16,fp32,61,8.0,100,8.0,100,0,0,57,0,0,9,1,0,1,70,9,17,0,4,13056,0,0,4.23,46.96,53248,36864,132384,331183,0.400,PASS +baseline,4,RS,32,8,4,fp16,fp32,90,7.9,99,32.0,100,0,0,66,0,0,17,3,0,0,69,10,16,1,4,1664,0,0,12.75,173.62,35328,23552,17760,169411,0.105,PASS +baseline,4,SS,32,8,4,fp16,fp32,90,7.9,99,32.0,100,0,0,66,0,0,18,3,0,0,68,10,17,1,4,1664,0,0,13.08,208.59,31232,23552,16992,168734,0.101,PASS +fedp2k,4,RS,8,8,4,fp16,fp32,65,8.0,100,8.0,100,0,0,59,1,0,8,2,0,0,69,10,17,0,4,13056,0,0,4.29,45.74,66048,37376,131104,360829,0.363,PASS +fedp2k,4,SS,8,8,4,fp16,fp32,68,8.0,100,8.0,100,0,0,64,0,0,9,4,0,1,68,10,18,0,4,13056,0,0,4.29,54.27,57856,37376,124960,382144,0.327,PASS +fedp2k,4,RS,32,8,4,fp16,fp32,90,7.9,98,32.0,100,0,0,64,0,0,16,4,0,0,68,11,17,1,4,1664,0,0,12.97,171.89,35328,23552,17376,169277,0.103,PASS +fedp2k,4,SS,32,8,4,fp16,fp32,91,7.8,98,32.0,100,0,0,67,0,0,18,2,0,0,67,10,18,1,4,1664,0,0,13.30,208.78,31232,23552,16608,170203,0.098,PASS +fedp2k,8,RS,8,8,4,fp16,fp32,64,7.9,99,8.0,100,0,0,60,1,0,9,2,0,1,69,10,17,0,4,13056,0,0,4.28,45.83,66048,37376,131104,356623,0.368,PASS +fedp2k,8,SS,8,8,4,fp16,fp32,68,8.0,100,8.0,100,0,0,64,0,0,9,4,0,1,68,10,18,0,4,13056,0,0,4.29,54.27,57856,37376,124960,382144,0.327,PASS +fedp2k,8,RS,32,8,4,fp16,fp32,90,7.9,98,32.0,100,0,0,64,0,0,16,4,0,0,68,11,17,1,4,1664,0,0,12.97,172.00,35328,23552,17376,169287,0.103,PASS +fedp2k,8,SS,32,8,4,fp16,fp32,91,7.9,98,32.0,100,0,0,67,0,0,17,2,0,0,67,10,18,1,4,1664,0,0,13.30,209.51,31232,23552,16608,170430,0.097,PASS diff --git a/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp8.csv b/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp8.csv new file mode 100644 index 0000000000..72313478fd --- /dev/null +++ b/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp8.csv @@ -0,0 +1,13 @@ +model,fedp_latency,mode,threads,warps,issue_width,itype,otype,idle_pct,occupancy,occupancy_pct,simt_util,simt_util_pct,stall_fetch_pct,stall_ibuf_pct,stall_scrb_pct,stall_opds_pct,stall_alu_pct,stall_lsu_pct,stall_sfu_pct,stall_fpu_pct,stall_tcu_pct,mix_alu_pct,mix_lsu_pct,mix_sfu_pct,mix_fpu_pct,mix_tcu_pct,branches,divergent_branches,divergence_pct,ifetch_latency,load_latency,loads,stores,instrs,cycles,ipc,status +baseline,4,RS,8,8,4,fp8,fp32,59,7.9,99,8.0,100,0,0,56,1,0,9,1,0,0,71,11,15,1,3,8192,0,0,4.35,36.23,45056,34816,90336,214859,0.420,PASS +baseline,4,SS,8,8,4,fp8,fp32,59,7.9,99,8.0,100,0,0,56,1,0,9,1,0,1,71,11,15,1,3,8192,0,0,4.36,39.24,40960,34816,88032,211130,0.417,PASS +baseline,4,RS,32,8,4,fp8,fp32,88,7.7,96,32.0,100,0,0,60,0,0,16,1,0,0,69,12,15,1,3,1056,0,0,14.43,138.98,24064,21504,11840,94625,0.125,PASS +baseline,4,SS,32,8,4,fp8,fp32,88,7.7,97,32.0,100,0,0,58,0,0,16,1,0,0,69,12,15,1,3,1056,0,0,14.61,138.39,22016,21504,11552,95324,0.121,PASS +fedp2k,4,RS,8,8,4,fp8,fp32,60,7.9,99,8.0,100,0,0,53,1,0,9,2,0,0,70,11,15,1,3,8192,0,0,4.36,36.50,45056,34816,87264,211279,0.413,PASS +fedp2k,4,SS,8,8,4,fp8,fp32,60,8.0,99,8.0,100,0,0,55,1,0,9,1,0,1,70,11,16,1,3,8192,0,0,4.40,39.41,40960,34816,84960,209156,0.406,PASS +fedp2k,4,RS,32,8,4,fp8,fp32,88,7.9,99,32.0,100,0,0,61,0,0,17,1,0,0,69,12,15,1,3,1056,0,0,14.43,128.49,24064,21504,11664,94822,0.123,PASS +fedp2k,4,SS,32,8,4,fp8,fp32,88,7.9,99,32.0,100,0,0,63,0,0,17,0,0,0,69,12,15,1,3,1056,0,0,14.74,145.60,22016,21504,11376,93752,0.121,PASS +fedp2k,8,RS,8,8,4,fp8,fp32,60,7.9,99,8.0,100,0,0,52,1,0,9,2,0,1,70,11,15,1,3,8192,0,0,4.36,35.99,45056,34816,87264,212092,0.411,PASS +fedp2k,8,SS,8,8,4,fp8,fp32,60,7.9,99,8.0,100,0,0,54,1,0,9,1,0,1,70,11,16,1,3,8192,0,0,4.40,38.96,40960,34816,84960,206133,0.412,PASS +fedp2k,8,RS,32,8,4,fp8,fp32,88,7.9,99,32.0,100,0,0,61,0,0,17,1,0,0,69,12,15,1,3,1056,0,0,14.43,128.49,24064,21504,11664,94866,0.123,PASS +fedp2k,8,SS,32,8,4,fp8,fp32,88,7.9,99,32.0,100,0,0,63,0,0,17,0,0,0,69,12,15,1,3,1056,0,0,14.74,147.17,22016,21504,11376,93702,0.121,PASS diff --git a/perf/results/tcu/sgemm_tcu_perf.md b/perf/results/tcu/sgemm_tcu_perf.md index f7a4b5837c..efb8641f64 100644 --- a/perf/results/tcu/sgemm_tcu_perf.md +++ b/perf/results/tcu/sgemm_tcu_perf.md @@ -30,7 +30,6 @@ CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -D CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32" ./ci/blackbox.sh --driver=simx --app=sgemm_tcu_mx --args="-m64 -n64 -k64" --perf=1 -CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxint8 -DOTYPE=int32" ./ci/blackbox.sh --driver=simx --app=sgemm_tcu_mx --args="-m64 -n64 -k64" --perf=1 # sgemm_tcu_sp_mx CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32" ./ci/blackbox.sh --driver=simx --app=sgemm_tcu_sp_mx --args="-m64 -n64 -k64" --perf=1 @@ -39,7 +38,6 @@ CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -D CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DTCU_MX_TLS -DITYPE=nvfp4 -DOTYPE=fp32" ./ci/blackbox.sh --driver=simx --app=sgemm_tcu_sp_mx --args="-m64 -n64 -k64" --perf=1 -CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxint8 -DOTYPE=int32" ./ci/blackbox.sh --driver=simx --app=sgemm_tcu_sp_mx --args="-m64 -n64 -k64" --perf=1 ``` ## RTLsim (DPI) Commands @@ -70,7 +68,6 @@ CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -D CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_TYPE_DPI -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32" ./ci/blackbox.sh --driver=rtlsim --app=sgemm_tcu_mx --args="-m64 -n64 -k64" --perf=1 -CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_TYPE_DPI -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxint8 -DOTYPE=int32" ./ci/blackbox.sh --driver=rtlsim --app=sgemm_tcu_mx --args="-m64 -n64 -k64" --perf=1 # sgemm_tcu_sp_mx CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_TYPE_DPI -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32" ./ci/blackbox.sh --driver=rtlsim --app=sgemm_tcu_sp_mx --args="-m64 -n64 -k64" --perf=1 @@ -79,7 +76,6 @@ CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -D CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_TYPE_DPI -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DTCU_MX_TLS -DITYPE=nvfp4 -DOTYPE=fp32" ./ci/blackbox.sh --driver=rtlsim --app=sgemm_tcu_sp_mx --args="-m64 -n64 -k64" --perf=1 -CONFIGS="-DVX_CFG_NUM_THREADS=8 -DVX_CFG_NUM_WARPS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_TYPE_DPI -DVX_CFG_TCU_SPARSE_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxint8 -DOTYPE=int32" ./ci/blackbox.sh --driver=rtlsim --app=sgemm_tcu_sp_mx --args="-m64 -n64 -k64" --perf=1 ``` ## FPGA Synthesis Configs diff --git a/perf/results/tcu/sgemm_tcu_rtlsim.csv b/perf/results/tcu/sgemm_tcu_rtlsim.csv index 4fe793e3fc..9dedd1a447 100644 --- a/perf/results/tcu/sgemm_tcu_rtlsim.csv +++ b/perf/results/tcu/sgemm_tcu_rtlsim.csv @@ -10,8 +10,6 @@ sgemm_tcu_sp,8,8,int4,int32,76,7.9,99,8.0,100,1,36,96,4,0,0,0,0,0,46,22,7,5,20,3 sgemm_tcu_mx,8,8,mxfp8,fp32,74,7.9,99,8.0,100,0,21,97,3,0,0,0,0,0,41,20,6,4,29,320,0,0,4.98,16.17,22016,4608,15264,38792,0.393,PASS sgemm_tcu_mx,8,8,mxfp4,fp32,78,7.9,99,8.0,100,0,17,96,4,0,0,0,0,0,42,21,7,6,23,256,0,0,5.14,21.51,12800,4608,9568,28382,0.337,PASS sgemm_tcu_mx,8,8,nvfp4,fp32,78,7.9,99,8.0,100,0,17,96,4,0,0,0,0,0,42,21,7,6,23,256,0,0,5.14,21.51,12800,4608,9568,28382,0.337,PASS -sgemm_tcu_mx,8,8,mxint8,int32,74,7.9,99,8.0,100,0,21,97,3,0,0,0,0,0,41,20,6,4,29,320,0,0,4.98,16.17,22016,4608,15264,38792,0.393,PASS sgemm_tcu_sp_mx,8,8,mxfp8,fp32,70,7.9,99,8.0,100,0,19,97,4,0,0,0,0,0,48,20,8,4,19,320,0,0,4.91,15.13,19456,4608,13024,32714,0.398,PASS sgemm_tcu_sp_mx,8,8,mxfp4,fp32,75,7.8,98,8.0,100,0,28,97,4,0,0,0,0,0,49,21,9,7,15,256,0,0,5.75,17.34,12288,4608,8736,25764,0.339,PASS sgemm_tcu_sp_mx,8,8,nvfp4,fp32,75,7.8,98,8.0,100,0,28,97,4,0,0,0,0,0,49,21,9,7,15,256,0,0,5.75,17.34,12288,4608,8736,25764,0.339,PASS -sgemm_tcu_sp_mx,8,8,mxint8,int32,70,7.9,99,8.0,100,0,19,97,4,0,0,0,0,0,48,20,8,4,19,320,0,0,4.91,15.13,19456,4608,13024,32714,0.398,PASS diff --git a/perf/results/tcu/sgemm_tcu_simx.csv b/perf/results/tcu/sgemm_tcu_simx.csv index 3029e72933..7c7dd6079d 100644 --- a/perf/results/tcu/sgemm_tcu_simx.csv +++ b/perf/results/tcu/sgemm_tcu_simx.csv @@ -10,8 +10,6 @@ sgemm_tcu_sp,8,8,int4,int32,64,7.8,98,8.0,100,0,0,87,4,3,57,0,1,15,46,22,7,5,20, sgemm_tcu_mx,8,8,mxfp8,fp32,61,7.8,98,8.0,100,0,0,94,5,3,47,0,0,28,41,20,6,4,29,320,0,0,4.19,16.13,19968,4608,15264,29156,0.524,PASS sgemm_tcu_mx,8,8,mxfp4,fp32,66,7.6,95,8.0,100,0,0,82,5,3,53,0,1,19,42,21,7,6,23,256,0,0,4.28,19.60,11776,4608,9568,22413,0.427,PASS sgemm_tcu_mx,8,8,nvfp4,fp32,66,7.6,95,8.0,100,0,0,82,5,3,53,0,1,19,42,21,7,6,23,256,0,0,4.28,19.60,11776,4608,9568,22413,0.427,PASS -sgemm_tcu_mx,8,8,mxint8,int32,61,7.8,98,8.0,100,0,0,94,5,3,47,0,0,28,41,20,6,4,29,320,0,0,4.19,16.13,19968,4608,15264,29156,0.524,PASS sgemm_tcu_sp_mx,8,8,mxfp8,fp32,61,7.9,99,8.0,100,0,0,93,4,3,56,0,0,14,48,20,8,4,19,320,0,0,4.20,21.96,16384,4608,13024,28793,0.452,PASS sgemm_tcu_sp_mx,8,8,mxfp4,fp32,66,7.9,99,8.0,100,0,0,90,4,4,65,0,1,10,49,21,9,7,15,256,0,0,4.31,28.30,10240,4608,8736,23192,0.377,PASS sgemm_tcu_sp_mx,8,8,nvfp4,fp32,66,7.9,99,8.0,100,0,0,90,4,4,65,0,1,10,49,21,9,7,15,256,0,0,4.31,28.30,10240,4608,8736,23192,0.377,PASS -sgemm_tcu_sp_mx,8,8,mxint8,int32,61,7.9,99,8.0,100,0,0,93,4,3,56,0,0,14,48,20,8,4,19,320,0,0,4.20,21.96,16384,4608,13024,28793,0.452,PASS diff --git a/sim/rtlsim/Makefile b/sim/rtlsim/Makefile index 97c7daf2e6..616331b261 100644 --- a/sim/rtlsim/Makefile +++ b/sim/rtlsim/Makefile @@ -67,7 +67,16 @@ ifneq (,$(filter -DVX_CFG_EXT_TCU_ENABLE, $(XCONFIGS))) RTL_INCLUDE += -I$(RTL_DIR)/tcu/dsp else ifneq (,$(filter -DVX_CFG_TCU_TYPE_BHF, $(XCONFIGS))) RTL_INCLUDE += -I$(RTL_DIR)/tcu/bhf - RTL_INCLUDE += -J$(THIRD_PARTY_DIR)/hardfloat/source/RISCV + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/fNToRecFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/addRecFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/bsg_hardfloat_pkg.sv + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/HardFloat_primitives.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/HardFloat_rawFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/isSigNaNRecFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/mulRecFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/recFNToFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/recFNToRecFN.v + RTL_INCLUDE += -I$(THIRD_PARTY_DIR)/hardfloat/source/RISCV RTL_INCLUDE += -I$(THIRD_PARTY_DIR)/hardfloat/source else ifneq (,$(filter -DVX_CFG_TCU_TYPE_FPNEW, $(XCONFIGS))) RTL_INCLUDE += -I$(RTL_DIR)/tcu/fpnew @@ -134,6 +143,7 @@ VL_FLAGS += -DVX_CFG_XLEN=$(XLEN) -DVX_CFG_XLEN_$(XLEN) VL_FLAGS += $(CONFIGS) VL_FLAGS += $(RTL_INCLUDE) VL_FLAGS += $(RTL_PKGS) +VL_FLAGS += $(RTL_PRE_SRCS) VL_FLAGS += --cc $(TOP) --top-module $(TOP) # Extract RTL directories from include directories diff --git a/sim/simx/decode.cpp b/sim/simx/decode.cpp index 75dcf1ce29..a9edb8d35a 100644 --- a/sim/simx/decode.cpp +++ b/sim/simx/decode.cpp @@ -899,7 +899,7 @@ Instr::Ptr Decoder::decode(uint32_t code, uint64_t uuid) { uint32_t fmt_d = rd, fmt_s = rs1; bool is_sparse = (rs2 & 1) != 0; instr->set_op_type(is_sparse ? TcuType::WMMA_SP : TcuType::WMMA); - instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, 0, 0, 0, 0, 0}); + instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, 0, 0, 0, 0, 0, 0}); instr->set_macro_op(); instr->set_wstall(true); } break; @@ -910,7 +910,7 @@ Instr::Ptr Decoder::decode(uint32_t code, uint64_t uuid) { uint32_t cd_nregs = (rs2 >> 1) & 0x3; bool is_a_smem = (rs2 >> 3) & 1; instr->set_op_type(is_sparse ? TcuType::WGMMA_SP : TcuType::WGMMA); - instr->set_args(IntrTcuArgs{is_a_smem ? 1u : 0u, cd_nregs, fmt_s, fmt_d, 0, 0, 0, 0, 0}); + instr->set_args(IntrTcuArgs{is_a_smem ? 1u : 0u, cd_nregs, fmt_s, fmt_d, 0, 0, 0, 0, 0, 0}); instr->set_macro_op(); instr->set_wstall(true); } break; @@ -920,7 +920,7 @@ Instr::Ptr Decoder::decode(uint32_t code, uint64_t uuid) { uint32_t fmt_s = rs2; uint32_t slot = rd; instr->set_op_type(TcuType::TCU_LD); - instr->set_args(IntrTcuArgs{0, 0, fmt_s, slot, 0, 0, 0, 0, 0}); + instr->set_args(IntrTcuArgs{0, 0, fmt_s, slot, 0, 0, 0, 0, 0, 0}); // rs1 holds the warp-broadcast base address (real I-reg read). instr->set_src_reg(0, rs1, RegType::Integer); } break; diff --git a/sim/simx/tcu/tcu_unit.cpp b/sim/simx/tcu/tcu_unit.cpp index a64343be90..12e6a9312c 100644 --- a/sim/simx/tcu/tcu_unit.cpp +++ b/sim/simx/tcu/tcu_unit.cpp @@ -32,6 +32,12 @@ using namespace vortex; namespace vt = vortex::tensor; using cfg = vt::wmma_config_t; using wg_cfg = vt::wgmma_config_t; +static constexpr uint32_t kFedpWords = wg_cfg::fedpK; +#ifdef VX_CFG_TCU_FEDP2K +static constexpr bool kFedp2K = true; +#else +static constexpr bool kFedp2K = false; +#endif // Dot-product pipeline depth of the configured tensor-PE type // (multiply / align / accumulate-reduce / round stage sum). @@ -164,11 +170,15 @@ template struct FEDP { using itype = typename It::dtype; static uint32_t eval(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val) { + return eval_n(a_row, b_col, c_val, cfg::tcK); + } + + static uint32_t eval_n(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val, uint32_t k_words) { constexpr uint32_t i_ratio = sizeof(uint32_t) / sizeof(itype); static_assert(i_ratio * sizeof(itype) == sizeof(uint32_t), "FEDP: tcK * i_ratio must be <= 32"); if constexpr (std::is_same_v) { uint32_t acc = 0; - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < k_words; ++z) { auto a = reinterpret_cast(&a_row[z].u32); auto b = reinterpret_cast(&b_col[z].u32); uint32_t prod = 0; @@ -180,7 +190,7 @@ struct FEDP { return rv_fadd_s(c_val, acc, 0, nullptr); } else { uint32_t acc = c_val; - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < k_words; ++z) { auto a = reinterpret_cast(&a_row[z].u32); auto b = reinterpret_cast(&b_col[z].u32); for (uint32_t i = 0; i < i_ratio; ++i) { @@ -195,8 +205,12 @@ struct FEDP { template <> struct FEDP{ static uint32_t eval(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val) { + return eval_n(a_row, b_col, c_val, cfg::tcK); + } + + static uint32_t eval_n(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val, uint32_t k_words) { auto acc = bit_cast(c_val); - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < k_words; ++z) { auto a = a_row[z].u32; auto b = b_col[z].u32; for (uint32_t i = 0; i < 8; ++i) { // 8 * 4 bits = 32 bits @@ -218,8 +232,12 @@ struct FEDP{ template <> struct FEDP{ static uint32_t eval(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val) { + return eval_n(a_row, b_col, c_val, cfg::tcK); + } + + static uint32_t eval_n(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val, uint32_t k_words) { auto acc = bit_cast(c_val); - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < k_words; ++z) { auto a = a_row[z].u32; auto b = b_col[z].u32; for (uint32_t i = 0; i < 8; ++i) { // 8 * 4 bits = 32 bits @@ -232,91 +250,65 @@ struct FEDP{ } }; -using PFN_FEDP = uint32_t (*)(const reg_data_t*, const reg_data_t*, uint32_t); +using PFN_FEDP_N = uint32_t (*)(const reg_data_t*, const reg_data_t*, uint32_t, uint32_t); -static PFN_FEDP select_FEDP(uint32_t IT, uint32_t OT) { +static PFN_FEDP_N select_FEDP_N(uint32_t IT, uint32_t OT) { switch (OT) { case vt::fp32::id: switch (IT) { case vt::fp16::id: - return FEDP::eval; + return FEDP::eval_n; case vt::bf16::id: - return FEDP::eval; + return FEDP::eval_n; case vt::fp8::id: - return FEDP::eval; + return FEDP::eval_n; case vt::bf8::id: - return FEDP::eval; + return FEDP::eval_n; case vt::tf32::id: - return FEDP::eval; + return FEDP::eval_n; default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); + break; } break; case vt::fp16::id: - switch (IT) { - case vt::fp16::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::fp16::id) + return FEDP::eval_n; break; case vt::bf16::id: - switch (IT) { - case vt::bf16::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::bf16::id) + return FEDP::eval_n; break; case vt::fp8::id: - switch (IT) { - case vt::fp8::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::fp8::id) + return FEDP::eval_n; break; case vt::bf8::id: - switch (IT) { - case vt::bf8::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::bf8::id) + return FEDP::eval_n; break; case vt::tf32::id: - switch (IT) { - case vt::tf32::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::tf32::id) + return FEDP::eval_n; break; case vt::int32::id: switch (IT) { case vt::int8::id: - return FEDP::eval; + return FEDP::eval_n; case vt::uint8::id: - return FEDP::eval; + return FEDP::eval_n; case vt::int4::id: - return FEDP::eval; + return FEDP::eval_n; case vt::uint4::id: - return FEDP::eval; + return FEDP::eval_n; default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); + break; } break; default: - std::cout << "Error: unsupported output type: " << OT << "!" << std::endl; - std::abort(); + break; } + std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; + std::abort(); } // Format-agnostic sparse gather: for each bword, iterate over its elem_count packed @@ -371,6 +363,7 @@ class TcuUnit::Impl { exec_done_.fill(false); wgmma_planned_warps_.fill(0); in_wgmma_.fill(false); + wgmma_desc_.fill({0, 0}); } ~Impl() {} @@ -391,6 +384,8 @@ class TcuUnit::Impl { exec_done_.fill(false); wgmma_planned_warps_.fill(0); in_wgmma_.fill(false); + lmem_desc_.clear(); + wgmma_desc_.fill({0, 0}); cta_owner_a_.fill(-1); cta_owner_b_ = -1; cur_block_ = 0; @@ -544,28 +539,21 @@ class TcuUnit::Impl { if (input.empty()) continue; auto trace = input.peek(); if (!tcu_is_wgmma(std::get(trace->op_type))) continue; - wgmma_active |= (1u << b); uint32_t wid = trace->wid; uint64_t wid_bit = (uint64_t(1) << wid); - if (wgmma_planned_warps_.at(b) & wid_bit) continue; + int32_t new_cta = (int32_t)core_->scheduler().warp(wid).cta_csrs.cta_id; auto& instr = *trace->instr_ptr; auto tpuArgs = std::get(instr.get_args()); - if (!(tpuArgs.step_m == 0 && tpuArgs.step_n == 0 && tpuArgs.step_k == 0)) { - // Non-first uop arrived without a prior plan: first uop already drained. - // Mark planned and continue (descriptors persist in lmem_desc_[wid]). - wgmma_planned_warps_.at(b) |= wid_bit; + + if (tpuArgs.is_setup_uop) continue; - } - auto& rs1_data = trace->src_data[0]; - auto& rs2_data = trace->src_data[1]; - uint32_t a_desc = rs1_data.empty() ? 0 : rs1_data.at(0).u32; - uint32_t b_desc = rs2_data.empty() ? 0 : rs2_data.at(0).u32; + + wgmma_active |= (1u << b); // CTA-overlap fence — defer this block's WGMMA if any other block // is mid-flight with a different CTA. The shared B buffer assumes // single-CTA occupancy across all blocks. - int32_t new_cta = (int32_t)core_->scheduler().warp(wid).cta_csrs.cta_id; bool block_other_cta_inflight = false; for (uint32_t k = 0; k < VX_CFG_NUM_TCU_BLOCKS; ++k) { if (k == b) continue; @@ -579,6 +567,25 @@ class TcuUnit::Impl { continue; } + if (wgmma_planned_warps_.at(b) & wid_bit) continue; + if (!(tpuArgs.step_m == 0 && tpuArgs.step_n == 0 && tpuArgs.step_k == 0)) { + // Non-first uop arrived without a prior plan: first uop already drained. + // Mark planned and continue (descriptors persist in lmem_desc_[wid]). + wgmma_planned_warps_.at(b) |= wid_bit; + continue; + } + uint32_t a_desc = wgmma_desc_[wid][0]; + uint32_t b_desc = wgmma_desc_[wid][1]; + bool needs_setup = kFedp2K && !tpuArgs.is_a_smem + && (std::get(trace->op_type) != TcuType::WGMMA_SP); + if (tpuArgs.is_first_uop && !needs_setup) { + if (tpuArgs.is_a_smem) + a_desc = trace->src_data.at(0).at(0).u32; + b_desc = trace->src_data.at(1).at(0).u32; + wgmma_desc_[wid][0] = a_desc; + wgmma_desc_[wid][1] = b_desc; + } + // Drop the shared B buffer only when no other block is mid-WGMMA — // otherwise we'd evict their resident bytes mid-flight. bool any_in_wgmma = false; @@ -629,19 +636,31 @@ class TcuUnit::Impl { continue; auto trace = input.peek(); auto tcu_type = std::get(trace->op_type); + auto tpuArgs = std::get(trace->instr_ptr->get_args()); - #ifdef VX_CFG_TCU_WGMMA_ENABLE + #ifdef VX_CFG_TCU_WGMMA_ENABLE // CTA-overlap fence deferred this block — skip until pass 1 plans it. - if (tcu_is_wgmma(tcu_type) && + if (tcu_is_wgmma(tcu_type) && !tpuArgs.is_setup_uop && !(wgmma_planned_warps_.at(b) & (uint64_t(1) << trace->wid))) continue; + if (tcu_is_wgmma(tcu_type) && !tpuArgs.is_setup_uop) { + int32_t this_cta = (int32_t)core_->scheduler().warp(trace->wid).cta_csrs.cta_id; + bool block_other_cta_inflight = false; + for (uint32_t k = 0; k < VX_CFG_NUM_TCU_BLOCKS; ++k) { + if (k == b) continue; + if (in_wgmma_.at(k) && cta_owner_a_.at(k) != this_cta) { + block_other_cta_inflight = true; + break; + } + } + if (block_other_cta_inflight) + continue; + } #endif // Execute once per trace; results persist across backpressure retries // via exec_done_[b]. if (!exec_done_.at(b)) { - auto& instr = *trace->instr_ptr; - auto tpuArgs = std::get(instr.get_args()); uint32_t wid = trace->wid; uint32_t num_threads = VX_CFG_NUM_THREADS; auto& rs1_data = trace->src_data[0]; @@ -666,7 +685,7 @@ class TcuUnit::Impl { cur_block_ = b; // CTA lockstep invariant: no block may execute a WGMMA uop for a // different cta_id while another block is mid-WGMMA. - { + if (!tpuArgs.is_setup_uop) { int32_t this_cta = (int32_t)core_->scheduler().warp(wid).cta_csrs.cta_id; for (uint32_t k = 0; k < VX_CFG_NUM_TCU_BLOCKS; ++k) { if (k == b) continue; @@ -681,9 +700,9 @@ class TcuUnit::Impl { } this->wgmma(wid, tpuArgs.fmt_s, tpuArgs.fmt_d, tpuArgs.step_m, tpuArgs.step_n, tpuArgs.step_k, - a_desc, b_desc, rs1_data, rs3_data, rd_data, + a_desc, b_desc, rs1_data, rs2_data, rs3_data, rd_data, tcu_is_sparse(tcu_type), - tpuArgs.cd_nregs, tpuArgs.is_a_smem); + tpuArgs.cd_nregs, tpuArgs.is_a_smem, tpuArgs.is_setup_uop); } break; #endif #ifdef TCU_META_ENABLE @@ -768,18 +787,16 @@ class TcuUnit::Impl { uint32_t fmt_s = args.fmt_s; bool is_a_smem = args.is_a_smem; uint32_t e_bits = elem_bits(fmt_s); - if (e_bits < 8) return; - uint32_t e_bytes = e_bits / 8; // NRC: cd_nregs 0/1/2 → 8/16/32; xtileN = NRC * NT / xtileM. uint32_t nrc = (args.cd_nregs == 0) ? 8 : (args.cd_nregs == 1) ? 16 : 32; uint32_t xtile_n = (nrc * VX_CFG_NUM_THREADS) / wg_cfg::xtileM; lmem_desc_t sd_a{}, sd_b{}; if (is_a_smem) { - sd_a = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (a_desc & 0xFFFF), (a_desc >> 16) / e_bytes, false}; + sd_a = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (a_desc & 0xFFFF), (a_desc >> 16) * 8 / e_bits, false}; lmem_desc_[wid][0] = sd_a; } - sd_b = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (b_desc & 0xFFFF), (b_desc >> 16) / e_bytes, false}; + sd_b = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (b_desc & 0xFFFF), (b_desc >> 16) * 8 / e_bits, false}; lmem_desc_[wid][1] = sd_b; // tileK = xtileK × ratio (ratio = 32/e_bits); sparse compresses K on A only. @@ -788,9 +805,11 @@ class TcuUnit::Impl { uint32_t a_k = is_sparse ? (tile_k / 2) : tile_k; // ldm==0 → block-major layout; ldm!=0 → row-major (stride in elements). - uint32_t k_blk_dim = cfg::tcK * ratio; - uint32_t a_blk_elems = cfg::tcM * k_blk_dim; - uint32_t b_blk_elems = k_blk_dim * cfg::tcN; + uint32_t fedp_words = kFedpWords; + uint32_t b_k_blk_dim = fedp_words * ratio; + uint32_t a_k_blk_dim = is_sparse ? (cfg::tcK * ratio) : b_k_blk_dim; + uint32_t a_blk_elems = cfg::tcM * a_k_blk_dim; + uint32_t b_blk_elems = b_k_blk_dim * cfg::tcN; uint32_t n_steps = xtile_n / cfg::tcN; auto& tbuf = simobject_->tbuf(); @@ -806,14 +825,14 @@ class TcuUnit::Impl { if (a_block_major) { uint32_t m_blk = r / cfg::tcM; uint32_t i_in = r % cfg::tcM; - uint32_t k_blk = c / k_blk_dim; - uint32_t k_in = c % k_blk_dim; + uint32_t k_blk = c / a_k_blk_dim; + uint32_t k_in = c % a_k_blk_dim; elem_off = (k_blk * wg_cfg::m_steps + m_blk) * a_blk_elems - + i_in * k_blk_dim + k_in; + + i_in * a_k_blk_dim + k_in; } else { elem_off = uint64_t(r) * sd_a.ldm + c; } - uint64_t addr = sd_a.base + elem_off * e_bytes; + uint64_t addr = sd_a.base + elem_off * e_bits / 8; a_lines.push_back(addr & ~uint64_t(VX_CFG_MEM_BLOCK_SIZE - 1)); } } @@ -831,17 +850,17 @@ class TcuUnit::Impl { for (uint32_t c = 0; c < xtile_n; ++c) { uint64_t elem_off; if (b_block_major) { - uint32_t k_blk = r / k_blk_dim; - uint32_t r_in = r % k_blk_dim; + uint32_t k_blk = r / b_k_blk_dim; + uint32_t r_in = r % b_k_blk_dim; uint32_t n_blk = c / cfg::tcN; uint32_t n_in = c % cfg::tcN; // Within-block layout: N outer, K inner. elem_off = (k_blk * n_steps + n_blk) * b_blk_elems - + n_in * k_blk_dim + r_in; + + n_in * b_k_blk_dim + r_in; } else { elem_off = uint64_t(c) * sd_b.ldm + r; } - uint64_t addr = sd_b.base + elem_off * e_bytes; + uint64_t addr = sd_b.base + elem_off * e_bits / 8; b_lines.push_back(addr & ~uint64_t(VX_CFG_MEM_BLOCK_SIZE - 1)); } } @@ -878,16 +897,15 @@ class TcuUnit::Impl { ? (step_n % cfg::b_sub_blocks_sp) * cfg::b_block_size_sp : (step_n % cfg::b_sub_blocks) * cfg::b_block_size; - // Prepare A tile [tcM][tcK] - reg_data_t a_tile[cfg::tcM * cfg::tcK]; + // WMMA occupies the low half of a widened FEDP and zeros the unused half. + reg_data_t a_tile[cfg::tcM * kFedpWords] = {}; for (uint32_t i = 0; i < cfg::tcM; ++i) { for (uint32_t z = 0; z < cfg::tcK; ++z) { - a_tile[i * cfg::tcK + z] = rs1_data.at(a_off + i * cfg::tcK + z); + a_tile[i * kFedpWords + z] = rs1_data.at(a_off + i * cfg::tcK + z); } } - // Prepare B tile [tcM][tcN][tcK] - reg_data_t b_tile[cfg::tcM * cfg::tcN * cfg::tcK]; + reg_data_t b_tile[cfg::tcM * cfg::tcN * kFedpWords] = {}; if (is_sparse) { constexpr uint32_t kCompression = 2; uint32_t ebits = elem_bits(fmt_s); @@ -907,7 +925,7 @@ class TcuUnit::Impl { lo |= meta_bit(row_base + meta_bits * z + b) << b; hi |= meta_bit(row_base + meta_bits * (cfg::tcK + z) + b) << b; } - b_tile[(i * cfg::tcN + j) * cfg::tcK + z].u32 = + b_tile[(i * cfg::tcN + j) * kFedpWords + z].u32 = gather_sparse(rs2_data.at(b_idx).u32, rs2_data.at(b_idx + 1).u32, lo, hi, ebits); } } @@ -916,14 +934,16 @@ class TcuUnit::Impl { for (uint32_t i = 0; i < cfg::tcM; ++i) { for (uint32_t j = 0; j < cfg::tcN; ++j) { for (uint32_t z = 0; z < cfg::tcK; ++z) { - b_tile[(i * cfg::tcN + j) * cfg::tcK + z] = rs2_data.at(b_off + j * cfg::tcK + z); + b_tile[(i * cfg::tcN + j) * kFedpWords + z] = rs2_data.at(b_off + j * cfg::tcK + z); } } } } fedp_tile(wid, step_m, step_n, step_k, fmt_s, fmt_d, - a_tile, b_tile, rs3_data, rd_data, is_sparse, true); + a_tile, b_tile, rs3_data, rd_data, + is_sparse, kFedpWords, cfg::tcK, + cfg::k_steps * cfg::tcK); } void wgmma(uint32_t wid, @@ -935,12 +955,13 @@ class TcuUnit::Impl { uint32_t a_desc, uint32_t b_desc, const std::vector& rs1_data, + const std::vector& rs2_data, const std::vector& rs3_data, std::vector& rd_data, bool is_sparse, uint32_t cd_nregs, - uint32_t is_a_smem) { - __unused(cd_nregs); + uint32_t is_a_smem, + uint32_t is_setup_uop) { if (is_sparse && !vt::sparse_format_supported(fmt_s)) { std::cout << "Error: WGMMA_SP unsupported input format: " << vt::fmt_string(fmt_s) << " (id=" << fmt_s << ")" << std::endl; @@ -948,46 +969,52 @@ class TcuUnit::Impl { } uint32_t ratio = elem_ratio(fmt_s); - uint32_t k_words = cfg::tcK; - uint32_t e_bytes = elem_bits(fmt_s) / 8; + uint32_t k_words = is_sparse ? cfg::tcK : kFedpWords; + uint32_t e_bits = elem_bits(fmt_s); // Decode smem descriptors (B always from smem, A optionally). - lmem_desc_t sd_a, sd_b; - if (step_k == 0 && step_m == 0 && step_n == 0) { + if (is_setup_uop) { + wgmma_desc_[wid][0] = a_desc; + wgmma_desc_[wid][1] = b_desc; if (is_a_smem) { - sd_a = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (a_desc & 0xFFFF), (a_desc >> 16) / e_bytes, false}; - lmem_desc_[wid][0] = sd_a; + lmem_desc_[wid][0] = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (a_desc & 0xFFFF), (a_desc >> 16) * 8 / e_bits, false}; } - sd_b = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (b_desc & 0xFFFF), (b_desc >> 16) / e_bytes, false}; - lmem_desc_[wid][1] = sd_b; - } else { - sd_a = lmem_desc_[wid][0]; - sd_b = lmem_desc_[wid][1]; + lmem_desc_[wid][1] = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (b_desc & 0xFFFF), (b_desc >> 16) * 8 / e_bits, false}; + uint32_t nrc = (cd_nregs == 0) ? 8 : (cd_nregs == 1) ? 16 : 32; + cur_xtile_n_ = (nrc * VX_CFG_NUM_THREADS) / wg_cfg::xtileM; + return; } + + lmem_desc_t sd_a = lmem_desc_[wid][0]; + lmem_desc_t sd_b = lmem_desc_[wid][1]; // load_lmem_word distinguishes A from B by descriptor base. cur_a_desc_base_ = is_a_smem ? sd_a.base : ~uint64_t(0); + cur_is_sparse_ = is_sparse; // NRC: cd_nregs 0/1/2 → 8/16/32; xtileN = NRC * NT / xtileM. { uint32_t nrc = (cd_nregs == 0) ? 8 : (cd_nregs == 1) ? 16 : 32; cur_xtile_n_ = (nrc * VX_CFG_NUM_THREADS) / wg_cfg::xtileM; } - // Prepare A tile [tcM][tcK] - reg_data_t a_tile[cfg::tcM * cfg::tcK]; + reg_data_t a_tile[cfg::tcM * kFedpWords] = {}; for (uint32_t i = 0; i < cfg::tcM; ++i) { uint32_t a_row_idx = step_m * cfg::tcM + i; for (uint32_t z = 0; z < k_words; ++z) { if (is_a_smem) { uint32_t k_elem = (step_k * k_words + z) * ratio; - a_tile[i * cfg::tcK + z].u32 = load_lmem_word(sd_a, a_row_idx, k_elem, fmt_s, false, false); + a_tile[i * k_words + z].u32 = + load_lmem_word(sd_a, a_row_idx, k_elem, fmt_s, false, false); } else { - a_tile[i * cfg::tcK + z] = rs1_data.at(i * cfg::tcK + z); + if (z < cfg::tcK) { + a_tile[i * k_words + z] = rs1_data.at(i * cfg::tcK + z); + } else { + a_tile[i * k_words + z] = rs2_data.at(i * cfg::tcK + (z - cfg::tcK)); + } } } } - // Prepare B tile [tcM][tcN][tcK] - reg_data_t b_tile[cfg::tcM * cfg::tcN * cfg::tcK]; + reg_data_t b_tile[cfg::tcM * cfg::tcN * kFedpWords] = {}; if (is_sparse) { uint32_t ebits = elem_bits(fmt_s); uint32_t rtl_i_ratio = 32 / ebits; @@ -1022,12 +1049,12 @@ class TcuUnit::Impl { uint32_t bword0 = load_lmem_word(sd_b, k_elem_b0, b_col_idx, fmt_s, true, true); uint32_t bword1 = load_lmem_word(sd_b, k_elem_b0 + ratio, b_col_idx, fmt_s, true, true); uint32_t gathered = gather_sparse(bword0, bword1, lo, hi, ebits); - b_tile[(i * cfg::tcN + j) * cfg::tcK + z].u32 = gathered; + b_tile[(i * cfg::tcN + j) * k_words + z].u32 = gathered; // Trace: B-gather (CSV: wid,step_m,step_n,i,lane,bword0,bword1,lo,hi,gathered). if (const char* p = std::getenv("VORTEX_TCU_TRACE")) { if (p[0] == '1') { fprintf(stderr, "GATHER,%u,%u,%u,%u,%u,0x%08x,0x%08x,%u,%u,0x%08x\n", - wid, step_m, step_n, i, j*cfg::tcK+z, + wid, step_m, step_n, i, j*k_words+z, bword0, bword1, lo, hi, gathered); } } @@ -1040,7 +1067,7 @@ class TcuUnit::Impl { uint32_t b_col_idx = step_n * cfg::tcN + j; for (uint32_t z = 0; z < k_words; ++z) { uint32_t k_elem = (step_k * k_words + z) * ratio; - b_tile[(i * cfg::tcN + j) * cfg::tcK + z].u32 = + b_tile[(i * cfg::tcN + j) * k_words + z].u32 = load_lmem_word(sd_b, k_elem, b_col_idx, fmt_s, true, false); } } @@ -1048,7 +1075,9 @@ class TcuUnit::Impl { } fedp_tile(wid, step_m, step_n, step_k, fmt_s, fmt_d, - a_tile, b_tile, rs3_data, rd_data, is_sparse, false); + a_tile, b_tile, rs3_data, rd_data, + is_sparse, k_words, k_words, + wg_cfg::k_steps * wg_cfg::fedpK); __unused(b_desc); } @@ -1074,7 +1103,6 @@ class TcuUnit::Impl { case vt::uint8::id: case vt::mxfp8::id: case vt::mxbf8::id: - case vt::mxint8::id: return 8; case vt::int4::id: case vt::uint4::id: @@ -1090,35 +1118,28 @@ class TcuUnit::Impl { return 32 / elem_bits(fmt_s); } - uint32_t mx_tile_scale_blocks(uint32_t fmt_s) const { - uint32_t logical_tile_k = cfg::k_steps * cfg::tcK * elem_ratio(fmt_s); + uint32_t mx_tile_scale_blocks(uint32_t fmt_s, uint32_t tile_k_words) const { + uint32_t logical_tile_k = tile_k_words * elem_ratio(fmt_s); uint32_t block_size = vt::mx_scale_block_size(fmt_s); - return std::max(1u, logical_tile_k / block_size); + return std::max(1u, (logical_tile_k + block_size - 1) / block_size); } static uint8_t meta_byte(const std::vector& words, uint32_t index) { return (words.at(index / 4) >> (8 * (index & 0x3))) & 0xff; } - static int32_t trunc_shift(int32_t value, int32_t shift) { - if (shift >= 0) - return value << shift; - uint32_t amount = static_cast(-shift); - uint32_t magnitude = value < 0 ? static_cast(-value) : static_cast(value); - int32_t scaled = static_cast(magnitude >> amount); - return value < 0 ? -scaled : scaled; - } - uint32_t eval_mx_fedp(uint32_t wid, uint32_t fmt_s, uint32_t fmt_d, uint32_t step_m, uint32_t step_n, uint32_t step_k, uint32_t i, uint32_t j, const reg_data_t* a_row, - const reg_data_t* b_col, uint32_t c_val, bool is_sparse) const { + const reg_data_t* b_col, uint32_t c_val, bool is_sparse, + uint32_t fedp_words, uint32_t tile_k_words) const { #ifndef VX_CFG_TCU_MX_ENABLE - __unused(wid, fmt_s, fmt_d, step_m, step_n, step_k, i, j, a_row, b_col, is_sparse); + __unused(wid, fmt_s, fmt_d, step_m, step_n, step_k, i, j, a_row, b_col, + is_sparse, fedp_words, tile_k_words); return c_val; #else uint32_t ratio = elem_ratio(fmt_s); - uint32_t scale_blocks_k = mx_tile_scale_blocks(fmt_s); + uint32_t scale_blocks_k = mx_tile_scale_blocks(fmt_s, tile_k_words); uint32_t block_size = vt::mx_scale_block_size(fmt_s); auto scale_a = [&](uint32_t elem_k) { uint32_t row = step_m * cfg::tcM + i; @@ -1131,10 +1152,10 @@ class TcuUnit::Impl { if (fmt_d == vt::fp32::id) { uint32_t acc = c_val; - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < fedp_words; ++z) { for (uint32_t e = 0; e < ratio; ++e) { uint32_t sparse_ratio = is_sparse ? 2 : 1; - uint32_t elem_k = ((step_k * cfg::tcK + z) * ratio + e) * sparse_ratio; + uint32_t elem_k = ((step_k * fedp_words + z) * ratio + e) * sparse_ratio; uint32_t xa, xb; if (fmt_s == vt::mxfp8::id) { xa = rv_mxfp8tof_s((a_row[z].u32 >> (8 * e)) & 0xff, scale_a(elem_k), 0, nullptr); @@ -1157,20 +1178,6 @@ class TcuUnit::Impl { return acc; } - if (fmt_s == vt::mxint8::id && fmt_d == vt::int32::id) { - int32_t acc = bit_cast(c_val); - for (uint32_t z = 0; z < cfg::tcK; ++z) { - for (uint32_t e = 0; e < ratio; ++e) { - uint32_t elem_k = ((step_k * cfg::tcK + z) * ratio + e) * (is_sparse ? 2 : 1); - auto a = static_cast((a_row[z].u32 >> (8 * e)) & 0xff); - auto b = static_cast((b_col[z].u32 >> (8 * e)) & 0xff); - int32_t shift = int32_t(scale_a(elem_k)) + int32_t(scale_b(elem_k)) - 266; - acc += trunc_shift(int32_t(a) * int32_t(b), shift); - } - } - return bit_cast(acc); - } - std::cout << "Error: unsupported MX mma format: " << fmt_s << " -> " << fmt_d << "!" << std::endl; std::abort(); #endif @@ -1184,19 +1191,21 @@ class TcuUnit::Impl { template uint32_t gather_word(ReadLine read_line, const lmem_desc_t& desc, uint32_t row, uint32_t col, - uint32_t fmt_s, bool pack_along_row, bool sparse_b) const { + uint32_t fmt_s, bool pack_along_row, bool sparse_b, + bool sparse_a_layout) const { uint32_t e_bits = elem_bits(fmt_s); uint32_t ratio = (e_bits >= 32) ? 1 : (32 / e_bits); - uint32_t e_bytes = (e_bits >= 8) ? (e_bits / 8) : 1; uint32_t result = 0; for (uint32_t r = 0; r < ratio; ++r) { uint32_t cur_row = pack_along_row ? (row + r) : row; uint32_t cur_col = pack_along_row ? col : (col + r); - uint64_t byte_addr; + uint64_t elem_off; if (desc.ldm == 0) { // Block-major SMEM. K dimension is along col for A (pack_along_row // false) and along row for B (pack_along_row true). - uint32_t k_blk_dim = cfg::tcK * ratio; + uint32_t k_blk_dim = (sparse_a_layout && !pack_along_row) + ? (cfg::tcK * ratio) + : (kFedpWords * ratio); if (pack_along_row && sparse_b) { // Sparse B in flat (candidate-pair) layout: block-contiguous // K-word-major / N-inner order [kw_in*tcN + n_in] (mirrors @@ -1213,8 +1222,7 @@ class TcuUnit::Impl { uint32_t n_steps = cur_xtile_n_ / cfg::tcN; uint64_t word_off = (k_blk * n_steps + n_blk) * blk_words + (kw_in * cfg::tcN + n_in); - uint64_t off = word_off * ratio + elem; - byte_addr = desc.base + off * e_bytes; + elem_off = word_off * ratio + elem; } else if (pack_along_row) { // Dense B (block-major): r is K coord, c is N coord; N outer, K inner. uint32_t k_blk = cur_row / k_blk_dim; @@ -1223,9 +1231,8 @@ class TcuUnit::Impl { uint32_t n_in = cur_col % cfg::tcN; uint32_t b_blk_elems = k_blk_dim * cfg::tcN; uint32_t n_steps = cur_xtile_n_ / cfg::tcN; - uint64_t off = (k_blk * n_steps + n_blk) * b_blk_elems - + n_in * k_blk_dim + r_in; - byte_addr = desc.base + off * e_bytes; + elem_off = (k_blk * n_steps + n_blk) * b_blk_elems + + n_in * k_blk_dim + r_in; } else { // A: r is M coord, c is K coord. uint32_t m_blk = cur_row / cfg::tcM; @@ -1233,20 +1240,20 @@ class TcuUnit::Impl { uint32_t k_blk = cur_col / k_blk_dim; uint32_t k_in = cur_col % k_blk_dim; uint32_t a_blk_elems = cfg::tcM * k_blk_dim; - uint64_t off = (k_blk * wg_cfg::m_steps + m_blk) * a_blk_elems - + i_in * k_blk_dim + k_in; - byte_addr = desc.base + off * e_bytes; + elem_off = (k_blk * wg_cfg::m_steps + m_blk) * a_blk_elems + + i_in * k_blk_dim + k_in; } } else if (desc.col_major) { - byte_addr = desc.base + (uint64_t(cur_col) * desc.ldm + cur_row) * e_bytes; + elem_off = uint64_t(cur_col) * desc.ldm + cur_row; } else if (pack_along_row) { // B: K-major (N-outer K-inner). cur_row=K, cur_col=N; // ldm = stride in elements between N rows. - byte_addr = desc.base + (uint64_t(cur_col) * desc.ldm + cur_row) * e_bytes; + elem_off = uint64_t(cur_col) * desc.ldm + cur_row; } else { // A: row-major (M-outer K-inner). cur_row=M, cur_col=K. - byte_addr = desc.base + (uint64_t(cur_row) * desc.ldm + cur_col) * e_bytes; + elem_off = uint64_t(cur_row) * desc.ldm + cur_col; } + uint64_t byte_addr = desc.base + elem_off * e_bits / 8; auto line = read_line(byte_addr); if (!line) { std::cout << "Error: TCU buffer miss at 0x" << std::hex << byte_addr @@ -1265,9 +1272,10 @@ class TcuUnit::Impl { result |= (val & 0xFFFF) << (r * 16); } else if (e_bits == 8) { result |= uint32_t((*line)[off]) << (r * 8); + } else if (e_bits == 4) { + uint32_t val = (uint32_t((*line)[off]) >> (4 * (elem_off & 1))) & 0xf; + result |= val << (r * 4); } else { - // 4-bit (int4/uint4) not supported. - std::cout << "Error: TCU 4-bit operand gather not supported" << std::endl; std::abort(); } } @@ -1275,17 +1283,17 @@ class TcuUnit::Impl { } // Routes A reads through the current block's A buffer; B through the shared B buffer. - // sparse_b selects the flat candidate-pair B layout (ignored for A reads). + // sparse_b selects the flat candidate-pair B layout. uint32_t load_lmem_word(const lmem_desc_t& desc, uint32_t row, uint32_t col, uint32_t fmt_s, bool pack_along_row, bool sparse_b) const { auto& tbuf = simobject_->tbuf(); if (desc.base == cur_a_desc_base_) { uint32_t b = cur_block_; return gather_word([&](uint64_t addr) { return tbuf->read_a(b, addr); }, - desc, row, col, fmt_s, pack_along_row, sparse_b); + desc, row, col, fmt_s, pack_along_row, false, cur_is_sparse_); } return gather_word([&](uint64_t addr) { return tbuf->read_b(addr); }, - desc, row, col, fmt_s, pack_along_row, sparse_b); + desc, row, col, fmt_s, pack_along_row, sparse_b, false); } static constexpr uint32_t kSparseKSteps = cfg::k_steps / 2; @@ -1293,8 +1301,8 @@ class TcuUnit::Impl { static constexpr uint32_t kMaxMetaCols = VX_CFG_NUM_THREADS / 2; // FEDP tile computation for both WMMA and WGMMA. - // a_tile: [tcM][tcK] flat array of pre-loaded A operand words. - // b_tile: [tcM][tcN][tcK] flat array of pre-loaded B operand words + // a_tile: [tcM][k_words] flat array of pre-loaded A operand words. + // b_tile: [tcM][tcN][k_words] flat array of pre-loaded B operand words // (for dense, each i-slice is identical; for sparse, already gathered). void fedp_tile(uint32_t wid, uint32_t step_m, uint32_t step_n, uint32_t step_k, @@ -1304,23 +1312,23 @@ class TcuUnit::Impl { const std::vector& rs3_data, std::vector& rd_data, bool is_sparse, - bool allow_mx) { - if (!allow_mx && vt::mx_scale_format(fmt_s)) { - std::cout << "Error: MX formats are supported only by WMMA." << std::endl; - std::abort(); - } - PFN_FEDP fedp = vt::mx_scale_format(fmt_s) ? nullptr : select_FEDP(fmt_s, fmt_d); + uint32_t k_words = cfg::tcK, + uint32_t fedp_words = cfg::tcK, + uint32_t tile_k_words = cfg::k_steps * cfg::tcK) { + PFN_FEDP_N fedp = vt::mx_scale_format(fmt_s) ? nullptr : select_FEDP_N(fmt_s, fmt_d); for (uint32_t i = 0; i < cfg::tcM; ++i) { for (uint32_t j = 0; j < cfg::tcN; ++j) { auto t = i * cfg::tcN + j; auto c_val = rs3_data.at(t).u32; - auto a_row = &a_tile[i * cfg::tcK]; - auto b_col = &b_tile[(i * cfg::tcN + j) * cfg::tcK]; + auto a_row = &a_tile[i * k_words]; + auto b_col = &b_tile[(i * cfg::tcN + j) * k_words]; - uint32_t d_val = (allow_mx && vt::mx_scale_format(fmt_s)) - ? eval_mx_fedp(wid, fmt_s, fmt_d, step_m, step_n, step_k, i, j, a_row, b_col, c_val, is_sparse) - : fedp(a_row, b_col, c_val); + uint32_t d_val = vt::mx_scale_format(fmt_s) + ? eval_mx_fedp(wid, fmt_s, fmt_d, step_m, step_n, step_k, + i, j, a_row, b_col, c_val, is_sparse, + fedp_words, tile_k_words) + : fedp(a_row, b_col, c_val, k_words); rd_data.at(t).u64 = nan_box(d_val); DTH(3, simobject_->name() << " FEDP" @@ -1341,6 +1349,7 @@ class TcuUnit::Impl { std::vector> mx_meta_b_; #endif std::unordered_map lmem_desc_; + std::array, VX_CFG_NUM_WARPS> wgmma_desc_; mutable PerfStats perf_stats_; // Per-block guard: execute already happened for this trace; reset on pop(). std::array exec_done_; @@ -1352,6 +1361,7 @@ class TcuUnit::Impl { uint32_t cur_block_ = 0; // A-descriptor base for the current wgmma(); distinguishes A from B in load_lmem_word. uint64_t cur_a_desc_base_ = ~uint64_t(0); + bool cur_is_sparse_ = false; // xtileN for the active WGMMA (derived from NRC). uint32_t cur_xtile_n_ = 8; // CTA owner per block's A buffer and the shared B buffer (-1 = unowned). @@ -1410,10 +1420,11 @@ uint32_t TcuUopGen::uop_count(const Instr& instr) { #ifdef VX_CFG_TCU_WGMMA_ENABLE if (tcu_is_wgmma(tcu_type)) { bool is_sparse = tcu_is_sparse(tcu_type); + bool needs_setup = kFedp2K && !args.is_a_smem && !is_sparse; uint32_t nrc = (args.cd_nregs == 0) ? 8 : (args.cd_nregs == 1) ? 16 : 32; - uint32_t k_count = is_sparse ? (wg_cfg::k_steps / 2) : wg_cfg::k_steps; + uint32_t k_count = is_sparse ? std::max(1u, wg_cfg::k_steps / 2) : wg_cfg::k_steps; uint32_t mma_uops = k_count * nrc; - return mma_uops; + return mma_uops + needs_setup; } #endif @@ -1471,7 +1482,7 @@ Instr::Ptr TcuUopGen::get(const Instr& macro_instr, uint32_t uop_index) { uint32_t actual_n = lg_k ? (n_sp >> lg_k) : n_sp; uint32_t reg_rs3 = rc_base + (mma_idx >> 1); uop_instr->set_op_type(TcuType::WMMA_SP); - uop_instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, m_sp, actual_n, 0, 0, 0}); + uop_instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, m_sp, actual_n, 0, 0, 0, 0}); uop_instr->set_dest_reg(reg_rs3, RegType::Float); uop_instr->set_src_reg(0, ra_base + m_sp, RegType::Float); uop_instr->set_src_reg(1, rb_base + n_sp, RegType::Float); @@ -1527,7 +1538,7 @@ Instr::Ptr TcuUopGen::get(const Instr& macro_instr, uint32_t uop_index) { reg_rs3 = rc_base + c_off; } uop_instr->set_op_type(is_sparse ? TcuType::WMMA_SP : TcuType::WMMA); - uop_instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, m, n, k, 0, 0}); + uop_instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, m, n, k, 0, 0, 0}); uop_instr->set_dest_reg(reg_rs3, RegType::Float); uop_instr->set_src_reg(0, reg_rs1, RegType::Float); uop_instr->set_src_reg(1, reg_rs2, RegType::Float); @@ -1539,48 +1550,60 @@ Instr::Ptr TcuUopGen::get(const Instr& macro_instr, uint32_t uop_index) { else if (tcu_is_wgmma(tcu_type)) { constexpr uint32_t m_steps = wg_cfg::m_steps; constexpr uint32_t k_steps = wg_cfg::k_steps; + constexpr uint32_t a_reg_k_steps = 2; uint32_t fmt_s = args.fmt_s; uint32_t fmt_d = args.fmt_d; bool is_sparse = tcu_is_sparse(tcu_type); bool is_a_smem = args.is_a_smem; + bool needs_setup = kFedp2K && !is_a_smem && !is_sparse; uint32_t cd_nregs = args.cd_nregs; - uint32_t k_count = is_sparse ? (k_steps / 2) : k_steps; + uint32_t k_count = is_sparse ? std::max(1u, k_steps / 2) : k_steps; constexpr uint32_t a0 = 10, a1 = 11; { + uop_instr->set_op_type(is_sparse ? TcuType::WGMMA_SP : TcuType::WGMMA); + + if (needs_setup && uop_index == 0) { + uop_instr->set_args(IntrTcuArgs{is_a_smem ? 1u : 0u, cd_nregs, + fmt_s, fmt_d, 0, 0, 0, 1, 0, 1}); + uop_instr->set_src_reg(0, a0, RegType::Integer); + uop_instr->set_src_reg(1, a1, RegType::Integer); + uop_instr->set_fu_lock(true); + uop_instr->set_fu_unlock(false); + return uop_instr; + } + // MMA phase - uint32_t mma_idx = uop_index; + uint32_t mma_idx = uop_index - needs_setup; uint32_t ra_base = is_a_smem ? 10 : 24; // Loop order: m (inner) -> n (middle) -> k (outer). K-outer maximizes // per-block A-buffer reuse: each A_w[m,k] is consumed across the entire // (n,m) inner sweep, and each shared B[k,n] is consumed for m_steps // consecutive uops. - uint32_t mn = total / k_count; + uint32_t compute_total = total - needs_setup; + uint32_t mn = compute_total / k_count; uint32_t k = mma_idx / mn; uint32_t rem = mma_idx % mn; uint32_t n = rem / m_steps; uint32_t m = rem % m_steps; uint32_t r = n * m_steps + m; - uop_instr->set_op_type(is_sparse ? TcuType::WGMMA_SP : TcuType::WGMMA); - bool first = (uop_index == 0); + bool first = (mma_idx == 0); bool last = (uop_index == (total - 1)); uop_instr->set_args(IntrTcuArgs{is_a_smem ? 1u : 0u, cd_nregs, - fmt_s, fmt_d, m, n, k, first ? 1u : 0u, last ? 1u : 0u}); + fmt_s, fmt_d, m, n, k, first ? 1u : 0u, last ? 1u : 0u, 0}); uop_instr->set_dest_reg(r, RegType::Float); - if (mma_idx == 0) { - if (is_a_smem) { - uop_instr->set_src_reg(0, a0, RegType::Integer); - } else { - uint32_t rs1_off = m * k_count + k; - uop_instr->set_src_reg(0, ra_base + rs1_off, RegType::Float); - } - uop_instr->set_src_reg(1, a1, RegType::Integer); - } else if (!is_a_smem) { - uint32_t rs1_off = m * k_count + k; + if (!is_a_smem) { + uint32_t rs1_off = is_sparse ? m : (m * a_reg_k_steps + k); uop_instr->set_src_reg(0, ra_base + rs1_off, RegType::Float); + if (needs_setup) + uop_instr->set_src_reg(1, ra_base + rs1_off + 1, RegType::Float); + } else if (first) { + uop_instr->set_src_reg(0, a0, RegType::Integer); } + if (first && !needs_setup) + uop_instr->set_src_reg(1, a1, RegType::Integer); uop_instr->set_src_reg(2, r, RegType::Float); } // fu_lock on first uop, fu_unlock on last uop @@ -1645,20 +1668,22 @@ void TcuUnit::wmma(uint32_t wid, } void TcuUnit::wgmma(uint32_t wid, - uint32_t fmt_s, - uint32_t fmt_d, - uint32_t step_m, - uint32_t step_n, - uint32_t step_k, - uint32_t a_desc, - uint32_t b_desc, - const std::vector& rs1_data, - const std::vector& rs3_data, - std::vector& rd_data, - bool is_sparse, - uint32_t cd_nregs, - uint32_t is_a_smem) { + uint32_t fmt_s, + uint32_t fmt_d, + uint32_t step_m, + uint32_t step_n, + uint32_t step_k, + uint32_t a_desc, + uint32_t b_desc, + const std::vector& rs1_data, + const std::vector& rs2_data, + const std::vector& rs3_data, + std::vector& rd_data, + bool is_sparse, + uint32_t cd_nregs, + uint32_t is_a_smem, + uint32_t is_setup_uop) { impl_->wgmma(wid, fmt_s, fmt_d, step_m, step_n, step_k, a_desc, b_desc, - rs1_data, rs3_data, rd_data, - is_sparse, cd_nregs, is_a_smem); + rs1_data, rs2_data, rs3_data, rd_data, + is_sparse, cd_nregs, is_a_smem, is_setup_uop); } diff --git a/sim/simx/tcu/tcu_unit.h b/sim/simx/tcu/tcu_unit.h index 88ea84aa75..5eca3c70fc 100644 --- a/sim/simx/tcu/tcu_unit.h +++ b/sim/simx/tcu/tcu_unit.h @@ -105,11 +105,13 @@ class TcuUnit : public FuncUnit { uint32_t a_desc, uint32_t b_desc, const std::vector& rs1_data, + const std::vector& rs2_data, const std::vector& rs3_data, std::vector& rd_data, bool is_sparse, uint32_t cd_nregs, - uint32_t is_a_smem); + uint32_t is_a_smem, + uint32_t is_setup_uop); // Tile-buffer subsystem (owns abuf×Q + bbuf + LMEM arb). // Exposed so that `Core` can bind its single LMEM port pair. diff --git a/sim/simx/types.h b/sim/simx/types.h index 5587a673ee..9b3253ce31 100644 --- a/sim/simx/types.h +++ b/sim/simx/types.h @@ -759,6 +759,7 @@ struct IntrTcuArgs { uint32_t step_k : 4; uint32_t is_first_uop : 1; // set per-uop by tcu_uops expansion (C4) uint32_t is_last_uop : 1; + uint32_t is_setup_uop : 1; // WGMMA descriptor setup uop, no FEDP compute }; // Helper: is_sparse derived from op_type (no per-uop bit). diff --git a/sw/common/rvfloats.cpp b/sw/common/rvfloats.cpp index e10f0f3d47..0648bd631a 100644 --- a/sw/common/rvfloats.cpp +++ b/sw/common/rvfloats.cpp @@ -12,8 +12,6 @@ // limitations under the License. #include "rvfloats.h" -#include -#include #include #include @@ -42,19 +40,6 @@ inline void rv_init(uint32_t frm) { softfloat_roundingMode = frm; } -static inline int32_t round_ties_to_even_i32(float value) { - float floor_v = std::floor(value); - float frac = value - floor_v; - int32_t base = static_cast(floor_v); - if (frac < 0.5f) { - return base; - } - if (frac > 0.5f) { - return base + 1; - } - return (base & 1) ? (base + 1) : base; -} - #ifdef __cplusplus extern "C" { #endif @@ -669,23 +654,6 @@ uint8_t rv_ftomxbf8_s(uint32_t a, uint8_t sf, uint32_t frm, uint32_t* fflags) { return mxbf8.v; } -uint8_t rv_ftomxint8_s(uint32_t a, uint8_t sf, uint32_t frm, uint32_t* fflags) { - rv_init(frm); - union { - uint32_t u; - float f; - } in{a}; - - int32_t scale_exp = static_cast(sf) - 127; - float inv_scale = std::ldexp(1.0f, -scale_exp); - float q_real = in.f * inv_scale * 64.0f; - int32_t q = round_ties_to_even_i32(q_real); - q = std::max(-127, std::min(127, q)); - - if (fflags) { *fflags = softfloat_exceptionFlags; } - return static_cast(static_cast(q)); -} - uint32_t rv_mxfp4tof_s(uint8_t a, uint8_t sf, uint32_t frm, uint32_t* fflags) { rv_init(frm); mxfloat4_t mxfp4 = {a, sf}; diff --git a/sw/common/rvfloats.h b/sw/common/rvfloats.h index a2ad935602..57d03079a5 100644 --- a/sw/common/rvfloats.h +++ b/sw/common/rvfloats.h @@ -127,9 +127,6 @@ uint8_t rv_ftomxfp4_s(uint32_t a, uint8_t sf, uint32_t frm, uint32_t* fflags); uint32_t rv_nvfp4tof_s(uint8_t a, uint8_t sf, uint32_t frm, uint32_t* fflags); uint8_t rv_ftonvfp4_s(uint32_t a, uint8_t sf, uint32_t frm, uint32_t* fflags); -// mxint8 conversion (to quantized INT8 representation) -uint8_t rv_ftomxint8_s(uint32_t a, uint8_t sf, uint32_t frm, uint32_t* fflags); - // e2m1 <--> fp32 conversions uint32_t rv_e2m1tof_s(uint8_t a, uint32_t frm, uint32_t* fflags); uint8_t rv_ftoe2m1_s(uint32_t a, uint32_t frm, uint32_t* fflags); diff --git a/sw/common/tensor_cfg.h b/sw/common/tensor_cfg.h index 18e5116fac..8ab90ed471 100644 --- a/sw/common/tensor_cfg.h +++ b/sw/common/tensor_cfg.h @@ -135,15 +135,6 @@ struct uint4 { static constexpr const char* name = "u4"; }; -struct mxint8 { - using dtype = int8_t; - static constexpr uint32_t id = 24; - static constexpr uint32_t bits = 8; - static constexpr uint32_t scale_bits = 8; - static constexpr uint32_t ele_block = 32; - static constexpr const char* name = "mxint8"; -}; - inline const char* fmt_string(uint32_t fmt) { switch (fmt) { case fp32::id: return fp32::name; @@ -161,7 +152,6 @@ inline const char* fmt_string(uint32_t fmt) { case uint8::id: return uint8::name; case int4::id: return int4::name; case uint4::id: return uint4::name; - case mxint8::id: return mxint8::name; default: return ""; } } @@ -172,7 +162,6 @@ inline constexpr bool mx_scale_format(uint32_t fmt) { case mxbf8::id: case mxfp4::id: case nvfp4::id: - case mxint8::id: return true; default: return false; @@ -184,7 +173,6 @@ inline constexpr uint32_t mx_scale_block_size(uint32_t fmt) { case mxfp8::id: case mxbf8::id: case mxfp4::id: - case mxint8::id: return 32; case nvfp4::id: return 16; @@ -204,7 +192,6 @@ inline constexpr bool sparse_format_supported(uint32_t fmt) { case uint8::id: case mxfp8::id: case mxbf8::id: - case mxint8::id: case int4::id: case uint4::id: case mxfp4::id: @@ -228,7 +215,6 @@ inline constexpr uint32_t sparse_meta_num_cols(uint32_t fmt, uint32_t nt) { case uint8::id: case mxfp8::id: case mxbf8::id: - case mxint8::id: return (nt + 3) / 4; case int4::id: case uint4::id: @@ -345,7 +331,7 @@ struct wmma_config_t { // All geometry derived from NT and NRC alone (NRA=4 fixed): // tcM = 2^ceil(log2(NT)/2), tcN = tcK = 2^floor(log2(NT)/2) // xtileM = 2*tcM, xtileN = NRC*NT/xtileM, xtileK = 2*tcK -// m_steps = k_steps = 2 (always) +// m_steps = 2; k_steps is 2 normally, or 1 with a doubled-K WGMMA FEDP. template struct wgmma_config_t { private: @@ -360,12 +346,18 @@ struct wgmma_config_t { static constexpr uint32_t tcM = 1u << ((lg_NT + 1) / 2); static constexpr uint32_t tcN = 1u << (lg_NT / 2); static constexpr uint32_t tcK = tcN; + static constexpr uint32_t fedpK = +#ifdef VX_CFG_TCU_FEDP2K + 2 * tcK; +#else + tcK; +#endif static constexpr uint32_t xtileM = 2 * tcM; static constexpr uint32_t xtileN = (NRC_ * NT) / xtileM; static constexpr uint32_t xtileK = 2 * tcK; static constexpr uint32_t tileK = xtileK * i_ratio; static constexpr uint32_t m_steps = 2; - static constexpr uint32_t k_steps = 2; + static constexpr uint32_t k_steps = xtileK / fedpK; static constexpr uint32_t NRC = NRC_; }; diff --git a/sw/kernel/include/vx_tensor.h b/sw/kernel/include/vx_tensor.h index 0662d14bb7..c132fd751d 100644 --- a/sw/kernel/include/vx_tensor.h +++ b/sw/kernel/include/vx_tensor.h @@ -715,6 +715,9 @@ struct wgmma_context { // Fragments using fragment_acc = typename ctx_c::fragment_acc; using fragment_a = typename ctx_a::fragment_a; + using fragment_b = typename ctx_c::fragment_b; + + static constexpr bool input_is_mx = mx_scale_format(It::id); // Block (micro-tile) geometry — derived from NT alone static constexpr uint32_t i_ratio = XB / sizeof(typename It::dtype); @@ -722,19 +725,27 @@ struct wgmma_context { static constexpr uint32_t tcN = 1u << (lg_NT / 2); static constexpr uint32_t tcK = tcN; - // Per-warp tile geometry — m_steps = k_steps = 2 always + // Per-warp tile geometry. WGMMA-2K keeps the logical tile K fixed and + // consumes both K halves inside one widened FEDP uop. + static constexpr uint32_t fedpK = +#ifdef VX_CFG_TCU_FEDP2K + 2 * tcK; +#else + tcK; +#endif static constexpr uint32_t m_steps = 2; - static constexpr uint32_t k_steps = 2; + static constexpr uint32_t k_steps = (2 * tcK) / fedpK; + static constexpr uint32_t a_reg_k_steps = 2; static constexpr uint32_t xtileM = m_steps * tcM; static constexpr uint32_t xtileN = (NRC_ * NT) / xtileM; - static constexpr uint32_t tileK = k_steps * tcK * i_ratio; + static constexpr uint32_t tileK = k_steps * fedpK * i_ratio; static constexpr uint32_t n_steps = xtileN / tcN; // Block-major SMEM constants. // BLOCK = micro-tile (tcM × tcK or tcK × tcN), measured in input_t element units. - static constexpr uint32_t a_blk_elems = tcM * tcK * i_ratio; // elements per A block + static constexpr uint32_t a_blk_elems = tcM * fedpK * i_ratio; // elements per A block static constexpr uint32_t a_warp_elems = xtileM * tileK; // elements per warp's A slice - static constexpr uint32_t b_blk_elems = tcK * i_ratio * tcN; // elements per B block + static constexpr uint32_t b_blk_elems = fedpK * i_ratio * tcN; // elements per B block // Cooperative-load index into A_smem for an (r, c) target in the // row-major-equivalent A view (r ∈ [0, cta_M), c ∈ [0, tileK)). @@ -744,11 +755,11 @@ struct wgmma_context { uint32_t r_in_warp = r % xtileM; uint32_t m_blk = r_in_warp / tcM; uint32_t i_in = r_in_warp % tcM; - uint32_t k_blk = c / (tcK * i_ratio); - uint32_t k_in = c % (tcK * i_ratio); + uint32_t k_blk = c / (fedpK * i_ratio); + uint32_t k_in = c % (fedpK * i_ratio); return warp_idx * a_warp_elems + (k_blk * m_steps + m_blk) * a_blk_elems - + i_in * (tcK * i_ratio) + k_in; + + i_in * (fedpK * i_ratio) + k_in; } // Cooperative-load index into B_smem for an (r, c) target in the @@ -756,12 +767,12 @@ struct wgmma_context { // Within-block layout: N outer, K inner — each 32-bit word packs i_ratio // K-elements at one (j, k_word) cell, matching tcu_core's b_off + j*TC_K + k. static __attribute__((always_inline)) uint32_t b_blockmajor_idx(uint32_t r, uint32_t c) { - uint32_t k_blk = r / (tcK * i_ratio); - uint32_t r_in = r % (tcK * i_ratio); + uint32_t k_blk = r / (fedpK * i_ratio); + uint32_t r_in = r % (fedpK * i_ratio); uint32_t n_blk = c / tcN; uint32_t n_in = c % tcN; return (k_blk * n_steps + n_blk) * b_blk_elems - + n_in * (tcK * i_ratio) + r_in; + + n_in * (fedpK * i_ratio) + r_in; } // Flat sparse-B SMEM layout: per (k_blk, n_blk) block, words are stored @@ -797,11 +808,12 @@ struct wgmma_context { // Sparse A is K/2 compressed; same per-block shape as dense A but only // K_STEPS/2 half-k blocks. Caller passes a per-warp pointer. static __attribute__((always_inline)) uint32_t a_sp_blockmajor_idx(uint32_t r, uint32_t c) { + constexpr uint32_t sparse_blk_elems = tcM * tcK * i_ratio; uint32_t half_k_blk = c / (tcK * i_ratio); uint32_t k_in_elem = c % (tcK * i_ratio); uint32_t m_blk = r / tcM; uint32_t i_in = r % tcM; - return (half_k_blk * m_steps + m_blk) * a_blk_elems + return (half_k_blk * m_steps + m_blk) * sparse_blk_elems + i_in * (tcK * i_ratio) + k_in_elem; } @@ -812,7 +824,8 @@ struct wgmma_context { // so wg_meta_total_bytes = NT*4. Unused SRAM cells are packed with zero // to keep the load shape uniform. static constexpr uint32_t sp_rtl_i_ratio = 32 / It::bits; - static constexpr uint32_t wg_meta_banks = m_steps * (k_steps / 2); + static constexpr uint32_t sparse_k_steps = (k_steps > 1) ? (k_steps / 2) : 1; + static constexpr uint32_t wg_meta_banks = m_steps * sparse_k_steps; static constexpr uint32_t wg_meta_row_bits = tcK * 2 * sp_rtl_i_ratio; static constexpr uint32_t wg_meta_stride_words = (tcM * wg_meta_row_bits + 31) / 32; static constexpr uint32_t wg_meta_stride_bytes = wg_meta_stride_words * 4; @@ -844,6 +857,7 @@ struct wgmma_context { static_assert(tcM * tcK == NT, "wgmma block-major load assumes canonical config (TC_M*TC_K == NT)"); constexpr uint32_t k_row_elems = tcK * i_ratio; + constexpr uint32_t sparse_blk_elems = tcM * k_row_elems; uint32_t lane = vx_thread_id(); uint32_t i_in = lane / tcK; uint32_t k_in_elem = (lane % tcK) * i_ratio; @@ -852,7 +866,8 @@ struct wgmma_context { uint32_t elem_off; if (ldm == 0) { // Block-major SMEM: blocks contiguous, k-block outer. - elem_off = (k_blk * m_steps + m_blk) * a_blk_elems + constexpr uint32_t block_elems = is_sparse ? sparse_blk_elems : a_blk_elems; + elem_off = (k_blk * m_steps + m_blk) * block_elems + i_in * k_row_elems + k_in_elem; } else { @@ -865,7 +880,7 @@ struct wgmma_context { }; if constexpr (is_sparse) { - constexpr uint32_t sp_k_steps_local = k_steps / 2; + constexpr uint32_t sp_k_steps_local = sparse_k_steps; constexpr uint32_t a_regs = m_steps * sp_k_steps_local; detail::unroll_for([&](auto r) { uint32_t m_blk = r / sp_k_steps_local; @@ -874,9 +889,24 @@ struct wgmma_context { }); } else { detail::unroll_for([&](auto r) { - uint32_t m_blk = r / k_steps; - uint32_t k_blk = r % k_steps; - dst.data[r] = load_reg(m_blk, k_blk); + uint32_t m_blk = r / a_reg_k_steps; + uint32_t k_blk = r % a_reg_k_steps; + if constexpr (k_steps == 1) { + uint32_t elem_off; + if (ldm == 0) { + elem_off = m_blk * a_blk_elems + + i_in * (fedpK * i_ratio) + + k_blk * k_row_elems + + k_in_elem; + } else { + elem_off = (m_blk * tcM + i_in) * uint32_t(ldm) + + (k_blk * k_row_elems + k_in_elem); + } + dst.data[r] = *reinterpret_cast( + reinterpret_cast(src) + elem_off); + } else { + dst.data[r] = load_reg(m_blk, k_blk); + } }); } } @@ -904,6 +934,11 @@ struct wgmma_context { ); } + template + static __attribute__((always_inline)) void load_mx_metadata(Frag& frag, const void* meta_mx_ptr) { + ctx_c::load_mx_metadata(frag, meta_mx_ptr); + } + // Store accumulator with n-major register layout: r = n * m_steps + m template static __attribute__((always_inline)) void store_matrix_sync(void *dst, const Frag &src, size_t ldm) { diff --git a/sw/runtime/include/tensor_mx.h b/sw/runtime/include/tensor_mx.h index 3353b9f289..ae59b14d89 100644 --- a/sw/runtime/include/tensor_mx.h +++ b/sw/runtime/include/tensor_mx.h @@ -127,20 +127,6 @@ struct mx_format_t { } }; -template <> -struct mx_format_t { - using storage_type = int8_t; - static constexpr bool needs_tensor_scale = false; - - static uint8_t select_scale(float max_abs) { - return select_e8m0_scale(max_abs, 127.0f / 64.0f); - } - - static storage_type convert(float v, uint8_t sf) { - return static_cast(rv_ftomxint8_s(bit_cast(v), sf, 0, nullptr)); - } -}; - template <> struct mx_format_t { using storage_type = uint8_t; diff --git a/sw/runtime/include/tensor_sp.h b/sw/runtime/include/tensor_sp.h index cad6380218..d8bba5345f 100644 --- a/sw/runtime/include/tensor_sp.h +++ b/sw/runtime/include/tensor_sp.h @@ -126,7 +126,7 @@ inline void select_top2(const typename TensorT::dtype* data, template inline float element_magnitude(const typename TensorT::dtype* data, uint32_t offset) { auto val = data_accessor_t::read(data, offset); - if constexpr (std::is_same_v || std::is_same_v) { + if constexpr (std::is_same_v) { return std::abs(static_cast(static_cast(val))); } else if constexpr (std::is_same_v || std::is_same_v diff --git a/tests/regression/Makefile b/tests/regression/Makefile index d11fca954e..be60b6363c 100644 --- a/tests/regression/Makefile +++ b/tests/regression/Makefile @@ -8,7 +8,7 @@ TESTS := \ vecadd vecadd_v1 sgemm sgemmx sgemm_v1 conv3 relu sgemv \ sgemm2 sgemm2_v1 madmax stencil3d raycast bfs jacobi pathfinder \ softmax \ - sgemm_tcu sgemm_tcu_mx sgemm_tcu_sp_mx sgemm_tcu_wg sgemm_tcu_sp sgemm_tcu_wg_sp \ + sgemm_tcu sgemm_tcu_mx sgemm_tcu_sp_mx sgemm_tcu_wg sgemm_tcu_wg_mx sgemm_tcu_sp sgemm_tcu_wg_sp \ sgemm2_dxa sgemm2_tcu sgemm_tcu_wg_dxa sgemm_tcu_wg_sp_dxa \ sgemm2_dxa_mcast sgemm_tcu_wg_dxa_mcast \ dxa_copy dxa_copy_mcast dxa_kmajor_check \ diff --git a/tests/regression/sgemm_tcu_mx/main.cpp b/tests/regression/sgemm_tcu_mx/main.cpp index 8887c60140..a0cff63253 100644 --- a/tests/regression/sgemm_tcu_mx/main.cpp +++ b/tests/regression/sgemm_tcu_mx/main.cpp @@ -59,10 +59,8 @@ static void pack_mx_a_metadata(std::vector &packed, uint32_t logical_tileK = cfg::tileK * kElemsPerByte; uint32_t num_k_tiles = K_logical / logical_tileK; uint32_t scale_blocks_k = K_logical / vt::ITYPE::ele_block; - uint32_t tile_scale_blocks_k = logical_tileK / vt::ITYPE::ele_block; - if (tile_scale_blocks_k == 0) { - tile_scale_blocks_k = 1; - } + uint32_t tile_scale_blocks_k = + std::max(1u, (logical_tileK + vt::ITYPE::ele_block - 1) / vt::ITYPE::ele_block); packed.assign(num_tile_rows * num_k_tiles * VX_CFG_NUM_THREADS, 0); for (uint32_t tr = 0; tr < num_tile_rows; ++tr) { @@ -91,10 +89,8 @@ static void pack_mx_b_metadata(std::vector &packed, uint32_t num_tile_cols = N / cfg::tileN; uint32_t logical_tileK = cfg::tileK * kElemsPerByte; uint32_t num_k_tiles = K_logical / logical_tileK; - uint32_t tile_scale_blocks_k = logical_tileK / vt::ITYPE::ele_block; - if (tile_scale_blocks_k == 0) { - tile_scale_blocks_k = 1; - } + uint32_t tile_scale_blocks_k = + std::max(1u, (logical_tileK + vt::ITYPE::ele_block - 1) / vt::ITYPE::ele_block); packed.assign(num_tile_cols * num_k_tiles * VX_CFG_NUM_THREADS, 0); for (uint32_t tc = 0; tc < num_tile_cols; ++tc) { @@ -127,11 +123,6 @@ static float dequantize_mx_value(const itype_t *data, return bit_cast(rv_mxfp8tof_s(data[offset], sf, 0, nullptr)); case vt::mxbf8::id: return bit_cast(rv_mxbf8tof_s(data[offset], sf, 0, nullptr)); - case vt::mxint8::id: - { - float scale = std::ldexp(1.0f, static_cast(sf) - 127); - return (static_cast(data[offset]) / 64.0f) * scale; - } case vt::mxfp4::id: { uint8_t q = read_nibble(reinterpret_cast(data), offset); @@ -147,16 +138,6 @@ static float dequantize_mx_value(const itype_t *data, } } -static int32_t trunc_shift(int32_t value, int32_t shift) { - if (shift >= 0) { - return value << shift; - } - uint32_t abs_shift = static_cast(-shift); - uint32_t mag = value < 0 ? static_cast(-value) : static_cast(value); - int32_t scaled = static_cast(mag >> abs_shift); - return value < 0 ? -scaled : scaled; -} - static void matmul_cpu(otype_t *C, const itype_t *A, const itype_t *B, @@ -170,19 +151,7 @@ static void matmul_cpu(otype_t *C, uint32_t scale_blocks_k = K_logical / vt::ITYPE::ele_block; for (uint32_t m = 0; m < M; ++m) { for (uint32_t n = 0; n < N; ++n) { - if constexpr (std::is_same::value) { - int32_t sum = 0; - for (uint32_t k = 0; k < K_logical; ++k) { - uint32_t scale_k = k / vt::ITYPE::ele_block; - uint8_t sf_a = scale_a[m * scale_blocks_k + scale_k]; - uint8_t sf_b = scale_b[scale_k * N + n]; - int32_t shift = static_cast(sf_a) - 133 + static_cast(sf_b) - 133; - int32_t product = static_cast(A[m * K_logical + k]) * static_cast(B[n * K_logical + k]); - sum += trunc_shift(product, shift); - } - C[m * N + n] = sum; - } else { - float sum = 0.0f; + float sum = 0.0f; for (uint32_t k = 0; k < K_logical; ++k) { uint32_t scale_k = k / vt::ITYPE::ele_block; auto a = dequantize_mx_value(A, scale_a, m * K_logical + k, @@ -191,8 +160,7 @@ static void matmul_cpu(otype_t *C, scale_k * N + n, tensor_scale_b); sum += a * b; } - C[m * N + n] = sum; - } + C[m * N + n] = sum; } } } @@ -434,14 +402,12 @@ int main(int argc, char *argv[]) { A_tensor_scale, B_tensor_scale, M, N, K_logical); int errors = 0; - float rel_tol = std::is_same::value ? 0.0f : - (std::is_same::value || std::is_same::value) ? 0.25f : 0.05f; + float rel_tol = (std::is_same::value || std::is_same::value) ? 0.25f : 0.05f; for (uint32_t i = 0; i < h_ref.size(); ++i) { float actual = static_cast(h_C[i]); float expected = static_cast(h_ref[i]); float diff = std::abs(actual - expected); - float tol = std::is_same::value ? 0.0f : - std::max(1.0e-3f, std::abs(expected) * rel_tol); + float tol = std::max(1.0e-3f, std::abs(expected) * rel_tol); if (!std::isfinite(actual) || !std::isfinite(expected) || !std::isfinite(diff) || diff > tol) { if (errors < MAX_ERRORS) { printf("*** error: [%d] expected=%f, actual=%f, diff=%f, tol=%f\n", diff --git a/tests/regression/sgemm_tcu_sp_mx/main.cpp b/tests/regression/sgemm_tcu_sp_mx/main.cpp index bab28a2871..8dabef023d 100644 --- a/tests/regression/sgemm_tcu_sp_mx/main.cpp +++ b/tests/regression/sgemm_tcu_sp_mx/main.cpp @@ -184,11 +184,6 @@ static float dequantize_mx_value(const itype_t *data, return bit_cast(rv_mxfp8tof_s(data[offset], sf, 0, nullptr)); case vt::mxbf8::id: return bit_cast(rv_mxbf8tof_s(data[offset], sf, 0, nullptr)); - case vt::mxint8::id: - { - float scale = std::ldexp(1.0f, static_cast(sf) - 127); - return (static_cast(data[offset]) / 64.0f) * scale; - } case vt::mxfp4::id: { uint8_t q = read_nibble(reinterpret_cast(data), offset); @@ -204,16 +199,6 @@ static float dequantize_mx_value(const itype_t *data, } } -static int32_t trunc_shift(int32_t value, int32_t shift) { - if (shift >= 0) { - return value << shift; - } - uint32_t abs_shift = static_cast(-shift); - uint32_t mag = value < 0 ? static_cast(-value) : static_cast(value); - int32_t scaled = static_cast(mag >> abs_shift); - return value < 0 ? -scaled : scaled; -} - static void matmul_cpu(otype_t *C, const itype_t *A, const itype_t *B, @@ -227,19 +212,7 @@ static void matmul_cpu(otype_t *C, uint32_t scale_blocks_k = K_logical / vt::ITYPE::ele_block; for (uint32_t m = 0; m < M; ++m) { for (uint32_t n = 0; n < N; ++n) { - if constexpr (std::is_same::value) { - int32_t sum = 0; - for (uint32_t k = 0; k < K_logical; ++k) { - uint32_t scale_k = k / vt::ITYPE::ele_block; - uint8_t sf_a = scale_a[m * scale_blocks_k + scale_k]; - uint8_t sf_b = scale_b[scale_k * N + n]; - int32_t shift = static_cast(sf_a) - 133 + static_cast(sf_b) - 133; - int32_t product = static_cast(A[m * K_logical + k]) * static_cast(B[n * K_logical + k]); - sum += trunc_shift(product, shift); - } - C[m * N + n] = sum; - } else { - float sum = 0.0f; + float sum = 0.0f; for (uint32_t k = 0; k < K_logical; ++k) { uint32_t scale_k = k / vt::ITYPE::ele_block; auto a = dequantize_mx_value(A, scale_a, m * K_logical + k, @@ -248,8 +221,7 @@ static void matmul_cpu(otype_t *C, scale_k * N + n, tensor_scale_b); sum += a * b; } - C[m * N + n] = sum; - } + C[m * N + n] = sum; } } } @@ -510,14 +482,12 @@ int main(int argc, char *argv[]) { A_tensor_scale, B_tensor_scale, M, N, K_logical); int errors = 0; - float rel_tol = std::is_same::value ? 0.0f : - (std::is_same::value || std::is_same::value) ? 0.25f : 0.05f; + float rel_tol = (std::is_same::value || std::is_same::value) ? 0.25f : 0.05f; for (uint32_t i = 0; i < h_ref.size(); ++i) { float actual = static_cast(h_C[i]); float expected = static_cast(h_ref[i]); float diff = std::abs(actual - expected); - float tol = std::is_same::value ? 0.0f : - std::max(1.0e-3f, std::abs(expected) * rel_tol); + float tol = std::max(1.0e-3f, std::abs(expected) * rel_tol); if (!std::isfinite(actual) || !std::isfinite(expected) || !std::isfinite(diff) || diff > tol) { if (errors < MAX_ERRORS) { printf("*** error: [%d] expected=%f, actual=%f, diff=%f, tol=%f\n", diff --git a/tests/regression/sgemm_tcu_wg/Makefile b/tests/regression/sgemm_tcu_wg/Makefile index 66934cf470..d67265b4c7 100644 --- a/tests/regression/sgemm_tcu_wg/Makefile +++ b/tests/regression/sgemm_tcu_wg/Makefile @@ -13,11 +13,9 @@ VX_SRCS := $(SRC_DIR)/kernel.cpp CONFIGS := $(if $(findstring -DVX_CFG_NUM_WARPS=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_NUM_WARPS=4) CONFIGS := $(if $(findstring -DVX_CFG_ISSUE_WIDTH=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_ISSUE_WIDTH=4) -CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE -DWGMMA_RS +CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE +CONFIGS := $(if $(findstring -DWGMMA_SS,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DWGMMA_RS) -# RTLsim is roughly two orders of magnitude slower than SimX; keep the -# default matrix size small enough for a 60s sweep budget. Override -# -m/-n/-k at the command line for full-size correctness sweeps. OPTS := $(OPTS) $(if $(findstring -m,$(OPTS)),,-m 16) OPTS := $(OPTS) $(if $(findstring -n,$(OPTS)),,-n 16) OPTS := $(OPTS) $(if $(findstring -k,$(OPTS)),,-k 16) diff --git a/tests/regression/sgemm_tcu_wg_mx/Makefile b/tests/regression/sgemm_tcu_wg_mx/Makefile new file mode 100644 index 0000000000..81d39703f5 --- /dev/null +++ b/tests/regression/sgemm_tcu_wg_mx/Makefile @@ -0,0 +1,24 @@ +ROOT_DIR := $(realpath ../../..) +include $(ROOT_DIR)/config.mk + +CONFIGS := $(if $(findstring -DVX_CFG_EXT_TCU_ENABLE,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_EXT_TCU_ENABLE) +CONFIGS := $(if $(findstring -DVX_CFG_TCU_MX_ENABLE,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_TCU_MX_ENABLE) +CONFIGS := $(if $(findstring -DVX_CFG_TCU_FP8_ENABLE,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_TCU_FP8_ENABLE) +CONFIGS := $(if $(findstring -DVX_CFG_NUM_WARPS=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_NUM_WARPS=4) +CONFIGS := $(if $(findstring -DVX_CFG_ISSUE_WIDTH=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_ISSUE_WIDTH=4) +CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE +CONFIGS := $(if $(findstring -DWGMMA_SS,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DWGMMA_RS) + +PROJECT := sgemm_tcu_wg_mx +SRC_DIR := $(VORTEX_HOME)/tests/regression/$(PROJECT) +SRCS := $(SRC_DIR)/main.cpp $(SW_COMMON_DIR)/rvfloats.cpp $(SW_COMMON_DIR)/softfloat_ext.cpp +VX_SRCS := $(SRC_DIR)/kernel.cpp + +OPTS := $(OPTS) $(if $(findstring -m,$(OPTS)),,-m 64) +OPTS := $(OPTS) $(if $(findstring -n,$(OPTS)),,-n 64) +OPTS := $(OPTS) $(if $(findstring -k,$(OPTS)),,-k 64) +CXXFLAGS += -I$(THIRD_PARTY_DIR)/softfloat/source/include +LDFLAGS += $(THIRD_PARTY_DIR)/softfloat/build/Linux-x86_64-GCC/softfloat.a +KERNEL_LIB := vortex2 + +include ../common.mk diff --git a/tests/regression/sgemm_tcu_wg_mx/common.h b/tests/regression/sgemm_tcu_wg_mx/common.h new file mode 100644 index 0000000000..f6cc3275d0 --- /dev/null +++ b/tests/regression/sgemm_tcu_wg_mx/common.h @@ -0,0 +1,25 @@ +#ifndef _COMMON_H_ +#define _COMMON_H_ + +#include + +#ifndef WGMMA_NRC +#define WGMMA_NRC 8 +#endif +#ifndef ITYPE +#define ITYPE mxfp8 +#endif +#ifndef OTYPE +#define OTYPE fp32 +#endif + +typedef struct { + uint32_t M, N, K; + uint64_t A_addr, B_addr, C_addr; + uint64_t MX_A_addr, MX_B_addr; +#ifdef TCU_MX_TLS + uint64_t A_tensor_scale_addr, B_tensor_scale_addr; +#endif +} kernel_arg_t; + +#endif diff --git a/tests/regression/sgemm_tcu_wg_mx/kernel.cpp b/tests/regression/sgemm_tcu_wg_mx/kernel.cpp new file mode 100644 index 0000000000..f370bbd583 --- /dev/null +++ b/tests/regression/sgemm_tcu_wg_mx/kernel.cpp @@ -0,0 +1,83 @@ +#include "common.h" +#include +#include +#include + +namespace vt = vortex::tensor; +using ctx = vt::wgmma_context; + +__kernel void kernel_main(kernel_arg_t* __UNIFORM__ arg) { + static_assert(sizeof(ctx::input_t) == 1, "MX WGMMA inputs use byte storage"); + auto pA = reinterpret_cast(arg->A_addr); + auto pB = reinterpret_cast(arg->B_addr); + auto pC = reinterpret_cast(arg->C_addr); + auto pMxA = reinterpret_cast(arg->MX_A_addr); + auto pMxB = reinterpret_cast(arg->MX_B_addr); +#ifdef TCU_MX_TLS + auto pATensorScale = reinterpret_cast(arg->A_tensor_scale_addr); + auto pBTensorScale = reinterpret_cast(arg->B_tensor_scale_addr); +#endif + + uint32_t N = arg->N; + uint32_t K = arg->K; + uint32_t tid = threadIdx.x; + uint32_t num_threads = blockDim.x; + uint32_t warp_rank = tid / VX_CFG_NUM_THREADS; + uint32_t num_warps = num_threads / VX_CFG_NUM_THREADS; + uint32_t cta_m = num_warps * ctx::xtileM; + uint32_t tile_row = blockIdx.y * cta_m; + uint32_t tile_col = blockIdx.x * ctx::xtileN; + uint32_t num_k_tiles = K / ctx::tileK; + + auto A_smem = reinterpret_cast(__local_mem()); + auto B_smem = A_smem + cta_m * ctx::tileK; + ctx::fragment_a fragA; + ctx::fragment_b fragB; + ctx::fragment_acc fragC; + ctx::fill_fragment(fragC, 0); + + for (uint32_t k = 0; k < K; k += ctx::tileK) { + uint32_t k_tile = k / ctx::tileK; + uint32_t a_bytes = cta_m * ctx::tileK; + for (uint32_t idx = tid; idx < a_bytes; idx += num_threads) { + uint32_t row = idx / ctx::tileK; + uint32_t col = idx % ctx::tileK; + A_smem[idx] = pA[(tile_row + row) * K + k + col]; + } + uint32_t b_bytes = ctx::xtileN * ctx::tileK; + for (uint32_t idx = tid; idx < b_bytes; idx += num_threads) { + uint32_t col = idx / ctx::tileK; + uint32_t k_byte = idx % ctx::tileK; + B_smem[idx] = pB[(tile_col + col) * K + k + k_byte]; + } + __syncthreads(); + + uint32_t tile_m = blockIdx.y * num_warps + warp_rank; + auto pTileMxA = pMxA + (tile_m * num_k_tiles + k_tile) * VX_CFG_NUM_THREADS; + auto pTileMxB = pMxB + (blockIdx.x * num_k_tiles + k_tile) * VX_CFG_NUM_THREADS; + ctx::load_mx_metadata(fragA, pTileMxA); + ctx::load_mx_metadata(fragB, pTileMxB); + + auto A_warp = A_smem + warp_rank * ctx::xtileM * ctx::tileK; + auto desc_b = vt::vx_make_smem_desc(B_smem, ctx::tileK); +#if defined(WGMMA_RS) && (WGMMA_NRC <= 16) + ctx::load_matrix_sync(fragA, A_warp, ctx::tileK); + ctx::wgmma_sync(fragC, fragA, desc_b, fragC); +#else + auto desc_a = vt::vx_make_smem_desc(A_warp, ctx::tileK); + ctx::wgmma_sync(fragC, desc_a, desc_b, fragC); +#endif + __syncthreads(); + } + +#ifdef TCU_MX_TLS + if constexpr (std::is_same::value) { + float tensor_scale = (*pATensorScale) * (*pBTensorScale); + for (uint32_t r = 0; r < ctx::fragment_acc::NR; ++r) + fragC.data[r] *= tensor_scale; + } +#endif + auto out = pC + (tile_row + warp_rank * ctx::xtileM) * N + tile_col; + ctx::store_matrix_sync(out, fragC, N); +} diff --git a/tests/regression/sgemm_tcu_wg_mx/main.cpp b/tests/regression/sgemm_tcu_wg_mx/main.cpp new file mode 100644 index 0000000000..f149b1dffc --- /dev/null +++ b/tests/regression/sgemm_tcu_wg_mx/main.cpp @@ -0,0 +1,435 @@ +#include "common.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define MAX_ERRORS 100 + +#define RT_CHECK(_expr) \ + do { \ + int _ret = _expr; \ + if (0 == _ret) \ + break; \ + printf("Error: '%s' returned %d!\n", #_expr, (int)_ret); \ + cleanup(); \ + exit(-1); \ + } while (false) + +using namespace vortex; +namespace vt = tensor; + +using cfg = vt::wgmma_config_t; +static constexpr uint32_t kTileM = cfg::xtileM; +static constexpr uint32_t kTileN = cfg::xtileN; +using itype_t = typename vt::ITYPE::dtype; +using otype_t = typename vt::OTYPE::dtype; + +static constexpr uint32_t kElemsPerByte = (vt::ITYPE::bits < 8) ? (8 / vt::ITYPE::bits) : 1; +static constexpr uint32_t kScalePack = 4; + +static uint8_t read_nibble(const uint8_t *ptr, uint32_t offset) { + uint8_t value = ptr[offset / 2]; + return (offset & 1) ? (value >> 4) : (value & 0x0f); +} + +static uint32_t pack4(const std::vector &scales, uint32_t start) { + uint32_t word = 0; + for (uint32_t i = 0; i < kScalePack; ++i) { + if ((start + i) < scales.size()) { + word |= static_cast(scales[start + i]) << (8 * i); + } + } + return word; +} + +static void pack_mx_a_metadata(std::vector &packed, + const std::vector &scales, + uint32_t M, + uint32_t K_logical) { + uint32_t num_tile_rows = M / kTileM; + uint32_t logical_tileK = cfg::tileK * kElemsPerByte; + uint32_t num_k_tiles = K_logical / logical_tileK; + uint32_t scale_blocks_k = K_logical / vt::ITYPE::ele_block; + uint32_t tile_scale_blocks_k = + std::max(1u, (logical_tileK + vt::ITYPE::ele_block - 1) / vt::ITYPE::ele_block); + + packed.assign(num_tile_rows * num_k_tiles * VX_CFG_NUM_THREADS, 0); + for (uint32_t tr = 0; tr < num_tile_rows; ++tr) { + for (uint32_t kt = 0; kt < num_k_tiles; ++kt) { + std::vector tile_scales; + tile_scales.reserve(kTileM * tile_scale_blocks_k); + uint32_t first_block = (kt * logical_tileK) / vt::ITYPE::ele_block; + for (uint32_t m = 0; m < kTileM; ++m) { + uint32_t row = tr * kTileM + m; + for (uint32_t kb = 0; kb < tile_scale_blocks_k; ++kb) { + tile_scales.push_back(scales[row * scale_blocks_k + first_block + kb]); + } + } + uint32_t base = (tr * num_k_tiles + kt) * VX_CFG_NUM_THREADS; + for (uint32_t w = 0; w < VX_CFG_NUM_THREADS && (w * kScalePack) < tile_scales.size(); ++w) { + packed[base + w] = pack4(tile_scales, w * kScalePack); + } + } + } +} + +static void pack_mx_b_metadata(std::vector &packed, + const std::vector &scales, + uint32_t N, + uint32_t K_logical) { + uint32_t num_tile_cols = N / kTileN; + uint32_t logical_tileK = cfg::tileK * kElemsPerByte; + uint32_t num_k_tiles = K_logical / logical_tileK; + uint32_t tile_scale_blocks_k = + std::max(1u, (logical_tileK + vt::ITYPE::ele_block - 1) / vt::ITYPE::ele_block); + + packed.assign(num_tile_cols * num_k_tiles * VX_CFG_NUM_THREADS, 0); + for (uint32_t tc = 0; tc < num_tile_cols; ++tc) { + for (uint32_t kt = 0; kt < num_k_tiles; ++kt) { + std::vector tile_scales; + tile_scales.reserve(kTileN * tile_scale_blocks_k); + uint32_t first_block = (kt * logical_tileK) / vt::ITYPE::ele_block; + for (uint32_t n = 0; n < kTileN; ++n) { + uint32_t col = tc * kTileN + n; + for (uint32_t kb = 0; kb < tile_scale_blocks_k; ++kb) { + tile_scales.push_back(scales[(first_block + kb) * N + col]); + } + } + uint32_t base = (tc * num_k_tiles + kt) * VX_CFG_NUM_THREADS; + for (uint32_t w = 0; w < VX_CFG_NUM_THREADS && (w * kScalePack) < tile_scales.size(); ++w) { + packed[base + w] = pack4(tile_scales, w * kScalePack); + } + } + } +} + +static float dequantize_mx_value(const itype_t *data, + const std::vector &scales, + uint32_t offset, + uint32_t scale_index, + float tensor_scale) { + uint8_t sf = scales[scale_index]; + switch (vt::ITYPE::id) { + case vt::mxfp8::id: + return bit_cast(rv_mxfp8tof_s(data[offset], sf, 0, nullptr)); + case vt::mxbf8::id: + return bit_cast(rv_mxbf8tof_s(data[offset], sf, 0, nullptr)); + case vt::mxfp4::id: + { + uint8_t q = read_nibble(reinterpret_cast(data), offset); + return bit_cast(rv_mxfp4tof_s(q, sf, 0, nullptr)); + } + case vt::nvfp4::id: + { + uint8_t q = read_nibble(reinterpret_cast(data), offset); + return bit_cast(rv_nvfp4tof_s(q, sf, 0, nullptr)) * tensor_scale; + } + default: + std::abort(); + } +} + +static void matmul_cpu(otype_t *C, + const itype_t *A, + const itype_t *B, + const std::vector &scale_a, + const std::vector &scale_b, + float tensor_scale_a, + float tensor_scale_b, + uint32_t M, + uint32_t N, + uint32_t K_logical) { + uint32_t scale_blocks_k = K_logical / vt::ITYPE::ele_block; + for (uint32_t m = 0; m < M; ++m) { + for (uint32_t n = 0; n < N; ++n) { + float sum = 0.0f; + for (uint32_t k = 0; k < K_logical; ++k) { + uint32_t scale_k = k / vt::ITYPE::ele_block; + auto a = dequantize_mx_value(A, scale_a, m * K_logical + k, + m * scale_blocks_k + scale_k, tensor_scale_a); + auto b = dequantize_mx_value(B, scale_b, n * K_logical + k, + scale_k * N + n, tensor_scale_b); + sum += a * b; + } + C[m * N + n] = sum; + } + } +} + +template +static bool quantize_inputs(typename FormatT::dtype *A, + typename FormatT::dtype *B, + std::vector &scale_a, + std::vector &scale_b, + float &A_tensor_scale, + float &B_tensor_scale, + const float *A_dense, + const float *B_dense, + uint32_t M, + uint32_t N, + uint32_t K_logical) { + if constexpr (std::is_same::value) { + return vt::quantize_mx_a_rowmajor( + reinterpret_cast(A), scale_a, A_tensor_scale, + A_dense, M, K_logical) + && vt::quantize_mx_b_colmajor( + reinterpret_cast(B), scale_b, B_tensor_scale, + B_dense, K_logical, N); + } else { + return vt::quantize_mx_a_rowmajor( + A, scale_a, A_dense, M, K_logical) + && vt::quantize_mx_b_colmajor( + B, scale_b, B_dense, K_logical, N); + } +} + +const char *kernel_file = "kernel.vxbin"; + +uint32_t xm = 32; +uint32_t xn = 32; +uint32_t xk = 32; + +vx_device_h device = nullptr; +vx_queue_h queue = nullptr; +vx_module_h module_ = nullptr; +vx_kernel_h kernel = nullptr; +vx_buffer_h A_buffer = nullptr; +vx_buffer_h B_buffer = nullptr; +vx_buffer_h C_buffer = nullptr; +vx_buffer_h MX_A_buffer = nullptr; +vx_buffer_h MX_B_buffer = nullptr; +#ifdef TCU_MX_TLS +vx_buffer_h A_tensor_scale_buffer = nullptr; +vx_buffer_h B_tensor_scale_buffer = nullptr; +#endif +kernel_arg_t kernel_arg = {}; + +static void show_usage() { + std::cout << "Vortex SGEMM TCU MX Test." << std::endl; + std::cout << "Usage: [-m M] [-n N] [-k K] [-h]" << std::endl; +} + +static void parse_args(int argc, char **argv) { + int c; + while ((c = getopt(argc, argv, "m:n:k:h")) != -1) { + switch (c) { + case 'm': xm = atoi(optarg); break; + case 'n': xn = atoi(optarg); break; + case 'k': xk = atoi(optarg); break; + case 'h': + show_usage(); + exit(0); + default: + show_usage(); + exit(-1); + } + } +} + +void cleanup() { + if (device) { + vx_mem_free(A_buffer); + vx_mem_free(B_buffer); + vx_mem_free(C_buffer); + vx_mem_free(MX_A_buffer); + vx_mem_free(MX_B_buffer); +#ifdef TCU_MX_TLS + vx_mem_free(A_tensor_scale_buffer); + vx_mem_free(B_tensor_scale_buffer); +#endif + if (kernel) vx_kernel_release(kernel); + if (module_) vx_module_release(module_); + if (queue) vx_queue_release(queue); + vx_dump_perf(device, stdout); + vx_dev_close(device); + } +} + +int main(int argc, char *argv[]) { + parse_args(argc, argv); + std::srand(50); + + RT_CHECK(vx_dev_open(&device)); + vx_queue_info_t qi = {sizeof(qi), nullptr, VX_QUEUE_PRIORITY_NORMAL, 0}; + RT_CHECK(vx_queue_create(device, &qi, &queue)); + + uint64_t isa_flags; + RT_CHECK(vx_dev_caps(device, VX_CAPS_ISA_FLAGS, &isa_flags)); + if ((isa_flags & VX_ISA_EXT_TCU) == 0) { + std::cout << "TCU extension not supported!" << std::endl; + cleanup(); + return -1; + } + + uint64_t NT; + RT_CHECK(vx_dev_caps(device, VX_CAPS_NUM_THREADS, &NT)); + if (NT != VX_CFG_NUM_THREADS) { + std::cout << "Error: device warp size (" << NT << ") must match VX_CFG_NUM_THREADS=" << VX_CFG_NUM_THREADS << "!" << std::endl; + return -1; + } + + uint32_t M = xm; + uint32_t N = xn; + uint32_t K_logical = xk; + uint32_t logical_tileK = cfg::tileK * kElemsPerByte; + uint32_t K_storage = K_logical / kElemsPerByte; + + uint32_t warps_per_cta = VX_CFG_ISSUE_WIDTH; + uint32_t cta_m = warps_per_cta * kTileM; + if ((M % cta_m) != 0 || (N % kTileN) != 0 || (K_logical % logical_tileK) != 0) { + std::cout << "Error: M/N/K must be multiples of CTA M=" << cta_m + << " N=" << kTileN << " K=" << logical_tileK << std::endl; + return -1; + } + if ((K_logical % vt::ITYPE::ele_block) != 0) { + std::cout << "Error: K must be a multiple of MX block size=" << vt::ITYPE::ele_block << std::endl; + return -1; + } + + size_t sizeA = M * K_storage; + size_t sizeB = K_storage * N; + size_t sizeC = M * N; + uint32_t grid_dim[2] = {N / kTileN, M / cta_m}; + uint32_t block_dim[2] = {warps_per_cta * (uint32_t)NT, 1}; + + std::cout << "input data type: " << vt::ITYPE::name << " (id=" << vt::ITYPE::id << ")" << std::endl; + std::cout << "output data type: " << vt::OTYPE::name << " (id=" << vt::OTYPE::id << ")" << std::endl; + std::cout << "MMA Tile Dimension: M=" << kTileM << ", N=" << kTileN + << ", K(logical)=" << logical_tileK << ", K(storage)=" << cfg::tileK << std::endl; + + kernel_arg.M = M; + kernel_arg.N = N; + kernel_arg.K = K_storage; +#ifdef TCU_MX_TLS + kernel_arg.A_tensor_scale_addr = 0; + kernel_arg.B_tensor_scale_addr = 0; +#endif + float A_tensor_scale = 1.0f; + float B_tensor_scale = 1.0f; + + RT_CHECK(vx_mem_alloc(device, sizeA * sizeof(itype_t), VX_MEM_READ, &A_buffer)); + RT_CHECK(vx_mem_address(A_buffer, &kernel_arg.A_addr)); + RT_CHECK(vx_mem_alloc(device, sizeB * sizeof(itype_t), VX_MEM_READ, &B_buffer)); + RT_CHECK(vx_mem_address(B_buffer, &kernel_arg.B_addr)); + RT_CHECK(vx_mem_alloc(device, sizeC * sizeof(otype_t), VX_MEM_WRITE, &C_buffer)); + RT_CHECK(vx_mem_address(C_buffer, &kernel_arg.C_addr)); + + std::vector h_A_dense(M * K_logical); + std::vector h_B_dense(K_logical * N); + for (auto &v : h_A_dense) { + v = (static_cast(std::rand()) / RAND_MAX) * 2.0f - 1.0f; + } + for (auto &v : h_B_dense) { + v = (static_cast(std::rand()) / RAND_MAX) * 2.0f - 1.0f; + } + + std::vector h_A(sizeA); + std::vector h_B(sizeB); + std::vector scale_a; + std::vector scale_b; + + bool ok = quantize_inputs( + h_A.data(), h_B.data(), scale_a, scale_b, A_tensor_scale, B_tensor_scale, + h_A_dense.data(), h_B_dense.data(), M, N, K_logical); + if (!ok) { + std::cout << "Error: MX quantization failed!" << std::endl; + return -1; + } + + std::vector h_mx_a; + std::vector h_mx_b; + pack_mx_a_metadata(h_mx_a, scale_a, M, K_logical); + pack_mx_b_metadata(h_mx_b, scale_b, N, K_logical); + + RT_CHECK(vx_mem_alloc(device, h_mx_a.size() * sizeof(uint32_t), VX_MEM_READ, &MX_A_buffer)); + RT_CHECK(vx_mem_address(MX_A_buffer, &kernel_arg.MX_A_addr)); + RT_CHECK(vx_mem_alloc(device, h_mx_b.size() * sizeof(uint32_t), VX_MEM_READ, &MX_B_buffer)); + RT_CHECK(vx_mem_address(MX_B_buffer, &kernel_arg.MX_B_addr)); +#ifdef TCU_MX_TLS + if constexpr (std::is_same::value) { + RT_CHECK(vx_mem_alloc(device, sizeof(float), VX_MEM_READ, &A_tensor_scale_buffer)); + RT_CHECK(vx_mem_address(A_tensor_scale_buffer, &kernel_arg.A_tensor_scale_addr)); + RT_CHECK(vx_mem_alloc(device, sizeof(float), VX_MEM_READ, &B_tensor_scale_buffer)); + RT_CHECK(vx_mem_address(B_tensor_scale_buffer, &kernel_arg.B_tensor_scale_addr)); + } +#endif + + RT_CHECK(vx_copy_to_dev(A_buffer, h_A.data(), 0, sizeA * sizeof(itype_t))); + RT_CHECK(vx_copy_to_dev(B_buffer, h_B.data(), 0, sizeB * sizeof(itype_t))); + RT_CHECK(vx_copy_to_dev(MX_A_buffer, h_mx_a.data(), 0, h_mx_a.size() * sizeof(uint32_t))); + RT_CHECK(vx_copy_to_dev(MX_B_buffer, h_mx_b.data(), 0, h_mx_b.size() * sizeof(uint32_t))); +#ifdef TCU_MX_TLS + if constexpr (std::is_same::value) { + RT_CHECK(vx_copy_to_dev(A_tensor_scale_buffer, &A_tensor_scale, 0, sizeof(float))); + RT_CHECK(vx_copy_to_dev(B_tensor_scale_buffer, &B_tensor_scale, 0, sizeof(float))); + } +#endif + + RT_CHECK(vx_module_load_file(device, kernel_file, &module_)); + RT_CHECK(vx_module_get_kernel(module_, "main", &kernel)); + + auto time_start = std::chrono::high_resolution_clock::now(); + vx_launch_info_t li = {}; + li.struct_size = sizeof(li); + li.kernel = kernel; + li.args_host = &kernel_arg; + li.args_size = sizeof(kernel_arg); + li.ndim = 2; + li.grid_dim[0] = grid_dim[0]; + li.grid_dim[1] = grid_dim[1]; + li.block_dim[0] = block_dim[0]; + li.block_dim[1] = block_dim[1]; + li.lmem_size = (cta_m + kTileN) * cfg::tileK * sizeof(itype_t); + + vx_event_h launch_ev = nullptr; + RT_CHECK(vx_enqueue_launch(queue, &li, 0, nullptr, &launch_ev)); + RT_CHECK(vx_event_wait_value(launch_ev, 1, VX_TIMEOUT_INFINITE)); + vx_event_release(launch_ev); + auto time_end = std::chrono::high_resolution_clock::now(); + double elapsed = std::chrono::duration_cast(time_end - time_start).count(); + printf("Elapsed time: %lg ms\n", elapsed); + + std::vector h_C(sizeC); + RT_CHECK(vx_copy_from_dev(h_C.data(), C_buffer, 0, sizeC * sizeof(otype_t))); + + std::vector h_ref(sizeC); + matmul_cpu(h_ref.data(), h_A.data(), h_B.data(), scale_a, scale_b, + A_tensor_scale, B_tensor_scale, M, N, K_logical); + + int errors = 0; + float rel_tol = (std::is_same::value || std::is_same::value) ? 0.25f : 0.05f; + for (uint32_t i = 0; i < h_ref.size(); ++i) { + float actual = static_cast(h_C[i]); + float expected = static_cast(h_ref[i]); + float diff = std::abs(actual - expected); + float tol = std::max(1.0e-3f, std::abs(expected) * rel_tol); + if (!std::isfinite(actual) || !std::isfinite(expected) || !std::isfinite(diff) || diff > tol) { + if (errors < MAX_ERRORS) { + printf("*** error: [%d] expected=%f, actual=%f, diff=%f, tol=%f\n", + i, expected, actual, diff, tol); + } + ++errors; + } + } + + cleanup(); + if (errors != 0) { + std::cout << "Found " << errors << " / " << sizeC << " errors!" << std::endl; + std::cout << "FAILED!" << std::endl; + return errors; + } + + std::cout << "PASSED!" << std::endl; + return 0; +} diff --git a/tests/regression/sgemm_tcu_wg_sp/Makefile b/tests/regression/sgemm_tcu_wg_sp/Makefile index fa25578b9c..54bb27c1a1 100644 --- a/tests/regression/sgemm_tcu_wg_sp/Makefile +++ b/tests/regression/sgemm_tcu_wg_sp/Makefile @@ -13,7 +13,10 @@ VX_SRCS := $(SRC_DIR)/kernel.cpp CONFIGS := $(if $(findstring -DVX_CFG_NUM_WARPS=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_NUM_WARPS=4) CONFIGS := $(if $(findstring -DVX_CFG_ISSUE_WIDTH=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_ISSUE_WIDTH=4) -CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_RS +CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE +CONFIGS := $(if $(findstring -DITYPE=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DITYPE=fp16) +CONFIGS := $(if $(findstring -DOTYPE=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DOTYPE=fp32) +CONFIGS := $(if $(findstring -DWGMMA_SS,$(CONFIGS)),$(CONFIGS),$(if $(findstring -DWGMMA_RS,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DWGMMA_RS)) CXXFLAGS += -I$(THIRD_PARTY_DIR)/softfloat/source/include diff --git a/tests/regression/sgemm_tcu_wg_sp/main.cpp b/tests/regression/sgemm_tcu_wg_sp/main.cpp index 372f5338af..30b2a1324e 100644 --- a/tests/regression/sgemm_tcu_wg_sp/main.cpp +++ b/tests/regression/sgemm_tcu_wg_sp/main.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -36,7 +37,7 @@ static constexpr uint32_t kTcK = wg_cfg_t::tcK; static constexpr uint32_t kTcM = wg_cfg_t::tcM; static constexpr uint32_t kMSteps = wg_cfg_t::m_steps; static constexpr uint32_t kKSteps = wg_cfg_t::k_steps; -static constexpr uint32_t kHalfKSteps = kKSteps / 2; +static constexpr uint32_t kHalfKSteps = (kKSteps > 1) ? (kKSteps / 2) : 1; static constexpr uint32_t kMetaRowBits = kTcK * 2 * kRtlIRatio; static constexpr uint32_t kMetaStrWords = (kTcM * kMetaRowBits + 31) / 32; static constexpr uint32_t kWgMetaBanks = kMSteps * kHalfKSteps; @@ -60,6 +61,67 @@ static constexpr uint32_t kDensePerSpStep = kTcK * kRtlIRatio * 2; using itype_t = vt::ITYPE::dtype; using otype_t = vt::OTYPE::dtype; +template +struct type_ops { + using dtype = typename T::dtype; + + static dtype generate() { + return dtype(float(rand()) / RAND_MAX); + } + + static float to_float(dtype value) { + return static_cast(value); + } +}; + +template <> +struct type_ops { + static uint16_t generate() { + auto value = float(rand()) / RAND_MAX; + return rv_ftoh_s(bit_cast(value), 0, nullptr); + } + + static float to_float(uint16_t value) { + return bit_cast(rv_htof_s(value, 0, nullptr)); + } +}; + +template <> +struct type_ops { + static uint16_t generate() { + auto value = float(rand()) / RAND_MAX; + return rv_ftob_s(bit_cast(value), 0, nullptr); + } + + static float to_float(uint16_t value) { + return bit_cast(rv_btof_s(value, 0, nullptr)); + } +}; + +template <> +struct type_ops { + static uint8_t generate() { + auto value = float(rand()) / RAND_MAX; + return rv_ftoe4m3_s(bit_cast(value), 0, nullptr); + } + + static float to_float(uint8_t value) { + return bit_cast(rv_e4m3tof_s(value, 0, nullptr)); + } +}; + +template <> +struct type_ops { + static uint8_t generate() { + auto value = float(rand()) / RAND_MAX; + return rv_ftoe5m2_s(bit_cast(value), 0, nullptr); + } + + static float to_float(uint8_t value) { + return bit_cast(rv_e5m2tof_s(value, 0, nullptr)); + } +}; + // CPU reference matmul using pruned (zero-padded) A static void matmul_cpu(otype_t *C, const itype_t *A_pruned, const itype_t *B, uint32_t M, uint32_t N, uint32_t K) { @@ -69,8 +131,8 @@ static void matmul_cpu(otype_t *C, const itype_t *A_pruned, const itype_t *B, for (uint32_t k = 0; k < K; ++k) { auto a = A_pruned[m * K + k]; auto b = B[k * N + n]; - auto fa = bit_cast(rv_htof_s(a, 0, nullptr)); - auto fb = bit_cast(rv_htof_s(b, 0, nullptr)); + auto fa = type_ops::to_float(a); + auto fb = type_ops::to_float(b); sum += fa * fb; } C[m * N + n] = sum; @@ -289,12 +351,10 @@ int main(int argc, char *argv[]) { std::vector h_A_full(sizeA_full); std::vector h_B(sizeB); for (uint32_t i = 0; i < sizeA_full; ++i) { - auto fv = float(rand()) / RAND_MAX; - h_A_full[i] = rv_ftoh_s(bit_cast(fv), 0, nullptr); + h_A_full[i] = type_ops::generate(); } for (uint32_t i = 0; i < sizeB; ++i) { - auto fv = float(rand()) / RAND_MAX; - h_B[i] = rv_ftoh_s(bit_cast(fv), 0, nullptr); + h_B[i] = type_ops::generate(); } // prune A in-place (2:4) then compress @@ -380,12 +440,20 @@ int main(int argc, char *argv[]) { int errors = 0; for (uint32_t i = 0; i < sizeC; ++i) { - union fi_t { float f; int32_t i; }; - fi_t fa, fb; - fa.f = h_C[i]; - fb.f = h_ref[i]; - auto d = std::abs(fa.i - fb.i); - if (d > FLOAT_ULP) { + bool match; + if constexpr (std::is_same::value || + std::is_same::value) { + match = (h_C[i] == 0.0f && h_ref[i] == 0.0f) + || (std::abs((h_C[i] - h_ref[i]) / h_ref[i]) < 0.01f); + } else { + union fi_t { float f; int32_t i; }; + fi_t fa, fb; + fa.f = h_C[i]; + fb.f = h_ref[i]; + auto d = std::abs(fa.i - fb.i); + match = d <= FLOAT_ULP; + } + if (!match) { if (errors < MAX_ERRORS) { printf("*** error: [%u] expected=%f, actual=%f\n", i, h_ref[i], h_C[i]); }