Skip to content

Commit

Permalink
Generalize GPU buffer create/read/update. Part 2
Browse files Browse the repository at this point in the history
  • Loading branch information
dfellis committed Nov 8, 2024
1 parent d4d90bc commit 99222b1
Show file tree
Hide file tree
Showing 3 changed files with 139 additions and 20 deletions.
8 changes: 4 additions & 4 deletions alan/src/compile/integration_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1104,7 +1104,7 @@ test_gpgpu!(hello_gpu_odd => r#"

test_gpgpu!(gpu_map => r#"
export fn main {
let b = GBuffer([1, 2, 3, 4]);
let b = GBuffer([1.i32, 2.i32, 3.i32, 4.i32]);
let out = b.map(fn (val: gi32) = val + 2);
out.read{i32}.print;
}"#;
Expand All @@ -1113,7 +1113,7 @@ test_gpgpu!(gpu_map => r#"

test_gpgpu!(gpu_if => r#"
export fn main {
let b = GBuffer([1, 2, 3, 4]);
let b = GBuffer([1.i32, 2.i32, 3.i32, 4.i32]);
let out = b.map(fn (val: gi32, i: gu32) = if(
i % 2 == 0,
val * i.gi32,
Expand All @@ -1125,9 +1125,9 @@ test_gpgpu!(gpu_if => r#"

test_gpgpu!(gpu_replace => r#"
export fn main {
let b = GBuffer([1, 2, 3, 4]);
let b = GBuffer([1.i32, 2.i32, 3.i32, 4.i32]);
b.map(fn (val: gi32) = val + 2).read{i32}.print;
b.replace([2, 4, 6, 8]);
b.replace([2.i32, 4.i32, 6.i32, 8.i32]);
b.map(fn (val: gi32) = val / 2).read{i32}.print;
}"#;
stdout "[3, 4, 5, 6]\n[1, 2, 3, 4]\n";
Expand Down
81 changes: 80 additions & 1 deletion alan/src/program/ctype.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3333,8 +3333,87 @@ impl CType {
// particularly for writing to disk or interfacing with network protocols, etc, so I'd
// prefer to keep it and have some compile-time guarantees we don't normally see.
match t {
CType::Void => CType::Int(0),
CType::Infer(..) => CType::Size(Box::new(t.clone())),
_ => CType::fail("TODO: Implement Size{T}!"),
CType::Type(_, t) => CType::size(t),
CType::Generic(..) => CType::fail("Cannot determine the size of an unbound generic"),
CType::Binds(t, ts) => {
if ts.len() > 0 {

Check warning

Code scanning / clippy

length comparison to zero Warning

length comparison to zero

Check warning

Code scanning / clippy

length comparison to zero Warning

length comparison to zero
CType::fail("Cannot determine the size of an unbound generic")
} else {
match &**t {
CType::TString(n) if n == "i8" => CType::Int(1),
CType::TString(n) if n == "u8" => CType::Int(1),
CType::TString(n) if n == "i16" => CType::Int(2),
CType::TString(n) if n == "u16" => CType::Int(2),
CType::TString(n) if n == "i32" => CType::Int(4),
CType::TString(n) if n == "u32" => CType::Int(4),
CType::TString(n) if n == "f32" => CType::Int(4),
CType::TString(n) if n == "i64" => CType::Int(8),
CType::TString(n) if n == "u64" => CType::Int(8),
CType::TString(n) if n == "f64" => CType::Int(8),
CType::TString(n) => {
CType::fail(&format!("Cannot determine the size of {}", n))
}
_ => CType::fail(&format!(
"Cannot determine the size of {}",
t.to_functional_string()
)),
}
}
}
CType::IntrinsicGeneric(..) => {
CType::fail("Cannot determine the size of an unbound generic")
}
CType::Int(_) | CType::Float(_) => CType::Int(8),
CType::Bool(_) => CType::Int(1),
CType::TString(s) => CType::Int(s.capacity() as i128),
CType::Group(t) | CType::Field(_, t) => CType::size(t),
CType::Tuple(ts) => {
let sizes = ts.iter().map(|t| CType::size(t)).collect::<Vec<CType>>();

Check warning

Code scanning / clippy

redundant closure Warning

redundant closure

Check warning

Code scanning / clippy

redundant closure Warning

redundant closure
let mut out_size = 0;
for t in sizes {
match t {
CType::Int(s) => out_size = out_size + s,

Check warning

Code scanning / clippy

manual implementation of an assign operation Warning

manual implementation of an assign operation

Check warning

Code scanning / clippy

manual implementation of an assign operation Warning

manual implementation of an assign operation
_ => unreachable!(),
}
}
CType::Int(out_size)
}
CType::Either(ts) => {
let sizes = ts.iter().map(|t| CType::size(t)).collect::<Vec<CType>>();

Check warning

Code scanning / clippy

redundant closure Warning

redundant closure

Check warning

Code scanning / clippy

redundant closure Warning

redundant closure
let mut out_size = 0;
for t in sizes {
match t {
CType::Int(s) => out_size = i128::max(out_size, s),
_ => unreachable!(),
}
}
CType::Int(out_size)
}
CType::Buffer(b, s) => {
let base_size = CType::size(b);
match (&base_size, &**s) {
(CType::Int(a), CType::Int(b)) => CType::Int(a + b),
(CType::Infer(..), _) | (_, CType::Infer(..)) => {
CType::Size(Box::new((&**b).clone()))

Check warning

Code scanning / clippy

this expression borrows a value the compiler would automatically borrow Warning

this expression borrows a value the compiler would automatically borrow

Check warning

Code scanning / clippy

this expression borrows a value the compiler would automatically borrow Warning

this expression borrows a value the compiler would automatically borrow
}
_ => unreachable!(),
}
}
CType::Array(_) => {
CType::fail("Cannot determine the size of an array, it's length is not static")
}
CType::Function(..)
| CType::Call(..)
| CType::Infix(_)
| CType::Prefix(_)
| CType::Method(_)
| CType::Property(_) => CType::fail("Cannot determine the size of a function"),
t => CType::fail(&format!(
"Getting the size of {} doesn't make any sense",
t.to_functional_string()
)),
}
}
pub fn filestr(f: &CType) -> CType {
Expand Down
70 changes: 55 additions & 15 deletions alan/src/std/root.ln
Original file line number Diff line number Diff line change
Expand Up @@ -1342,13 +1342,13 @@ export fn{Rs} mapWriteBuffer "alan_std::map_write_buffer_type" <- RootBacking ::
export fn{Js} mapWriteBuffer "alan_std.mapWriteBufferType" <- RootBacking :: () -> BufferUsages;
export fn{Rs} storageBuffer "alan_std::storage_buffer_type" <- RootBacking :: () -> BufferUsages;
export fn{Js} storageBuffer "alan_std.storageBufferType" <- RootBacking :: () -> BufferUsages;
export fn{Rs} GBuffer "alan_std::create_buffer_init" <- RootBacking :: (BufferUsages, i32[]) -> GBuffer;
export fn{Js} GBuffer "alan_std.createBufferInit" <- RootBacking :: (BufferUsages, i32[]) -> GBuffer;
export fn{Rs} GBuffer "alan_std::create_empty_buffer" <- RootBacking :: (BufferUsages, i64) -> GBuffer;
export fn{Js} GBuffer "alan_std.createEmptyBuffer" <- RootBacking :: (BufferUsages, i64) -> GBuffer;
export fn GBuffer(vals: i32[]) = GBuffer(storageBuffer(), vals);
export fn GBuffer{T}(vals: T[]) = GBuffer(storageBuffer(), vals.map(fn (v: T) -> i32 = v.i32));
export fn GBuffer(size: i64) = GBuffer(storageBuffer(), size);
export fn{Rs} GBuffer{T}(bu: BufferUsages, arr: T[]) = {"alan_std::create_buffer_init" <- RootBacking :: (BufferUsages, T[], i8) -> GBuffer}(bu, arr, {Size{T}}().i8);
export fn{Js} GBuffer{T}(bu: BufferUsages, arr: T[]) = {"alan_std.createBufferInit" <- RootBacking :: (BufferUsages, T[], i8) -> GBuffer}(bu, arr, {Size{T}}().i8);
export fn{Rs} GBuffer{T}(bu: BufferUsages, size: i64) = {"alan_std::create_empty_buffer" <- RootBacking :: (BufferUsages, i64, i8) -> GBuffer}(bu, size, {Size{T}}().i8);
// TODO: Get the type into JS
export fn{Js} GBuffer{T}(bu: BufferUsages, size: i64) = {"alan_std.createEmptyBuffer" <- RootBacking :: (BufferUsages, i64) -> GBuffer}(bu, size);
export fn GBuffer{T}(vals: T[]) = GBuffer{T}(storageBuffer(), vals);
export fn GBuffer{T}(size: i64) = GBuffer{T}(storageBuffer(), size);
export fn{Rs} len "alan_std::bufferlen" <- RootBacking :: GBuffer -> i64;
export fn{Js} len "alan_std.bufferlen" <- RootBacking :: GBuffer -> i64;
export fn{Rs} id "alan_std::buffer_id" <- RootBacking :: GBuffer -> string;
Expand Down Expand Up @@ -1387,9 +1387,8 @@ export fn{Rs} run "alan_std::gpu_run" <- RootBacking :: GPGPU;
export fn{Js} run "alan_std.gpuRun" <- RootBacking :: GPGPU;
export fn{Rs} read{T} "alan_std::read_buffer" <- RootBacking :: GBuffer -> T[];
export fn{Js} read{T} "alan_std.readBuffer" <- RootBacking :: GBuffer -> T[];
export fn{Rs} replace "alan_std::replace_buffer" <- RootBacking :: (GBuffer, i32[]) -> ()!;
export fn{Js} replace "alan_std.replaceBuffer" <- RootBacking :: (GBuffer, i32[]) -> ()!;
export fn replace{T}(g: GBuffer, v: T[]) -> ()! = g.replace(v.map(i32));
export fn{Rs} replace{T} "alan_std::replace_buffer" <- RootBacking :: (GBuffer, T[]) -> ()!;
export fn{Js} replace{T} "alan_std.replaceBuffer" <- RootBacking :: (GBuffer, T[]) -> ()!;

// The WgpuType provides a mechanism to build up the logic for a compute shader with normal-looking
// Alan code, not requiring to think in two different languages at the same time. All of the GPU
Expand All @@ -1400,6 +1399,37 @@ export type WgpuType{N} =
statements: Dict{string, string},
buffers: Set{GBuffer};

// It would be nice to not have this "lookup" type, but this seems to be the easiest way to
// accomplish the desired goal of termining the constructor type fn to use when casting between CPU
// and GPU memory blocks. This approach locks in only the GPU types defined in the root scope as
// automatically converting back and forth. That may not be a bad thing, but it is annoying.
export type WgpuTypeMap =
u32: u32,
i32: i32,
f32: f32,
bool: bool,
vec2u: Buffer{u32, 2}, // Until Rust has SIMD support these are just normal buffers
vec2i: Buffer{i32, 2},
vec2f: Buffer{f32, 2},
"vec2<bool>": Buffer{bool, 2},
vec3u: Buffer{u32, 3},
vec3i: Buffer{i32, 3},
vec3f: Buffer{f32, 3},
"vec3<bool>": Buffer{bool, 3},
vec4u: Buffer{u32, 4},
vec4i: Buffer{i32, 4},
vec4f: Buffer{f32, 4},
"vec4<bool>": Buffer{bool, 4},
mat2x2f: Buffer{f32, 4}, // It may make sense to make CPU-native matrix types, though
mat2x3f: Buffer{f32, 6},
mat2x4f: Buffer{f32, 8},
mat3x2f: Buffer{f32, 6},
mat3x3f: Buffer{f32, 9},
mat3x4f: Buffer{f32, 12},
mat4x2f: Buffer{f32, 8},
mat4x3f: Buffer{f32, 12},
mat4x4f: Buffer{f32, 16};

// Scalar types and constructors
export type gu32 = WgpuType{"u32"};
export type gi32 = WgpuType{"i32"};
Expand Down Expand Up @@ -4248,19 +4278,29 @@ export fn if(c: gvec4b, t: gvec4b, f: gvec4b) = piecewiseIf{gvec4b, gvec4b}(c, t

// GBuffer methods

// TODO: Support more than i32 for GBuffer
export fn map(gb: GBuffer, f: (gi32) -> gi32) {
/* This one doesn't seem to work yet for some reason
export fn map{G, G2}(gb: GBuffer, f: G -> G2) {
let idx = gFor(gb.len);
let val = gb[idx];
let out = GBuffer{i32}(gb.len);
let compute = out[idx].store(f(val));
compute.build.run;
return out;
}*/
export fn map(gb: GBuffer, f: gi32 -> gi32) {
let idx = gFor(gb.len);
let val = gb[idx];
let out = GBuffer(gb.len.mul(4));
let out = GBuffer{i32}(gb.len);
let compute = out[idx].store(f(val));
compute.build.run;
return out;
}
export fn map(gb: GBuffer, f: (gi32, gu32) -> gi32) {
// This one technically only works for i32/u32/f32 by chance. Once the issue above is fixed, this
// one can be fixed, too.
export fn map{G, G2}(gb: GBuffer, f: (G, gu32) -> G2) {
let idx = gFor(gb.len);
let val = gb[idx];
let out = GBuffer(gb.len.mul(4));
let out = GBuffer{i32}(gb.len);
let compute = out[idx].store(f(val, idx));
compute.build.run;
return out;
Expand Down

0 comments on commit 99222b1

Please sign in to comment.