RustHPC/rust-gpu
1
0
Fork 0
mirror of https://github.com/Rust-GPU/rust-gpu.git synced 2026-06-05 17:19:51 +09:00
Code Issues Actions Packages Projects Releases Wiki Activity

builder: use SPIR-V instructions for checked_{add,sub,mul} and saturating_{add,sub}.

Browse Source
This commit is contained in:
Eduard-Mihai Burtescu
2025-10-28 13:42:42 +02:00
parent 003dde7b4c
commit 39eec00a8f
2 changed files with 108 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
crates/rustc_codegen_spirv/src/builder/builder_methods.rs
View File

@@ -27,7 +27,7 @@ use rustc_codegen_ssa::traits::{
};
use rustc_middle::bug;
use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
use rustc_middle::ty::layout::TyAndLayout;
use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
use rustc_middle::ty::{self, AtomicOrdering, Ty};
use rustc_span::Span;
use rustc_target::callconv::FnAbi;
@@ -1721,30 +1721,15 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
fn checked_binop(
&mut self,
oop: OverflowOp,
ty: Ty<'_>,
ty: Ty<'tcx>,
lhs: Self::Value,
rhs: Self::Value,
) -> (Self::Value, Self::Value) {
// adopted partially from https://github.com/ziglang/zig/blob/master/src/codegen/spirv.zig
let is_add = match oop {
OverflowOp::Add => true,
OverflowOp::Sub => false,
OverflowOp::Mul => {
// NOTE(eddyb) this needs to be `undef`, not `false`/`true`, because
// we don't want the user's boolean constants to keep the zombie alive.
let bool = SpirvType::Bool.def(self.span(), self);
let overflowed = self.undef(bool);
let result = (self.mul(lhs, rhs), overflowed);
self.zombie(result.1.def(self), "checked mul is not supported yet");
return result;
}
};
let signed = match ty.kind() {
ty::Int(_) => true,
ty::Uint(_) => false,
other => self.fatal(format!(
"Unexpected {} type: {other:#?}",
_ => self.fatal(format!(
"unexpected {} type: {ty}",
match oop {
OverflowOp::Add => "checked add",
OverflowOp::Sub => "checked sub",
@@ -1753,13 +1738,17 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
)),
};
let result = if is_add {
self.add(lhs, rhs)
} else {
self.sub(lhs, rhs)
};
// HACK(eddyb) SPIR-V `OpIAddCarry`/`OpISubBorrow` are specifically for
// unsigned overflow, so signed overflow still needs this custom logic.
if signed && let OverflowOp::Add | OverflowOp::Sub = oop {
let result = match oop {
OverflowOp::Add => self.add(lhs, rhs),
OverflowOp::Sub => self.sub(lhs, rhs),
OverflowOp::Mul => unreachable!(),
};
// adopted partially from https://github.com/ziglang/zig/blob/master/src/codegen/spirv.zig
let overflowed = if signed {
// when adding, overflow could happen if
// - rhs is positive and result < lhs; or
// - rhs is negative and result > lhs
@@ -1771,30 +1760,80 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
// this is equivalent to (rhs < 0) == (result < lhs)
let rhs_lt_zero = self.icmp(IntPredicate::IntSLT, rhs, self.constant_int(rhs.ty, 0));
let result_gt_lhs = self.icmp(
if is_add {
IntPredicate::IntSGT
} else {
IntPredicate::IntSLT
match oop {
OverflowOp::Add => IntPredicate::IntSGT,
OverflowOp::Sub => IntPredicate::IntSLT,
OverflowOp::Mul => unreachable!(),
},
result,
lhs,
);
self.icmp(IntPredicate::IntEQ, rhs_lt_zero, result_gt_lhs)
} else {
// for unsigned addition, overflow occurred if the result is less than any of the operands.
// for subtraction, overflow occurred if the result is greater.
self.icmp(
if is_add {
IntPredicate::IntULT
} else {
IntPredicate::IntUGT
},
result,
lhs,
)
let overflowed = self.icmp(IntPredicate::IntEQ, rhs_lt_zero, result_gt_lhs);
return (result, overflowed);
}
let result_type = self.layout_of(ty).spirv_type(self.span(), self);
let pair_result_type = {
let field_types = [result_type, result_type];
let (field_offsets, size, align) = crate::abi::auto_struct_layout(self, &field_types);
SpirvType::Adt {
def_id: None,
size,
align,
field_types: &field_types,
field_offsets: &field_offsets,
field_names: None,
}
.def(self.span(), self)
};
(result, overflowed)
let lhs = lhs.def(self);
let rhs = rhs.def(self);
let pair_result = match oop {
OverflowOp::Add => self
.emit()
.i_add_carry(pair_result_type, None, lhs, rhs)
.unwrap(),
OverflowOp::Sub => self
.emit()
.i_sub_borrow(pair_result_type, None, lhs, rhs)
.unwrap(),
OverflowOp::Mul => {
if signed {
self.emit()
.s_mul_extended(pair_result_type, None, lhs, rhs)
.unwrap()
} else {
self.emit()
.u_mul_extended(pair_result_type, None, lhs, rhs)
.unwrap()
}
}
}
.with_type(pair_result_type);
let result_lo = self.extract_value(pair_result, 0);
let result_hi = self.extract_value(pair_result, 1);
// HACK(eddyb) SPIR-V lacks any `(T, T) -> (T, bool)` instructions,
// so instead `result_hi` is compared with the value expected in the
// non-overflow case (`0`, or `-1` for negative signed multiply result).
let expected_nonoverflowing_hi = match (oop, signed) {
(OverflowOp::Add | OverflowOp::Sub, _) | (OverflowOp::Mul, false) => {
self.const_uint(result_type, 0)
}
(OverflowOp::Mul, true) => {
// HACK(eddyb) `(x: iN) >> (N - 1)` will spread the sign bit
// across all `N` bits of `iN`, and should be equivalent to
// `if x < 0 { -1 } else { 0 }`, without needing compare+select).
let result_width = u32::try_from(self.int_width(result_type)).unwrap();
self.ashr(result_lo, self.const_u32(result_width - 1))
}
};
let overflowed = self.icmp(IntPredicate::IntNE, result_hi, expected_nonoverflowing_hi);
(result_lo, overflowed)
}
// rustc has the concept of an immediate vs. memory type - bools are compiled to LLVM bools as

46
crates/rustc_codegen_spirv/src/builder/intrinsics.rs
View File

@@ -118,33 +118,39 @@ impl<'a, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'a, 'tcx> {
sym::saturating_add => {
assert_eq!(arg_tys[0], arg_tys[1]);
let result = match arg_tys[0].kind() {
TyKind::Int(_) | TyKind::Uint(_) => {
self.add(args[0].immediate(), args[1].immediate())
}
TyKind::Float(_) => self.fadd(args[0].immediate(), args[1].immediate()),
match arg_tys[0].kind() {
TyKind::Int(_) | TyKind::Uint(_) => self
.emit()
.i_add_sat_intel(
ret_ty,
None,
args[0].immediate().def(self),
args[1].immediate().def(self),
)
.unwrap()
.with_type(ret_ty),
other => self.fatal(format!(
"Unimplemented saturating_add intrinsic type: {other:#?}"
"unimplemented saturating_add intrinsic type: {other:#?}"
)),
};
// TODO: Implement this
self.zombie(result.def(self), "saturating_add is not implemented yet");
result
}
}
sym::saturating_sub => {
assert_eq!(arg_tys[0], arg_tys[1]);
let result = match &arg_tys[0].kind() {
TyKind::Int(_) | TyKind::Uint(_) => {
self.sub(args[0].immediate(), args[1].immediate())
}
TyKind::Float(_) => self.fsub(args[0].immediate(), args[1].immediate()),
match &arg_tys[0].kind() {
TyKind::Int(_) | TyKind::Uint(_) => self
.emit()
.i_sub_sat_intel(
ret_ty,
None,
args[0].immediate().def(self),
args[1].immediate().def(self),
)
.unwrap()
.with_type(ret_ty),
other => self.fatal(format!(
"Unimplemented saturating_sub intrinsic type: {other:#?}"
"unimplemented saturating_sub intrinsic type: {other:#?}"
)),
};
// TODO: Implement this
self.zombie(result.def(self), "saturating_sub is not implemented yet");
result
}
}
sym::sqrtf32 | sym::sqrtf64 | sym::sqrtf128 => {