Skip to content

Re-enabled Search #199

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
jafioti merged 6 commits into from
Jan 4, 2026
Merged

Re-enabled Search #199

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
fbf3e81
search
jafioti Jan 2, 2026
62170cb
cstruct name change
jafioti Jan 4, 2026
6cb5e39
re-enabled timings
jafioti Jan 4, 2026
eeec8d4
fixed cuda test
jafioti Jan 4, 2026
2bcc5d5
symbolic test disables
jafioti Jan 4, 2026
d6321f3
changed cuda constants
jafioti Jan 4, 2026
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
4 changes: 3 additions & 1 deletion crates/luminal_cuda/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ license = "MIT OR Apache-2.0"
[dependencies]
luminal = { path = "../.." }
cudarc = {version="0.17.3", features=["cuda-12080"]}
rustc-hash = "2.1.1"
as-any = "0.3.2"
itertools = "0.12.1"
fixedbitset = "0.5.7"
Expand All @@ -19,3 +18,6 @@ tracing = "0.1.43"
tracing-perfetto-sdk-schema = "0.13.0"
prost = "0.14.1"
half = "2.7.1"
pretty-duration = "0.1.1"
bytemuck = "1.24.0"
memmap2 = "0.9.9"
3 changes: 1 addition & 2 deletions crates/luminal_cuda/src/block/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,7 @@ mod ops;
pub use ops::*;

use cudarc::driver::CudaStream;
use luminal::{shape::Expression, utils::EgglogOp};
use rustc_hash::FxHashMap;
use luminal::{prelude::FxHashMap, shape::Expression, utils::EgglogOp};
use std::fmt::Debug;

use crate::runtime::CustomState;
Expand Down
130 changes: 118 additions & 12 deletions crates/luminal_cuda/src/block/ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,13 @@ use cudarc::driver::{CudaStream, DevicePtr};
use itertools::Itertools;
use luminal::{
graph::{extract_expr, extract_expr_list},
prelude::ENodeId,
prelude::*,
serialized_egraph::SerializedEGraph,
shape::Expression,
utils::{
flatten_mul_strides, CStructBuilder, EgglogOp, LLIROp,
flatten_mul_strides, EgglogOp, LLIROp,
OpParam::{self, *},
},
};
use rustc_hash::FxHashMap;

use crate::block::BlockOp;

Expand Down Expand Up @@ -94,11 +92,15 @@ impl EgglogOp for RowAdd {

impl BlockOp for RowAdd {
fn launch_range(&self) -> Vec<Expression> {
self.range.clone()
if self.range.is_empty() {
vec![1.into()]
} else {
self.range.clone()
}
}

fn output_size(&self) -> Expression {
self.range.iter().copied().product::<Expression>() * self.row_width
self.range.iter().copied().product::<Expression>().max(1) * self.row_width
}

fn consumer_barriers_seperate(&self) -> Vec<Vec<bool>> {
Expand Down Expand Up @@ -128,7 +130,7 @@ impl BlockOp for RowAdd {
_: &CudaStream,
expressions: &FxHashMap<Expression, i32>,
) -> Vec<u8> {
CStructBuilder::new()
CStruct::new()
.int(expressions[&flatten_mul_strides(&self.range, &self.a_stride)])
.int(expressions[&flatten_mul_strides(&self.range, &self.b_stride)])
.int(expressions[&flatten_mul_strides(&self.range, &self.out_stride)])
Expand Down Expand Up @@ -270,7 +272,7 @@ impl BlockOp for RowSwishMul {
_: &CudaStream,
expressions: &FxHashMap<Expression, i32>,
) -> Vec<u8> {
CStructBuilder::new()
CStruct::new()
.int(expressions[&flatten_mul_strides(&self.range, &self.a_stride)])
.int(expressions[&flatten_mul_strides(&self.range, &self.b_stride)])
.int(expressions[&flatten_mul_strides(&self.range, &self.a_stride)])
Expand Down Expand Up @@ -491,7 +493,7 @@ impl BlockOp for RowRMSNorm {
_: &CudaStream,
expressions: &FxHashMap<Expression, i32>,
) -> Vec<u8> {
CStructBuilder::new()
CStruct::new()
.int(expressions[&flatten_mul_strides(&self.range, &self.a_stride)])
.int(expressions[&flatten_mul_strides(&self.range, &self.a_stride)])
.int(expressions[&self.row_width])
Expand Down Expand Up @@ -623,7 +625,7 @@ impl BlockOp for RowRope {
_: &CudaStream,
expressions: &FxHashMap<Expression, i32>,
) -> Vec<u8> {
CStructBuilder::new()
CStruct::new()
.int(expressions[&flatten_mul_strides(&self.range, &self.a_stride)])
.int(expressions[&flatten_mul_strides(&self.range, &self.a_stride)])
.int(expressions[&self.row_width])
Expand Down Expand Up @@ -938,7 +940,7 @@ impl BlockOp for TileMatmul {
m_pos_stride[self.range.len() - 2] = 1.into();
let mut n_pos_stride = vec![0.into(); self.range.len()];
n_pos_stride[self.range.len() - 1] = 1.into();
CStructBuilder::new()
CStruct::new()
.ints(
&self
.untiled_range
Expand Down Expand Up @@ -1277,7 +1279,7 @@ impl BlockOp for GQAAttention {
group_pos_stride[self.range.len() - 2] = 1.into();
let mut head_pos_stride = vec![0.into(); self.range.len()];
head_pos_stride[self.range.len() - 3] = 1.into();
CStructBuilder::new()
CStruct::new()
.int(expressions[&self.head_dim])
.int(expressions[&self.cur_seq])
.int(expressions[&self.kv_row_stride])
Expand Down Expand Up @@ -1316,3 +1318,107 @@ impl BlockOp for GQAAttention {
]
}
}

#[derive(Debug)]
pub struct CStruct {
buf: Vec<u8>,
max_align: usize,
}

impl Default for CStruct {
fn default() -> Self {
Self {
buf: Vec::new(),
max_align: 1,
}
}
}

impl CStruct {
pub fn new() -> Self {
Self::default()
}

fn align_to(&mut self, align: usize) {
self.max_align = self.max_align.max(align);

let len = self.buf.len();
let rem = len % align;
if rem != 0 {
let pad = align - rem;
self.buf.extend(std::iter::repeat_n(0u8, pad));
}
}

pub fn int(mut self, v: i32) -> Self {
self.align_to(4);
self.buf.extend_from_slice(&v.to_ne_bytes());
self
}

pub fn ints(mut self, vs: &[i32]) -> Self {
self.align_to(4);
for &v in vs {
self.buf.extend_from_slice(&v.to_ne_bytes());
}
self
}

pub fn float(mut self, v: f32) -> Self {
self.align_to(4);
self.buf.extend_from_slice(&v.to_ne_bytes());
self
}

pub fn floats(mut self, vs: &[f32]) -> Self {
self.align_to(4);
for &v in vs {
self.buf.extend_from_slice(&v.to_ne_bytes());
}
self
}

pub fn bool(mut self, v: bool) -> Self {
self.align_to(1);
self.buf.push(if v { 1 } else { 0 });
self
}

pub fn ptr_const_f32(mut self, p: *const f32) -> Self {
let ptr_size = std::mem::size_of::<usize>(); // usually 8
let ptr_align = ptr_size;
self.align_to(ptr_align);

let addr = p as usize;
let bytes = addr.to_ne_bytes();

self.buf.extend_from_slice(&bytes[..ptr_size]);
self
}

pub fn ptr_mut_f32(self, p: *mut f32) -> Self {
self.ptr_const_f32(p as *const f32)
}

/// Pad the struct size to a multiple of max_align.
pub fn finish_struct(mut self) -> Vec<u8> {
let align = self.max_align;
if align > 1 {
let len = self.buf.len();
let rem = len % align;
if rem != 0 {
let pad = align - rem;
self.buf.extend(std::iter::repeat_n(0u8, pad));
}
}
self.buf
}

/// Insert a raw byte field (e.g., another struct).
/// `align` must be the alignment of the nested struct.
pub fn bytes(mut self, align: usize, data: &[u8]) -> Self {
self.align_to(align);
self.buf.extend_from_slice(data);
self
}
}
6 changes: 3 additions & 3 deletions crates/luminal_cuda/src/kernel/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,8 @@

use std::sync::Arc;

use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream};
use luminal::{shape::Expression, utils::EgglogOp};
use rustc_hash::FxHashMap;
use cudarc::driver::{CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream};
use luminal::prelude::*;

pub mod ops;
pub use ops::Ops;
Expand All @@ -16,6 +15,7 @@ pub trait KernelOp: EgglogOp {
stream: &Arc<CudaStream>,
) -> (
CudaFunction,
Arc<CudaModule>,
String,
(Expression, Expression, Expression),
(Expression, Expression, Expression),
Expand Down
18 changes: 11 additions & 7 deletions crates/luminal_cuda/src/kernel/ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,24 +1,20 @@
use std::sync::Arc;

use crate::{cuda_dtype, kernel::KernelOp};
use cudarc::{
driver::{CudaContext, CudaFunction, CudaSlice, CudaStream},
driver::{CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream},
nvrtc::{compile_ptx, CompileOptions},
};
use itertools::Itertools;
use luminal::{
graph::{extract_dtype, extract_expr, extract_expr_list},
op::DType,
prelude::ENodeId,
prelude::*,
serialized_egraph::SerializedEGraph,
shape::Expression,
utils::{
flatten_mul_strides, EgglogOp, LLIROp,
OpParam::{self, *},
},
};
use rustc_hash::{FxHashMap, FxHashSet};

use crate::{cuda_dtype, kernel::KernelOp};

pub type Ops = (KernelAdd, KernelMul, KernelIota, KernelGather);

Expand Down Expand Up @@ -85,6 +81,7 @@ impl KernelOp for KernelAdd {
stream: &Arc<CudaStream>,
) -> (
CudaFunction,
Arc<CudaModule>,
String,
(Expression, Expression, Expression),
(Expression, Expression, Expression),
Expand Down Expand Up @@ -125,6 +122,7 @@ extern \"C\" {{
.collect();
(
func,
module,
kernel,
(
self.out_shape.iter().copied().product::<Expression>(),
Expand Down Expand Up @@ -205,6 +203,7 @@ impl KernelOp for KernelMul {
stream: &Arc<CudaStream>,
) -> (
CudaFunction,
Arc<CudaModule>,
String,
(Expression, Expression, Expression),
(Expression, Expression, Expression),
Expand Down Expand Up @@ -245,6 +244,7 @@ extern \"C\" {{
.collect();
(
func,
module,
kernel,
(
self.out_shape.iter().copied().product::<Expression>(),
Expand Down Expand Up @@ -328,6 +328,7 @@ impl KernelOp for KernelGather {
stream: &Arc<CudaStream>,
) -> (
CudaFunction,
Arc<CudaModule>,
String,
(Expression, Expression, Expression),
(Expression, Expression, Expression),
Expand Down Expand Up @@ -370,6 +371,7 @@ extern \"C\" {{
.collect();
(
func,
module,
kernel,
(self.out_shape.iter().copied().product(), 1.into(), 1.into()),
(1.into(), 1.into(), 1.into()),
Expand Down Expand Up @@ -438,6 +440,7 @@ impl KernelOp for KernelIota {
stream: &Arc<CudaStream>,
) -> (
CudaFunction,
Arc<CudaModule>,
String,
(Expression, Expression, Expression),
(Expression, Expression, Expression),
Expand Down Expand Up @@ -468,6 +471,7 @@ extern \"C\" {{
.collect();
(
func,
module,
kernel,
(self.range, 1.into(), 1.into()),
(1.into(), 1.into(), 1.into()),
Expand Down
Loading