Skip to content

[codex] Unify GPU renderer FMA detection #72

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

Merged
adpena merged 1 commit into from
Jul 6, 2026
Merged

[codex] Unify GPU renderer FMA detection #72

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
50 changes: 10 additions & 40 deletions runtime/molt-gpu/src/render/cuda.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,9 @@ use crate::ops::PrimitiveOp;
use crate::render::indexing::{
IndexDialect, render_reduction_input_index, render_shapetracker_index, zero_literal_for_dtype,
};
use crate::render::{BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer};
use crate::render::{
BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer, detect_fma_pattern,
};

/// CUDA C renderer for all 26 primitive ops.
pub struct CudaRenderer;
Expand Down Expand Up @@ -170,45 +172,13 @@ impl CudaRenderer {
kernel: &FusedKernel,
idx_var: &str,
) -> Option<(String, String, String)> {
if op.op() != PrimitiveOp::Add {
return None;
}
if !op.dst_dtype().is_float() {
return None;
}

for (mul_src_pos, add_src_pos) in [(0, 1), (1, 0)] {
if let FusedSrc::Op(prior_idx) = &op.srcs()[mul_src_pos]
&& *prior_idx < op_idx
{
let prior_op = &kernel.ops[*prior_idx];
if prior_op.op() == PrimitiveOp::Mul {
let a = match &prior_op.srcs()[0] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
let b = match &prior_op.srcs()[1] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
let c = match &op.srcs()[add_src_pos] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
return Some((a, b, c));
}
}
}
None
let pattern = detect_fma_pattern(op, op_idx, kernel, op.dst_dtype().is_float())?;
let prior_op = &kernel.ops[pattern.mul_op_idx];
Some((
Self::render_src(&prior_op.srcs()[0], kernel, idx_var),
Self::render_src(&prior_op.srcs()[1], kernel, idx_var),
Self::render_src(&op.srcs()[pattern.add_src_pos], kernel, idx_var),
))
}
}

Expand Down
50 changes: 10 additions & 40 deletions runtime/molt-gpu/src/render/hip.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@ use crate::ops::PrimitiveOp;
use crate::render::indexing::{
IndexDialect, render_reduction_input_index, render_shapetracker_index, zero_literal_for_dtype,
};
use crate::render::{BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer};
use crate::render::{
BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer, detect_fma_pattern,
};

/// HIP C renderer for all 26 primitive ops.
pub struct HipRenderer;
Expand Down Expand Up @@ -170,45 +172,13 @@ impl HipRenderer {
kernel: &FusedKernel,
idx_var: &str,
) -> Option<(String, String, String)> {
if op.op() != PrimitiveOp::Add {
return None;
}
if !op.dst_dtype().is_float() {
return None;
}

for (mul_src_pos, add_src_pos) in [(0, 1), (1, 0)] {
if let FusedSrc::Op(prior_idx) = &op.srcs()[mul_src_pos]
&& *prior_idx < op_idx
{
let prior_op = &kernel.ops[*prior_idx];
if prior_op.op() == PrimitiveOp::Mul {
let a = match &prior_op.srcs()[0] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
let b = match &prior_op.srcs()[1] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
let c = match &op.srcs()[add_src_pos] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
return Some((a, b, c));
}
}
}
None
let pattern = detect_fma_pattern(op, op_idx, kernel, op.dst_dtype().is_float())?;
let prior_op = &kernel.ops[pattern.mul_op_idx];
Some((
Self::render_src(&prior_op.srcs()[0], kernel, idx_var),
Self::render_src(&prior_op.srcs()[1], kernel, idx_var),
Self::render_src(&op.srcs()[pattern.add_src_pos], kernel, idx_var),
))
}
}

Expand Down
36 changes: 36 additions & 0 deletions runtime/molt-gpu/src/render/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,42 @@ use crate::shapetracker::ShapeTracker;
mod fused_op;
pub use fused_op::{FusedOp, FusedOpDomain, ReductionDomain};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct FmaPattern {
pub mul_op_idx: usize,
pub add_src_pos: usize,
}

pub(crate) fn detect_fma_pattern(
op: &FusedOp,
op_idx: usize,
kernel: &FusedKernel,
dst_is_float: bool,
) -> Option<FmaPattern> {
if op.op() != PrimitiveOp::Add || !dst_is_float {
return None;
}

for (mul_src_pos, add_src_pos) in [(0, 1), (1, 0)] {
let FusedSrc::Op(prior_idx) = &op.srcs()[mul_src_pos] else {
continue;
};
if *prior_idx < op_idx
&& kernel
.ops
.get(*prior_idx)
.is_some_and(|prior_op| prior_op.op() == PrimitiveOp::Mul)
{
return Some(FmaPattern {
mul_op_idx: *prior_idx,
add_src_pos,
});
}
}

None
}

/// Executable body carried by a scheduled kernel.
///
/// Compute kernels evaluate the ordered [`FusedOp`] chain. Materialization
Expand Down
53 changes: 11 additions & 42 deletions runtime/molt-gpu/src/render/msl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,9 @@ use crate::ops::PrimitiveOp;
use crate::render::indexing::{
IndexDialect, render_reduction_input_index, render_shapetracker_index, zero_literal_for_dtype,
};
use crate::render::{BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer};
use crate::render::{
BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer, detect_fma_pattern,
};

/// Metal Shading Language renderer for all 26 primitive ops.
pub struct MslRenderer;
Expand Down Expand Up @@ -138,47 +140,14 @@ impl MslRenderer {
kernel: &FusedKernel,
idx_var: &str,
) -> Option<(String, String, String)> {
if op.op() != PrimitiveOp::Add {
return None;
}
// Only emit FMA for float types
if !op.dst_dtype().narrow_metal().is_float() {
return None;
}

// Check if either source is a MUL op result
for (mul_src_pos, add_src_pos) in [(0, 1), (1, 0)] {
if let FusedSrc::Op(prior_idx) = &op.srcs()[mul_src_pos]
&& *prior_idx < op_idx
{
let prior_op = &kernel.ops[*prior_idx];
if prior_op.op() == PrimitiveOp::Mul {
let a = match &prior_op.srcs()[0] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
let b = match &prior_op.srcs()[1] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
let c = match &op.srcs()[add_src_pos] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
return Some((a, b, c));
}
}
}
None
let pattern =
detect_fma_pattern(op, op_idx, kernel, op.dst_dtype().narrow_metal().is_float())?;
let prior_op = &kernel.ops[pattern.mul_op_idx];
Some((
Self::render_src(&prior_op.srcs()[0], kernel, idx_var),
Self::render_src(&prior_op.srcs()[1], kernel, idx_var),
Self::render_src(&op.srcs()[pattern.add_src_pos], kernel, idx_var),
))
}
}

Expand Down
55 changes: 15 additions & 40 deletions runtime/molt-gpu/src/render/opencl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ use crate::ops::PrimitiveOp;
use crate::render::indexing::{
IndexDialect, render_reduction_input_index, render_shapetracker_index, zero_literal_for_dtype,
};
use crate::render::{BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer};
use crate::render::{
BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer, detect_fma_pattern,
};

/// OpenCL C renderer for all 26 primitive ops.
pub struct OpenClRenderer {
Expand Down Expand Up @@ -200,45 +202,18 @@ impl OpenClRenderer {
kernel: &FusedKernel,
idx_var: &str,
) -> Option<(String, String, String)> {
if op.op() != PrimitiveOp::Add {
return None;
}
if !op.dst_dtype().narrow_opencl(self.has_fp64).is_float() {
return None;
}

for (mul_src_pos, add_src_pos) in [(0, 1), (1, 0)] {
if let FusedSrc::Op(prior_idx) = &op.srcs()[mul_src_pos]
&& *prior_idx < op_idx
{
let prior_op = &kernel.ops[*prior_idx];
if prior_op.op() == PrimitiveOp::Mul {
let a = match &prior_op.srcs()[0] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => self.format_const(*val, *dtype),
};
let b = match &prior_op.srcs()[1] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => self.format_const(*val, *dtype),
};
let c = match &op.srcs()[add_src_pos] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => self.format_const(*val, *dtype),
};
return Some((a, b, c));
}
}
}
None
let pattern = detect_fma_pattern(
op,
op_idx,
kernel,
op.dst_dtype().narrow_opencl(self.has_fp64).is_float(),
)?;
let prior_op = &kernel.ops[pattern.mul_op_idx];
Some((
self.render_src(&prior_op.srcs()[0], kernel, idx_var),
self.render_src(&prior_op.srcs()[1], kernel, idx_var),
self.render_src(&op.srcs()[pattern.add_src_pos], kernel, idx_var),
))
}

/// Check whether any buffer in the kernel uses Float64 (pre-narrowing).
Expand Down
55 changes: 15 additions & 40 deletions runtime/molt-gpu/src/render/wgsl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,9 @@ use crate::ops::PrimitiveOp;
use crate::render::indexing::{
IndexDialect, render_reduction_input_index, render_shapetracker_index, zero_literal_for_dtype,
};
use crate::render::{BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer};
use crate::render::{
BufferAccess, FusedKernel, FusedOp, FusedSrc, KernelBody, Renderer, detect_fma_pattern,
};

/// Configuration for the WGSL renderer.
#[derive(Debug, Clone, Default)]
Expand Down Expand Up @@ -172,45 +174,18 @@ impl WgslRenderer {
kernel: &FusedKernel,
idx_var: &str,
) -> Option<(String, String, String)> {
if op.op() != PrimitiveOp::Add {
return None;
}
if !op.dst_dtype().narrow_webgpu().is_float() {
return None;
}

for (mul_src_pos, add_src_pos) in [(0, 1), (1, 0)] {
if let FusedSrc::Op(prior_idx) = &op.srcs()[mul_src_pos]
&& *prior_idx < op_idx
{
let prior_op = &kernel.ops[*prior_idx];
if prior_op.op() == PrimitiveOp::Mul {
let a = match &prior_op.srcs()[0] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
let b = match &prior_op.srcs()[1] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
let c = match &op.srcs()[add_src_pos] {
FusedSrc::Buf(buf_idx) => {
Self::render_buf_read(*buf_idx, &kernel.bufs[*buf_idx], idx_var)
}
FusedSrc::Op(p) => format!("v{}", p),
FusedSrc::Const { val, dtype } => Self::format_const(*val, *dtype),
};
return Some((a, b, c));
}
}
}
None
let pattern = detect_fma_pattern(
op,
op_idx,
kernel,
op.dst_dtype().narrow_webgpu().is_float(),
)?;
let prior_op = &kernel.ops[pattern.mul_op_idx];
Some((
Self::render_src(&prior_op.srcs()[0], kernel, idx_var),
Self::render_src(&prior_op.srcs()[1], kernel, idx_var),
Self::render_src(&op.srcs()[pattern.add_src_pos], kernel, idx_var),
))
}
}

Expand Down
Loading