diff --git a/runtime/molt-gpu/src/render/cuda.rs b/runtime/molt-gpu/src/render/cuda.rs index b69b15c2f..5c5010428 100644 --- a/runtime/molt-gpu/src/render/cuda.rs +++ b/runtime/molt-gpu/src/render/cuda.rs @@ -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; @@ -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), + )) } } diff --git a/runtime/molt-gpu/src/render/hip.rs b/runtime/molt-gpu/src/render/hip.rs index f92e84ccb..5b1a23529 100644 --- a/runtime/molt-gpu/src/render/hip.rs +++ b/runtime/molt-gpu/src/render/hip.rs @@ -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; @@ -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), + )) } } diff --git a/runtime/molt-gpu/src/render/mod.rs b/runtime/molt-gpu/src/render/mod.rs index 9f4ad5f44..88c7e8555 100644 --- a/runtime/molt-gpu/src/render/mod.rs +++ b/runtime/molt-gpu/src/render/mod.rs @@ -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 { + 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 diff --git a/runtime/molt-gpu/src/render/msl.rs b/runtime/molt-gpu/src/render/msl.rs index f13a277f3..6fe174c7e 100644 --- a/runtime/molt-gpu/src/render/msl.rs +++ b/runtime/molt-gpu/src/render/msl.rs @@ -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; @@ -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), + )) } } diff --git a/runtime/molt-gpu/src/render/opencl.rs b/runtime/molt-gpu/src/render/opencl.rs index c08aa1c87..3c556a987 100644 --- a/runtime/molt-gpu/src/render/opencl.rs +++ b/runtime/molt-gpu/src/render/opencl.rs @@ -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 { @@ -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). diff --git a/runtime/molt-gpu/src/render/wgsl.rs b/runtime/molt-gpu/src/render/wgsl.rs index db92b1aa6..c58e72683 100644 --- a/runtime/molt-gpu/src/render/wgsl.rs +++ b/runtime/molt-gpu/src/render/wgsl.rs @@ -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)] @@ -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), + )) } }