builder_spirv: add SpirvConst::Scalar to replace SpirvConst float/integer variants.

eddyb · eddyb · commit 929b227053cc · 2024-12-03T19:51:34.000+02:00
diff --git a/crates/rustc_codegen_spirv/src/builder/builder_methods.rs b/crates/rustc_codegen_spirv/src/builder/builder_methods.rs
@@ -49,8 +49,7 @@ macro_rules! simple_op {
                             let size = Size::from_bits(bits);
                             let as_u128 = |const_val| {
                                 let x = match const_val {
-                                    SpirvConst::U32(x) => x as u128,
-                                    SpirvConst::U64(x) => x as u128,
+                                    SpirvConst::Scalar(x) => x,
                                     _ => return None,
                                 };
                                 Some(if signed {
@@ -225,7 +224,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             }
             SpirvType::Array { element, count } => {
                 let elem_pat = self.memset_const_pattern(&self.lookup_type(element), fill_byte);
-                let count = self.builder.lookup_const_u64(count).unwrap() as usize;
+                let count = self.builder.lookup_const_scalar(count).unwrap() as usize;
                 self.constant_composite(
                     ty.def(self.span(), self),
                     iter::repeat(elem_pat).take(count),
@@ -269,7 +268,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             SpirvType::Adt { .. } => self.fatal("memset on structs not implemented yet"),
             SpirvType::Array { element, count } => {
                 let elem_pat = self.memset_dynamic_pattern(&self.lookup_type(element), fill_var);
-                let count = self.builder.lookup_const_u64(count).unwrap() as usize;
+                let count = self.builder.lookup_const_scalar(count).unwrap() as usize;
                 self.emit()
                     .composite_construct(
                         ty.def(self.span(), self),
@@ -327,7 +326,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             .lookup_type(pat.ty)
             .sizeof(self)
             .expect("Unable to memset a dynamic sized object");
-        let size_elem_const = self.constant_int(size_bytes.ty, size_elem.bytes());
+        let size_elem_const = self.constant_int(size_bytes.ty, size_elem.bytes().into());
         let zero = self.constant_int(size_bytes.ty, 0);
         let one = self.constant_int(size_bytes.ty, 1);
         let zero_align = Align::from_bytes(0).unwrap();
@@ -595,8 +594,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         // - dynamic indexing of a single array
         let const_ptr_offset = self
             .builder
-            .lookup_const_u64(ptr_base_index)
-            .and_then(|idx| Some(idx * self.lookup_type(ty).sizeof(self)?));
+            .lookup_const_scalar(ptr_base_index)
+            .and_then(|idx| Some(u64::try_from(idx).ok()? * self.lookup_type(ty).sizeof(self)?));
         if let Some(const_ptr_offset) = const_ptr_offset {
             if let Some((base_indices, base_pointee_ty)) = self.recover_access_chain_from_offset(
                 original_pointee_ty,
@@ -707,7 +706,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         let mut emit = self.emit();
 
         let non_zero_ptr_base_index =
-            ptr_base_index.filter(|&idx| self.builder.lookup_const_u64(idx) != Some(0));
+            ptr_base_index.filter(|&idx| self.builder.lookup_const_scalar(idx) != Some(0));
         if let Some(ptr_base_index) = non_zero_ptr_base_index {
             let result = if is_inbounds {
                 emit.in_bounds_ptr_access_chain(
@@ -1083,8 +1082,8 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
         // HACK(eddyb) constant-fold branches early on, as the `core` library is
         // starting to get a lot of `if cfg!(debug_assertions)` added to it.
         match self.builder.lookup_const_by_id(cond) {
-            Some(SpirvConst::Bool(true)) => self.br(then_llbb),
-            Some(SpirvConst::Bool(false)) => self.br(else_llbb),
+            Some(SpirvConst::Scalar(1)) => self.br(then_llbb),
+            Some(SpirvConst::Scalar(0)) => self.br(else_llbb),
             _ => {
                 self.emit()
                     .branch_conditional(cond, then_llbb, else_llbb, empty())
@@ -2232,7 +2231,10 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
                 "memcpy with mem flags is not supported yet: {flags:?}"
             ));
         }
-        let const_size = self.builder.lookup_const_u64(size).map(Size::from_bytes);
+        let const_size = self
+            .builder
+            .lookup_const_scalar(size)
+            .and_then(|size| Some(Size::from_bytes(u64::try_from(size).ok()?)));
         if const_size == Some(Size::ZERO) {
             // Nothing to do!
             return;
@@ -2306,6 +2308,12 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
                 "memset with mem flags is not supported yet: {flags:?}"
             ));
         }
+
+        let const_size = self
+            .builder
+            .lookup_const_scalar(size)
+            .and_then(|size| Some(Size::from_bytes(u64::try_from(size).ok()?)));
+
         let elem_ty = match self.lookup_type(ptr.ty) {
             SpirvType::Pointer { pointee } => pointee,
             _ => self.fatal(format!(
@@ -2314,13 +2322,13 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
             )),
         };
         let elem_ty_spv = self.lookup_type(elem_ty);
-        let pat = match self.builder.lookup_const_u64(fill_byte) {
+        let pat = match self.builder.lookup_const_scalar(fill_byte) {
             Some(fill_byte) => self.memset_const_pattern(&elem_ty_spv, fill_byte as u8),
             None => self.memset_dynamic_pattern(&elem_ty_spv, fill_byte.def(self)),
         }
         .with_type(elem_ty);
-        match self.builder.lookup_const_u64(size) {
-            Some(size) => self.memset_constant_size(ptr, pat, size),
+        match const_size {
+            Some(size) => self.memset_constant_size(ptr, pat, size.bytes()),
             None => self.memset_dynamic_size(ptr, pat, size),
         }
     }
@@ -2354,7 +2362,7 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
             SpirvType::Vector { element, .. } => element,
             other => self.fatal(format!("extract_element not implemented on type {other:?}")),
         };
-        match self.builder.lookup_const_u64(idx) {
+        match self.builder.lookup_const_scalar(idx) {
             Some(const_index) => self.emit().composite_extract(
                 result_type,
                 None,
@@ -2781,7 +2789,7 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
                 let mut decoded_format_args = DecodedFormatArgs::default();
 
                 let const_u32_as_usize = |ct_id| match self.builder.lookup_const_by_id(ct_id)? {
-                    SpirvConst::U32(x) => Some(x as usize),
+                    SpirvConst::Scalar(x) => Some(u32::try_from(x).ok()? as usize),
                     _ => None,
                 };
                 let const_slice_as_elem_ids = |slice_ptr_and_len_ids: &[Word]| {
@@ -2948,10 +2956,10 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
                             match (inst.class.opcode, inst.result_id, &id_operands[..]) {
                                 (Op::Bitcast, Some(r), &[x]) => Inst::Bitcast(r, x),
                                 (Op::InBoundsAccessChain, Some(r), &[p, i]) => {
-                                    if let Some(SpirvConst::U32(i)) =
+                                    if let Some(SpirvConst::Scalar(i)) =
                                         self.builder.lookup_const_by_id(i)
                                     {
-                                        Inst::InBoundsAccessChain(r, p, i)
+                                        Inst::InBoundsAccessChain(r, p, i as u32)
                                     } else {
                                         Inst::Unsupported(inst.class.opcode)
                                     }
diff --git a/crates/rustc_codegen_spirv/src/builder/byte_addressable_buffer.rs b/crates/rustc_codegen_spirv/src/builder/byte_addressable_buffer.rs
@@ -111,7 +111,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     count,
                 ),
             SpirvType::Array { element, count } => {
-                let count = match self.builder.lookup_const_u64(count) {
+                let count = match self.builder.lookup_const_scalar(count) {
                     Some(count) => count as u32,
                     None => return self.load_err(original_type, result_type),
                 };
@@ -301,7 +301,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     count,
                 ),
             SpirvType::Array { element, count } => {
-                let count = match self.builder.lookup_const_u64(count) {
+                let count = match self.builder.lookup_const_scalar(count) {
                     Some(count) => count as u32,
                     None => return self.store_err(original_type, value),
                 };
diff --git a/crates/rustc_codegen_spirv/src/builder/mod.rs b/crates/rustc_codegen_spirv/src/builder/mod.rs
@@ -113,8 +113,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 other.debug(shift.ty, self)
             )),
         };
-        let int_size = self.constant_int(shift.ty, width as u64);
-        let mask = self.constant_int(shift.ty, (width - 1) as u64);
+        let int_size = self.constant_int(shift.ty, width.into());
+        let mask = self.constant_int(shift.ty, (width - 1).into());
         let zero = self.constant_int(shift.ty, 0);
         let bool = SpirvType::Bool.def(self.span(), self);
         // https://stackoverflow.com/a/10134877
diff --git a/crates/rustc_codegen_spirv/src/builder_spirv.rs b/crates/rustc_codegen_spirv/src/builder_spirv.rs
@@ -4,12 +4,13 @@ use crate::spirv_type::SpirvType;
 use crate::symbols::Symbols;
 use crate::target::SpirvTarget;
 use crate::target_feature::TargetFeature;
-use rspirv::dr::{Block, Builder, Module, Operand};
+use rspirv::dr::{Block, Builder, Instruction, Module, Operand};
 use rspirv::spirv::{
     AddressingModel, Capability, MemoryModel, Op, SourceLanguage, StorageClass, Word,
 };
 use rspirv::{binary::Assemble, binary::Disassemble};
 use rustc_arena::DroplessArena;
+use rustc_codegen_ssa::traits::ConstMethods as _;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::sync::Lrc;
 use rustc_middle::bug;
@@ -18,6 +19,7 @@ use rustc_middle::ty::TyCtxt;
 use rustc_span::source_map::SourceMap;
 use rustc_span::symbol::Symbol;
 use rustc_span::{FileName, FileNameDisplayPreference, SourceFile, Span, DUMMY_SP};
+use rustc_target::abi::Size;
 use std::assert_matches::assert_matches;
 use std::cell::{RefCell, RefMut};
 use std::hash::{Hash, Hasher};
@@ -221,13 +223,8 @@ impl SpirvValueExt for Word {
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub enum SpirvConst<'a, 'tcx> {
-    U32(u32),
-    U64(u64),
-    /// f32 isn't hash, so store bits
-    F32(u32),
-    /// f64 isn't hash, so store bits
-    F64(u64),
-    Bool(bool),
+    /// Constants of boolean, integer or floating-point type (up to 128-bit).
+    Scalar(u128),
 
     Null,
     Undef,
@@ -273,11 +270,7 @@ impl<'tcx> SpirvConst<'_, 'tcx> {
 
         match self {
             // FIXME(eddyb) these are all noop cases, could they be automated?
-            SpirvConst::U32(v) => SpirvConst::U32(v),
-            SpirvConst::U64(v) => SpirvConst::U64(v),
-            SpirvConst::F32(v) => SpirvConst::F32(v),
-            SpirvConst::F64(v) => SpirvConst::F64(v),
-            SpirvConst::Bool(v) => SpirvConst::Bool(v),
+            SpirvConst::Scalar(v) => SpirvConst::Scalar(v),
             SpirvConst::Null => SpirvConst::Null,
             SpirvConst::Undef => SpirvConst::Undef,
             SpirvConst::ZombieUndefForFnAddr => SpirvConst::ZombieUndefForFnAddr,
@@ -570,8 +563,26 @@ impl<'tcx> BuilderSpirv<'tcx> {
         val: SpirvConst<'_, 'tcx>,
         cx: &CodegenCx<'tcx>,
     ) -> SpirvValue {
+        let scalar_ty = match val {
+            SpirvConst::Scalar(_) => Some(cx.lookup_type(ty)),
+            _ => None,
+        };
+
+        // HACK(eddyb) this is done so late (just before interning `val`) to
+        // minimize any potential misuse from direct `def_constant` calls.
+        let val = match (val, scalar_ty) {
+            (SpirvConst::Scalar(val), Some(SpirvType::Integer(bits, signed))) => {
+                let size = Size::from_bits(bits);
+                SpirvConst::Scalar(if signed {
+                    size.sign_extend(val)
+                } else {
+                    size.truncate(val)
+                })
+            }
+            _ => val,
+        };
+
         let val_with_type = WithType { ty, val };
-        let mut builder = self.builder(BuilderCursor::default());
         if let Some(entry) = self.const_to_id.borrow().get(&val_with_type) {
             // FIXME(eddyb) deduplicate this `if`-`else` and its other copies.
             let kind = if entry.legal.is_ok() {
@@ -582,16 +593,99 @@ impl<'tcx> BuilderSpirv<'tcx> {
             return SpirvValue { kind, ty };
         }
         let val = val_with_type.val;
+
+        // FIXME(eddyb) make this an extension method on `rspirv::dr::Builder`?
+        let const_op = |builder: &mut Builder, op, lhs, maybe_rhs: Option<_>| {
+            // HACK(eddyb) remove after `OpSpecConstantOp` support gets added to SPIR-T.
+            let spirt_has_const_op = false;
+
+            if !spirt_has_const_op {
+                let zombie = builder.undef(ty, None);
+                cx.zombie_with_span(
+                    zombie,
+                    DUMMY_SP,
+                    &format!("unsupported constant of type `{}`", cx.debug_type(ty)),
+                );
+                return zombie;
+            }
+
+            let id = builder.id();
+            builder
+                .module_mut()
+                .types_global_values
+                .push(Instruction::new(
+                    Op::SpecConstantOp,
+                    Some(ty),
+                    Some(id),
+                    [
+                        Operand::LiteralSpecConstantOpInteger(op),
+                        Operand::IdRef(lhs),
+                    ]
+                    .into_iter()
+                    .chain(maybe_rhs.map(Operand::IdRef))
+                    .collect(),
+                ));
+            id
+        };
+
+        let mut builder = self.builder(BuilderCursor::default());
         let id = match val {
-            SpirvConst::U32(v) | SpirvConst::F32(v) => builder.constant_bit32(ty, v),
-            SpirvConst::U64(v) | SpirvConst::F64(v) => builder.constant_bit64(ty, v),
-            SpirvConst::Bool(v) => {
-                if v {
-                    builder.constant_true(ty)
-                } else {
-                    builder.constant_false(ty)
+            SpirvConst::Scalar(v) => match scalar_ty.unwrap() {
+                SpirvType::Integer(..=32, _) | SpirvType::Float(..=32) => {
+                    builder.constant_bit32(ty, v as u32)
                 }
-            }
+                SpirvType::Integer(64, _) | SpirvType::Float(64) => {
+                    builder.constant_bit64(ty, v as u64)
+                }
+                SpirvType::Integer(128, false) => {
+                    // HACK(eddyb) avoid borrow conflicts.
+                    drop(builder);
+
+                    let const_64_u32_id = cx.const_u32(64).def_cx(cx);
+                    let [lo_id, hi_id] =
+                        [v as u64, (v >> 64) as u64].map(|half| cx.const_u64(half).def_cx(cx));
+
+                    builder = self.builder(BuilderCursor::default());
+                    let mut const_op =
+                        |op, lhs, maybe_rhs| const_op(&mut builder, op, lhs, maybe_rhs);
+                    let [lo_u128_id, hi_shifted_u128_id] =
+                        [(lo_id, None), (hi_id, Some(const_64_u32_id))].map(
+                            |(half_u64_id, shift)| {
+                                let mut half_u128_id = const_op(Op::UConvert, half_u64_id, None);
+                                if let Some(shift_amount_id) = shift {
+                                    half_u128_id = const_op(
+                                        Op::ShiftLeftLogical,
+                                        half_u128_id,
+                                        Some(shift_amount_id),
+                                    );
+                                }
+                                half_u128_id
+                            },
+                        );
+                    const_op(Op::BitwiseOr, lo_u128_id, Some(hi_shifted_u128_id))
+                }
+                SpirvType::Integer(128, true) | SpirvType::Float(128) => {
+                    // HACK(eddyb) avoid borrow conflicts.
+                    drop(builder);
+
+                    let v_u128_id = cx.const_u128(v).def_cx(cx);
+
+                    builder = self.builder(BuilderCursor::default());
+                    const_op(&mut builder, Op::Bitcast, v_u128_id, None)
+                }
+                SpirvType::Bool => match v {
+                    0 => builder.constant_false(ty),
+                    1 => builder.constant_true(ty),
+                    _ => cx
+                        .tcx
+                        .dcx()
+                        .fatal(format!("invalid constant value for bool: {v}")),
+                },
+                other => cx.tcx.dcx().fatal(format!(
+                    "SpirvConst::Scalar does not support type {}",
+                    other.debug(ty, cx)
+                )),
+            },
 
             SpirvConst::Null => builder.constant_null(ty),
             SpirvConst::Undef
@@ -606,11 +700,7 @@ impl<'tcx> BuilderSpirv<'tcx> {
         };
         #[allow(clippy::match_same_arms)]
         let legal = match val {
-            SpirvConst::U32(_)
-            | SpirvConst::U64(_)
-            | SpirvConst::F32(_)
-            | SpirvConst::F64(_)
-            | SpirvConst::Bool(_) => Ok(()),
+            SpirvConst::Scalar(_) => Ok(()),
 
             SpirvConst::Null => {
                 // FIXME(eddyb) check that the type supports `OpConstantNull`.
@@ -712,10 +802,9 @@ impl<'tcx> BuilderSpirv<'tcx> {
         }
     }
 
-    pub fn lookup_const_u64(&self, def: SpirvValue) -> Option<u64> {
+    pub fn lookup_const_scalar(&self, def: SpirvValue) -> Option<u128> {
         match self.lookup_const(def)? {
-            SpirvConst::U32(v) => Some(v as u64),
-            SpirvConst::U64(v) => Some(v),
+            SpirvConst::Scalar(v) => Some(v),
             _ => None,
         }
     }
diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/constant.rs b/crates/rustc_codegen_spirv/src/codegen_cx/constant.rs
diff --git a/crates/rustc_codegen_spirv/src/spirv_type.rs b/crates/rustc_codegen_spirv/src/spirv_type.rs
