Skip to content
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.

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

Implement Pack/Unpack for HLSL #2353

Merged
merged 21 commits into from
Jun 23, 2023
Merged
Show file tree
Hide file tree
Changes from 13 commits
Commits
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
291 changes: 290 additions & 1 deletion src/back/hlsl/writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,17 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
crate::MathFunction::Asinh
| crate::MathFunction::Acosh
| crate::MathFunction::Atanh
| crate::MathFunction::Unpack2x16float => {
| crate::MathFunction::Unpack2x16float
| crate::MathFunction::Unpack2x16snorm
| crate::MathFunction::Unpack2x16unorm
| crate::MathFunction::Unpack4x8snorm
| crate::MathFunction::Unpack4x8unorm
// TODO: These use multiple args, unsure how to bake them
Elabajaba marked this conversation as resolved.
Show resolved Hide resolved
| crate::MathFunction::Pack2x16float
| crate::MathFunction::Pack2x16snorm
| crate::MathFunction::Pack2x16unorm
| crate::MathFunction::Pack4x8snorm
| crate::MathFunction::Pack4x8unorm => {
self.need_bake_expressions.insert(arg);
}
crate::MathFunction::CountLeadingZeros => {
Expand Down Expand Up @@ -2590,7 +2600,12 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
enum Function {
Asincosh { is_sin: bool },
Atanh,
ExtractBits,
InsertBits,
Pack2x16float,
PackBits { signed: bool, dims: u32, scale: u32 },
Unpack2x16float,
UnpackBits { signed: bool, dims: u32, scale: u32 },
Regular(&'static str),
MissingIntOverload(&'static str),
MissingIntReturnType(&'static str),
Expand Down Expand Up @@ -2664,7 +2679,52 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
Mf::ReverseBits => Function::MissingIntOverload("reversebits"),
Mf::FindLsb => Function::MissingIntReturnType("firstbitlow"),
Mf::FindMsb => Function::MissingIntReturnType("firstbithigh"),
Mf::ExtractBits => Function::ExtractBits,
Mf::InsertBits => Function::InsertBits,
// Data Packing
Mf::Pack2x16float => Function::Pack2x16float,
Mf::Pack2x16snorm => Function::PackBits {
signed: true,
dims: 2,
scale: 32767,
},
Mf::Pack2x16unorm => Function::PackBits {
signed: false,
dims: 2,
scale: 65535,
},
Mf::Pack4x8snorm => Function::PackBits {
signed: true,
dims: 4,
scale: 127,
},
Mf::Pack4x8unorm => Function::PackBits {
signed: false,
dims: 4,
scale: 255,
},
// Data Unpacking
Mf::Unpack2x16float => Function::Unpack2x16float,
Mf::Unpack2x16snorm => Function::UnpackBits {
signed: true,
dims: 2,
scale: 32767,
},
Mf::Unpack2x16unorm => Function::UnpackBits {
signed: false,
dims: 2,
scale: 65535,
},
Mf::Unpack4x8snorm => Function::UnpackBits {
signed: true,
dims: 4,
scale: 127,
},
Mf::Unpack4x8unorm => Function::UnpackBits {
signed: false,
dims: 4,
scale: 255,
},
_ => return Err(Error::Unimplemented(format!("write_expr_math {fun:?}"))),
};

Expand All @@ -2688,13 +2748,242 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "))")?;
}
Function::ExtractBits => {
// e: T,
// offset: u32,
// count: u32
// T is u32 or i32 or vecN<u32> or vecN<i32>
if let (Some(offset), Some(count)) = (arg1, arg2) {
let inner = func_ctx.info[expr].ty.inner_with(&module.types);
let scalar_kind = inner.scalar_kind();
let scalar_width = inner.scalar_width().unwrap_or(32);
let scalar_max: u32 = match scalar_width {
8 => 0xff,
16 => 0xffff,
32 => 0xffffffff,
_ => {
return Err(Error::Unimplemented(format!(
"write_expr_math extract_bits for scalar_width {}",
scalar_width
)))
}
};

if let Some(ScalarKind::Uint) = scalar_kind {
// Unsigned
// ((e >> offset) & (count == 32u ? 0xffffffffu : ((1 << count) - 1)))
write!(self.out, "((")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " >> ")?;
self.write_expr(module, offset, func_ctx)?;
write!(self.out, ") & (")?;
self.write_expr(module, count, func_ctx)?;
write!(self.out, " == {scalar_width}u ? {scalar_max}u : ((1u << ")?;
self.write_expr(module, count, func_ctx)?;
write!(self.out, ") - 1)))")?;
} else {
// Signed
// (count == 0 ? 0 : (e << (32 - count - offset)) >> (32 - count))
write!(self.out, "(")?;
self.write_expr(module, count, func_ctx)?;
write!(self.out, " == 0 ? 0 : (")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " << ({scalar_width} - ")?;
self.write_expr(module, count, func_ctx)?;
write!(self.out, " - ")?;
self.write_expr(module, offset, func_ctx)?;
write!(self.out, ")) >> ({scalar_width} - ")?;
self.write_expr(module, count, func_ctx)?;
write!(self.out, "))")?;
Elabajaba marked this conversation as resolved.
Show resolved Hide resolved
}
}
}
Function::InsertBits => {
// e: T,
// newbits: T,
// offset: u32,
// count: u32
// returns T
// T is i32, u32, vecN<i32>, or vecN<u32>
if let (Some(newbits), Some(offset), Some(count)) = (arg1, arg2, arg3) {
let inner = func_ctx.info[expr].ty.inner_with(&module.types);
let scalar_width = inner.scalar_width().unwrap_or(32);
let scalar_max: u64 = match scalar_width {
8 => 0xff,
16 => 0xffff,
32 => 0xffffffff,
64 => 0xffffffffffffffff,
_ => {
return Err(Error::Unimplemented(format!(
"write_expr_math extract_bits for scalar_width {}",
scalar_width
)))
}
};
// mask = ((0xFFFFFFFFu >> (32 - count)) << offset)
// return (count == 0 ? e : ((e & ~mask) | ((newbits << offset) & mask)))
write!(self.out, "(")?;
self.write_expr(module, count, func_ctx)?;
write!(self.out, " == 0 ? ")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " : ")?;
write!(self.out, "(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " & ~")?;
// mask
write!(self.out, "(({scalar_max}u >> ({scalar_width}u - ")?;
self.write_expr(module, count, func_ctx)?;
write!(self.out, ")) << ")?;
self.write_expr(module, offset, func_ctx)?;
write!(self.out, ")")?;
// end mask
write!(self.out, ") | ((")?;
self.write_expr(module, newbits, func_ctx)?;
write!(self.out, " << ")?;
self.write_expr(module, offset, func_ctx)?;
write!(self.out, ") & ")?;
// // mask
write!(self.out, "(({scalar_max}u >> ({scalar_width}u - ")?;
self.write_expr(module, count, func_ctx)?;
write!(self.out, ")) << ")?;
self.write_expr(module, offset, func_ctx)?;
write!(self.out, ")")?;
// // end mask
write!(self.out, "))")?;
}
}
Function::Pack2x16float => {
write!(self.out, "f32tof16(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "[0]) | (f32tof16(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "[1]) << 16)")?;
Elabajaba marked this conversation as resolved.
Show resolved Hide resolved
}
Function::PackBits {
teoxoy marked this conversation as resolved.
Show resolved Hide resolved
signed,
dims,
scale,
} => {
if dims == 4 {
if signed {
// pack4x8snorm
write!(self.out, "uint((int(floor(clamp(")?;
self.write_expr(module, arg, func_ctx)?;
write!(
self.out,
"[0], -1.0, 1.0) * {scale}.0 + 0.5)) & 0xFF) | ((int(floor(clamp("
)?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "[1], -1.0, 1.0) * {scale}.0 + 0.5)) & 0xFF) << 8) | ((int(floor(clamp(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "[2], -1.0, 1.0) * {scale}.0 + 0.5)) & 0xFF) << 16) | ((int(floor(clamp(")?;
self.write_expr(module, arg, func_ctx)?;
write!(
self.out,
"[3], -1.0, 1.0) * {scale}.0 + 0.5)) & 0xFF) << 24))",
)?;
} else {
// pack4x8unorm
write!(self.out, "(uint(clamp(")?;
self.write_expr(module, arg, func_ctx)?;
write!(
self.out,
"[0], 0.0, 1.0) * {scale}.0 + 0.5) | uint(clamp("
)?;
self.write_expr(module, arg, func_ctx)?;
write!(
self.out,
"[1], 0.0, 1.0) * {scale}.0 + 0.5) << 8 | uint(clamp("
)?;
self.write_expr(module, arg, func_ctx)?;
write!(
self.out,
"[2], 0.0, 1.0) * {scale}.0 + 0.5) << 16 | uint(clamp("
)?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "[3], 0.0, 1.0) * {scale}.0 + 0.5) << 24)")?;
}
} else if dims == 2 {
if signed {
// pack2x16snorm
write!(self.out, "uint((int(floor(clamp(")?;
self.write_expr(module, arg, func_ctx)?;
write!(
self.out,
"[0], -1.0, 1.0) * {scale}.0 + 0.5))) & 0xFFFF) | (((int(floor(clamp("
)?;
self.write_expr(module, arg, func_ctx)?;
write!(
self.out,
"[1], -1.0, 1.0) * {scale}.0 + 0.5)) & 0xFFFF) << 16))",
)?;
} else {
// pack2x16unorm
write!(self.out, "uint(uint(float(clamp(")?;
self.write_expr(module, arg, func_ctx)?;
write!(
self.out,
"[0], 0, 1) * {scale}.0) + 0.5) | uint(float(clamp("
)?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "[1], 0, 1) * {scale}.0) + 0.5) << 16)")?;
}
}
}
Function::Unpack2x16float => {
write!(self.out, "float2(f16tof32(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "), f16tof32((")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, ") >> 16))")?;
}
Function::UnpackBits {
signed,
dims,
scale,
} => {
if dims == 4 {
if signed {
// Unpack4x8snorm
write!(self.out, "clamp(float4((int4(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "<< 24, ")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " << 16, ")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " << 8, ")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, ") >> 24)) / {scale}.0, -1.0, 1.0)")?;
} else {
// Unpack4x8unorm
write!(self.out, "clamp(float4(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " & 0xFF, ")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " >> 8 & 0xFF, ")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " >> 16 & 0xFF, ")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " >> 24 & 0xFF) / {scale}.0, 0.0, 1.0)")?;
}
} else if dims == 2 {
if signed {
// Unpack2x16snorm
write!(self.out, "clamp(float2(int2(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, "<< 16, ")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, ") >> 16) / {scale}.0, -1.0, 1.0)")?;
} else {
// Unpack2x16unorm
write!(self.out, "clamp(float2(float(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " & 0xFFFF), float(")?;
self.write_expr(module, arg, func_ctx)?;
write!(self.out, " >> 16 & 0xFFFF)) / {scale}.0, 0.0, 1.0)")?;
}
}
}
Function::Regular(fun_name) => {
write!(self.out, "{fun_name}(")?;
self.write_expr(module, arg, func_ctx)?;
Expand Down
11 changes: 11 additions & 0 deletions src/proc/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,17 @@ impl super::TypeInner {
}
}

pub const fn scalar_width(&self) -> Option<u8> {
// Multiply by 8 to get the bit width
match *self {
super::TypeInner::Scalar { width, .. } | super::TypeInner::Vector { width, .. } => {
Some(width * 8)
}
super::TypeInner::Matrix { width, .. } => Some(width * 8),
_ => None,
}
}

pub const fn pointer_space(&self) -> Option<crate::AddressSpace> {
match *self {
Self::Pointer { space, .. } => Some(space),
Expand Down
Loading