plumb per-build wave-size flags from CMake into codegen (#6103)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/3008

Add the CMake-to-codegen plumbing that a later change uses to emit only the wave-size-specific kernel instantiations and host dispatch tables a given ROCm wheel needs. `cmake/Hip.cmake` derives `FBGEMM_HAS_WAVE32` / `FBGEMM_HAS_WAVE64` from `PYTORCH_ROCM_ARCH` (gfx9* is wave64; gfx10/11/12 is wave32), `CMakeLists.txt` forwards them to the gen scripts as `--has_wave32` / `--has_wave64`, and `scripts_argsparse.py` parses them.

This change is pure plumbing: the flags are parsed but not yet consumed by the Jinja templates, so generated code is byte-identical to before on every build configuration. The consumer lands in the next PR in the chain. Splitting the plumbing out lets it be reviewed and landed without touching any generated kernel output.

Fourth in the chain splitting #5804 into reviewable pieces; stacked on #6037.

Authored with assistance from Claude (Anthropic).

Pull Request resolved: #6103

Reviewed By: cthi

Differential Revision: D114618233

Pulled By: q10

fbshipit-source-id: 306b374d49e7dcfd0b3b3513520001d2d5afe1dc

Author: jeffdaily

fbgemm_gpu/CMakeLists.txt

@@ -182,8 +182,15 @@ endif()
 set(CMAKE_CODEGEN_DIR ${CMAKE_CURRENT_SOURCE_DIR}/codegen)
 
 macro(RUN_GEN_SCRIPT SCRIPT)
+  set(rocm_flag)
   if(FBGEMM_BUILD_VARIANT STREQUAL BUILD_VARIANT_ROCM)
-    set(rocm_flag --is_rocm)
+    list(APPEND rocm_flag --is_rocm)
+    if(FBGEMM_HAS_WAVE32)
+      list(APPEND rocm_flag --has_wave32)
+    endif()
+    if(FBGEMM_HAS_WAVE64)
+      list(APPEND rocm_flag --has_wave64)
+    endif()
   endif()
 
   BLOCK_PRINT(

fbgemm_gpu/cmake/Hip.cmake

@@ -34,6 +34,22 @@ if(PYTORCH_ROCM_ARCH STREQUAL "")
 endif()
 message("Building FBGEMM for GPU arch: ${PYTORCH_ROCM_ARCH}")
 
+# Derive the wave-size set in scope for this build from PYTORCH_ROCM_ARCH.
+# gfx9* archs (CDNA: gfx90a, gfx940/941/942, gfx950) run wave64; gfx10/11/12
+# archs (RDNA3/4: gfx1100, gfx1200, ...) run wave32. The codegen uses these
+# to emit only the host-side dispatch branches needed by this wheel, so
+# single-arch wheels stay free of the wrong wave's bracket table.
+set(FBGEMM_HAS_WAVE32 OFF)
+set(FBGEMM_HAS_WAVE64 OFF)
+foreach(fbgemm_rocm_arch ${PYTORCH_ROCM_ARCH})
+  if(fbgemm_rocm_arch MATCHES "^gfx9")
+    set(FBGEMM_HAS_WAVE64 ON)
+  else()
+    set(FBGEMM_HAS_WAVE32 ON)
+  endif()
+endforeach()
+message("FBGEMM wave-size set: WAVE32=${FBGEMM_HAS_WAVE32} WAVE64=${FBGEMM_HAS_WAVE64}")
+
 ADD_DEFINITIONS(-DNDEBUG)
 # USE_ROCM flag is used inside FBGEMM_GPU C++ code
 ADD_DEFINITIONS(-DUSE_ROCM)

fbgemm_gpu/codegen/genscript/scripts_argsparse.py

@@ -21,6 +21,11 @@
 )
 parser.add_argument("--opensource", action="store_false", dest="is_fbcode")
 parser.add_argument("--is_rocm", action="store_true")
+# CMake-derived wave-size set for the current ROCm build. CUDA builds ignore
+# these (always wave32). Both unset on a ROCm build defaults to wave64-only
+# to preserve pre-port behavior.
+parser.add_argument("--has_wave32", action="store_true")
+parser.add_argument("--has_wave64", action="store_true")
 
 args: argparse.Namespace
 _: list[str]