FEDP2K + WGMMA dse + opts

[NikhilRout]

Isolating fedp2k and wgmma changes from #376

---

1. Widened FEDP to 2K elements:

• Number of elements each FEDP processes is doubled when VX_CFG_TCU_FEDP2K is enabled
• This reduces number of uops in k-dimension by half for WGMMA instructions
• Achieves significant speedups (1.5-1.8x) over baseline because of pipelining (visit perf/dse/wgmma)
• Upper half of elements are hardwired to 0 for WMMA instructions when FEDP2K is enabled

---

2a. Extra desc uop vs TCU_LD for desc WGMMA DSE

• While feeding double the operands to FEDPs when FEDP2K is enabled was straightforward in SS mode, we weren't able to do the same in RS mode by only using rs1 for A operands
• To keep rs2 free to feed the upper half of FEDP A operands, we need to finish rs2 descriptor copy to x11 before starting compute uops
• There's 2 ways to do this:
  i) Add an additional desc uop before starting the main compute uops solely for rs1 -> x10 and rs2 -> x11 descriptor copy
  ii) Handle descriptor copy into TCU SRAM and use it from there directly (no x10/x11), similar to SP/MX metadata copy in the kernel via TCU_LD instruction + new load_wg_metadata intrinsic
• Prototyping both in SimX showed arch i) is 5-10% faster across all configs and modes than arch ii). visit perf/dse/wgmma/desc
• Experimental arch ii) prototype pushed to archival feature_wgld branch for ref
• Extra desc uop before compute uops only adds +/- 1% overhead compared to baseline where desc copy + 1st compute uop were both happening in the first uop

2b. WGMMA + Sparsity/MX support

• Added sgemm_tcu_wg_mx test
• Modified sub-FEDP MX scale factor block indexing for FEDP2K enabled
• Worked on sgemm_tcu_wg_sp test w/wo FEDP2K in RS/SS mode
• Enabling FEDP2K now results in RS > SS for sparse fp16 wgmma kernels, aligns with hopper ubenchmarks! (https://arxiv.org/pdf/2501.12084). visit perf/tcu/dse/sp_rs_ss

---

Notes:

• Includes minor MX ci bug fixes + scale factor application RTL opt
• Updated tensor_wg ci with fedp2k and sgemm_tcu_wg_mx + removed duplicate tests

[tinebp]

The master branch now has the FPGA fixed.
I recommend that you run and pass the following synthesis
dut/top + NT+NW=16 + TFR + all formats + sparsity+wgmma + dxa + 1k @ 250 MHz
dut/top + NT+NW=16 + TFR + all formats + sparsity+wgmma + dxa + 2k @ 250 MHz

[tinebp]

tier: fast doesn't exist

[tinebp]

looks too expensive
reg [NUM_EX_UNITS-1:0][PER_ISSUE_WARPS-1:0] fu_locked;
either
reg [NUM_EX_UNITS-1:0] fu_locked;
or
reg [PER_ISSUE_WARPS-1:0] fu_locked;

the end goal is to lock a warp to a given FU or lock a FU simply
efficienty is very important insice the scoreboard - this is one of the most important module in the pipeline - if you can make it better in term of efficiency - it is a huge win

[NikhilRout]

I'm going to delete the MXINT8 datapath as well. It's on the critical path and should've actually been accumulating in FP32 anyway. I'll add it back alongside MXINT4, NVINT4 changes in a future PR

[NikhilRout]

passes NW=NT=16 fpga synthesis now with all formats enabled + mx + sp + wg + dxa @ 250 MHz
however, it (understandably) fails with -0.143 WNS when FEDP2K is enabled

[tinebp]

since scorebaord and TCU TFR changes are ortogonal to WGMMA, I strongly recommend pushing them a separate PRs.

[tinebp]

Are you fixing a genuine bug in the scoreboard, or is this an optimization?
I prefer that you make this a separate pull request, bug or not so that it is fully isolated

[NikhilRout]

its an optimization that's tied into the wgmma optimization. moving this to another PR will break the wgmma uop expansion here

previously, fu_locked was a warp mask shared by all execution units
so even if a warp "X" only locked TCU for wgmma uop expansion, only warp X could pass through the scoreboard
this unnecessarily blocked any other warps (Y, Z ...) from issuing unrelated ALU, FPU, LSU instructions

the new change now allows fu_locked to track locking per execution unit instead
so now, warp X's first wgmma uop only locks TCU
warp Y (say doing LSU load) is allowed
warp Z (doing another wgmma) remains blocked

so now FU blocked = selected FU locked && trying to get that FU lock

10 -> first uop, get lock
00 -> intermediate uop
01 -> final uop, release lock
11 -> single uop instruction

[NikhilRout]

TFR and sw/ folder changes are just deleting MXINT8 datapath (for timing + shouldve been acc in FP32 anyway)
will add MXINT8 back alongside MXINT4 and NVINT4 in a future PR

and there's just one tiny MXFP8 scale factor application parallelization in vx_tcu_tfr_mul_f8

[tinebp]

Do not modify the coverage of an existing baseline; we need to have perfect args and config match to catch regression.
You are free to expand the perf gate coverage with new tests

[NikhilRout]

this test was defaulting to RS mode previously (from sgemm_tcu_wg makefile)
so we were essentially never testing SS mode in perf_gate

[NikhilRout]

and I also changed 128^3 to 64^3 cause large rtlsim runs would just crash verilator on my local setup
so I couldnt even generate the SS baseline with 128^3

but i guess if we really want large dims for this I can fail this particular perf_gate test once in CI on purpose and then get the value from there and then manually edit baseline instead of through pytest?

[NikhilRout]

also NW=NT=16 TCU top dut synthesis with

TFR + all fmts + wg + sp + mx + dxa passes @ 250 MHz with +0.09 WNS
TFR + all fmts + wg + sp + mx + dxa + fedp2k fails @ 250 MHz with -0.051 WNS