Skip to content

feat(backend): HIP + Vulkan backend touches (#13 carve-out 3/N) #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
TheTom wants to merge 1 commit into from
Closed

feat(backend): HIP + Vulkan backend touches (#13 carve-out 3/N) #30

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 46 additions & 56 deletions crates/metaltile-codegen/src/hip/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,34 +2,23 @@
//! SPDX-License-Identifier: Apache-2.0
//! HIP / ROCm codegen backend (`AMD_BACKEND_SPEC.md`).
//!
//! There is **one kernel-source layer** for the C++-dialect GPU targets, not
//! a CUDA original with a HIP copy. At the kernel level, `__global__`,
//! HIP is a CUDA-portable C++ dialect: at the **kernel** level, `__global__`,
//! `blockIdx`/`threadIdx`/`blockDim`, `__shared__`, `__syncthreads`,
//! `__syncwarp`, `__shfl_*_sync`, atomics, and the precise single-precision
//! math intrinsics (`expf`, `exp2f`, `rsqrtf`, `__frcp_rn`, `fmaxf`/`fminf`,
//! `fmaf`) are a **shared vocabulary** that NVIDIA and AMD define
//! identically. The op-walker (hosted in [`crate::cuda`]) emits that shared
//! source once; each vendor then gets a thin dialect lowering:
//! `fmaf`) are **bit-identical** to the CUDA emitter's output. The only deltas
//! are:
//!
//! * **CUDA** — the identity (the shared spelling is already valid CUDA).
//! * **HIP** — [`to_hip_dialect`]: header includes (`cuda_fp16.h` →
//! `hip/hip_fp16.h`, `cuda_bf16.h` → `hip/hip_bf16.h`), the bf16 type name
//! (`__nv_bfloat16` → `__hip_bfloat16`), 64-bit shuffle-mask containers,
//! and the missing bf16 `__ldg` shim.
//! 1. Header includes — `cuda_fp16.h` / `cuda_bf16.h` → `hip/hip_fp16.h` /
//! `hip/hip_bf16.h`.
//! 2. The bf16 type name — `__nv_bfloat16` → `__hip_bfloat16`.
//! 3. The `MT_INF` literal — `__int_as_float` is in HIP, but reusing the same
//! expression keeps the preamble identical.
//!
//! Sharing the one op-walker (rather than forking it) keeps the invariants
//! that took 4164/4164 corpus passes to establish: a numerics fix lands once
//! and both vendors inherit it.
//!
//! TODO(follow-up, post-#274/#275): the shared layer still *lives in* the
//! `cuda` module and the lowering is textual, which makes CUDA read as the
//! first-class target. The structural fix is to extract the op-walker into a
//! vendor-neutral `gpu_cpp` emitter module that `CudaGenerator` and
//! `HipGenerator` both wrap with their dialect lowering (CUDA's being the
//! identity), with the lowering applied to the emitter's structured
//! preamble/body parts instead of post-processing a flat string. Deferred
//! here because it churns the hardware-validated CUDA emitter this stack is
//! built on; behavior-preserving naming/docs land now.
//! So this generator **does not fork the CUDA op-walker**. It delegates to
//! [`crate::cuda::CudaGenerator`] for the full kernel emission, then runs a
//! small textual transform (`cuda_to_hip`) over the result. That keeps the
//! invariants that took the CUDA backend 4164/4164 passes to establish.
//!
//! **Wavefront width:** the default profile is wave32 (RDNA1+; gfx10/11/12 —
//! including the RX 9070 XT / gfx1201). CDNA wave64 needs the broader 64-bit
Expand All @@ -45,15 +34,15 @@ use crate::{
cuda::CudaGenerator,
};

/// HIP code generator: the shared op-walker plus the HIP dialect lowering
/// ([`to_hip_dialect`]).
/// HIP code generator. Thin wrapper around [`CudaGenerator`] that post-
/// processes the emitted CUDA source into HIP-flavored C++.
#[derive(Debug, Clone)]
pub struct HipGenerator {
profile: TargetProfile,
/// The shared op-walker (hosted by [`CudaGenerator`] today — see the
/// module-level TODO). Constructed with a CUDA-target profile so it
/// emits the shared spellings the lowering recognises; the *outer*
/// `profile` is what callers see via the `CodegenBackend` trait.
/// We construct a CUDA-profiled inner generator (lane_width 32, shared
/// idioms) so the **op-walker** uses CUDA spellings the transform
/// recognises; the *outer* `profile` is what callers see via the
/// `CodegenBackend` trait.
inner: CudaGenerator,
}

Expand All @@ -65,14 +54,15 @@ impl HipGenerator {
pub fn new() -> Self { Self::with_profile(TargetProfile::hip()) }

/// Pin to a specific HIP profile. For [`TargetProfile::hip_wave64`]
/// (CDNA — MI200/MI300/MI350), we *also* give the shared op-walker a
/// lane_width=64 profile so its reduction tree sizes the shared
/// scratch to 64 warps and emits the right per-warp boundaries. The
/// dialect lowering then widens the shuffle masks to 64 bits.
/// (CDNA — MI200/MI300/MI350), we *also* swap the inner CUDA
/// generator's profile to a lane_width=64 CUDA variant so its
/// reduction tree sizes the shared scratch to 64 warps and emits the
/// right per-warp boundaries. The textual transform then widens the
/// shuffle masks to 64 bits.
///
/// The inner op-walker profile inherits `mma` and `lane_width` from
/// the HIP profile — that's how `SoftwareLocalC` reaches the CoopTile
/// op-walker.
/// The inner CudaGenerator profile inherits `mma` and `lane_width`
/// from the HIP profile — that's how `SoftwareLocalC` reaches the
/// CoopTile op-walker.
pub fn with_profile(profile: TargetProfile) -> Self {
debug_assert_eq!(profile.target, Target::Hip);
let mut inner_profile = TargetProfile::cuda();
Expand All @@ -86,9 +76,9 @@ impl HipGenerator {
Self { profile, inner: CudaGenerator::with_profile(inner_profile) }
}

/// Dynamic shared-memory bytes for a launch, forwarded to the shared
/// op-walker (the layout is identical — `__shared__` decls are pure C
/// and survive [`to_hip_dialect`] unchanged).
/// Dynamic shared-memory bytes for a launch, forwarded to the inner
/// CUDA generator (the layout is identical — `__shared__` decls are
/// pure C and survive `cuda_to_hip` unchanged).
pub fn shared_bytes(&self, kernel: &Kernel, block_x: u32) -> usize {
self.inner.shared_bytes(kernel, block_x)
}
Expand All @@ -100,10 +90,10 @@ impl CodegenBackend for HipGenerator {
fn profile(&self) -> &TargetProfile { &self.profile }

fn generate(&self, kernel: &Kernel) -> Result<String> {
let shared_src = self.inner.generate(kernel)?;
let src = to_hip_dialect(&shared_src);
let cuda_src = self.inner.generate(kernel)?;
let src = cuda_to_hip(&cuda_src);
// Wave64 (CDNA) — widen any 32-bit shuffle masks to 64-bit
// ALL-LANES-ACTIVE. The wave32 lowering already swapped
// ALL-LANES-ACTIVE. The wave32 transform already swapped
// `0xffffffffu` to `0xffffffffull` (the low-32 mask in a 64-bit
// container), but on wave64 the upper 32 lanes are real, so the
// mask needs all-1s.
Expand All @@ -115,15 +105,15 @@ impl CodegenBackend for HipGenerator {
}
}

/// The HIP dialect lowering over the shared kernel source. Surgical: only
/// touches the lines where the vendor dialects genuinely differ. The kernel
/// body itself is untouched because both dialects define it identically.
/// Textual CUDA → HIP transform. Surgical: only touches the lines that
/// genuinely differ between the dialects. The kernel body itself is
/// untouched because the CUDA syntax HIP recognises 1:1.
///
/// The rewrites are **conservative** — each only matches text the shared
/// op-walker actually produces, so unrelated code (e.g. user-supplied
/// The transforms are **conservative** — each only matches text the CUDA
/// emitter actually produces, so unrelated code (e.g. user-supplied
/// preambles in future inline kernels) is not silently rewritten.
pub fn to_hip_dialect(shared_src: &str) -> String {
let mut s = shared_src.to_string();
pub fn cuda_to_hip(cuda_src: &str) -> String {
let mut s = cuda_src.to_string();

// Header includes. hipRTC auto-includes `hip/hip_runtime.h`, so the
// `__global__`/`__shared__` keywords resolve without an explicit include
Expand Down Expand Up @@ -273,19 +263,19 @@ mod tests {
}

#[test]
fn hip_dialect_lowering_is_idempotent() {
// Running the lowering twice produces the same output (each rule
// only fires against the shared spellings, not the HIP ones). Use
// the reduction kernel: it emits `__shfl_down_sync` masks, the rule
// most prone to matching its own output (`0xffffffffull` contains
fn cuda_to_hip_is_idempotent() {
// Running the transform twice produces the same output (each rule
// only fires against the CUDA spellings, not the HIP ones). Use the
// reduction kernel: it emits `__shfl_down_sync` masks, the rule most
// prone to matching its own output (`0xffffffffull` contains
// `0xffffffffu`).
let src = HipGenerator::new().generate(&row_reduce_sum_ir()).unwrap();
assert!(src.contains("0xffffffffull"), "expected widened shuffle mask");
let twice = to_hip_dialect(&src);
let twice = cuda_to_hip(&src);
assert_eq!(src, twice);
// And the simple elementwise kernel stays covered.
let src = HipGenerator::new().generate(&vector_add_ir()).unwrap();
assert_eq!(src, to_hip_dialect(&src));
assert_eq!(src, cuda_to_hip(&src));
}

/// Build a simple reduction-mode kernel so the emitted source contains
Expand Down
63 changes: 0 additions & 63 deletions crates/metaltile-runtime/src/device/hip/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -543,69 +543,6 @@ impl HipDevice {
Ok(out)
}

/// Prepare once, launch `warmup + iters` times, and return per-launch GPU
/// times in µs measured with `hipEvent` pairs (CUDA `bench_kernel` analog).
///
/// No read-back: throughput is data-independent, and skipping the DtoH
/// copy keeps the timed region pure kernel execution. Outputs stay
/// resident (overwritten each iter — fine, we only time).
pub fn bench_kernel(
&self,
kernel: &Kernel,
buffers: &BTreeMap<String, Vec<u8>>,
grid: [u32; 3],
block: [u32; 3],
warmup: u32,
iters: u32,
) -> Result<Vec<f64>, MetalTileError> {
let prep = self.prepare(kernel, buffers, block)?;
let mut args = prep.args();

// Warmup — first launch pays hipRTC/cache/clock-ramp costs.
for _ in 0..warmup {
self.launch(prep.func, grid, block, prep.shared_bytes, &mut args)?;
}

// Two events bracket each timed launch on the null stream — the
// same stream every `launch` rides.
let (mut start, mut stop): (hipEvent_t, hipEvent_t) = (ptr::null_mut(), ptr::null_mut());
hip_check(unsafe { hipEventCreate(&mut start) }, "hipEventCreate(start)")?;
hip_check(unsafe { hipEventCreate(&mut stop) }, "hipEventCreate(stop)")?;

// Closure owns its sample vec (returned on success) so no outer
// borrow outlives it — lets us unconditionally destroy events after.
let mut timed = || -> Result<Vec<f64>, MetalTileError> {
let mut samples = Vec::with_capacity(iters as usize);
for _ in 0..iters {
hip_check(
unsafe { hipEventRecord(start, ptr::null_mut()) },
"hipEventRecord(start)",
)?;
self.launch(prep.func, grid, block, prep.shared_bytes, &mut args)?;
hip_check(
unsafe { hipEventRecord(stop, ptr::null_mut()) },
"hipEventRecord(stop)",
)?;
hip_check(unsafe { hipEventSynchronize(stop) }, "hipEventSynchronize")?;
let mut ms: f32 = 0.0;
hip_check(
unsafe { hipEventElapsedTime(&mut ms, start, stop) },
"hipEventElapsedTime",
)?;
samples.push(ms as f64 * 1000.0); // ms → µs
}
Ok(samples)
};
let res = timed();

// Always destroy events, even on error.
unsafe {
hipEventDestroy(start);
hipEventDestroy(stop);
}
res
}

pub fn synchronize(&self) -> Result<(), MetalTileError> {
hip_check(unsafe { hipDeviceSynchronize() }, "hipDeviceSynchronize")
}
Expand Down
Loading
Loading