leading_zeros: support 8bit, 16bit and emulate 64bit

Author: Firestar99

crates/rustc_codegen_spirv/src/builder/intrinsics.rs

@@ -383,10 +383,9 @@ impl Builder<'_, '_> {
     ) -> SpirvValue {
         let ty = arg.ty;
         match self.cx.lookup_type(ty) {
-            SpirvType::Integer(bits, _) => {
-                let int_0 = self.constant_int(ty, 0);
-                let int_bits = self.constant_int(ty, bits as u128).def(self);
+            SpirvType::Integer(bits, signed) => {
                 let bool = SpirvType::Bool.def(self.span(), self);
+                let u32 = SpirvType::Integer(32, false).def(self.span(), self);
 
                 let gl_op = if trailing {
                     // rust is always unsigned
@@ -396,24 +395,87 @@ impl Builder<'_, '_> {
                 };
 
                 let glsl = self.ext_inst.borrow_mut().import_glsl(self);
-                let find_xsb = self
-                    .emit()
-                    .ext_inst(ty, None, glsl, gl_op as u32, [Operand::IdRef(
-                        arg.def(self),
-                    )])
-                    .unwrap();
+                let find_xsb = |arg| {
+                    self.emit()
+                        .ext_inst(u32, None, glsl, gl_op as u32, [Operand::IdRef(arg)])
+                        .unwrap()
+                };
+
+                let converted = match bits {
+                    8 | 16 => {
+                        if trailing {
+                            let arg = self.emit().s_convert(u32, None, arg.def(self)).unwrap();
+                            find_xsb(arg)
+                        } else {
+                            let arg = arg.def(self);
+                            let arg = if signed {
+                                let unsigned =
+                                    SpirvType::Integer(bits, false).def(self.span(), self);
+                                self.emit().bitcast(unsigned, None, arg).unwrap()
+                            } else {
+                                arg
+                            };
+                            let arg = self.emit().u_convert(u32, None, arg).unwrap();
+                            let xsb = find_xsb(arg);
+                            let subtrahend = self.constant_u32(self.span(), 32 - bits).def(self);
+                            self.emit().i_sub(u32, None, xsb, subtrahend).unwrap()
+                        }
+                    }
+                    32 => find_xsb(arg.def(self)),
+                    64 => {
+                        let u32_0 = self.constant_int(u32, 0).def(self);
+                        let u32_32 = self.constant_u32(self.span(), 32).def(self);
+
+                        let arg = arg.def(self);
+                        let lower = self.emit().s_convert(u32, None, arg).unwrap();
+                        let higher = self
+                            .emit()
+                            .shift_left_logical(ty, None, arg, u32_32)
+                            .unwrap();
+                        let higher = self.emit().s_convert(u32, None, higher).unwrap();
+
+                        let lower_bits = find_xsb(lower);
+                        let higher_bits = find_xsb(higher);
+
+                        if trailing {
+                            let use_lower = self.emit().i_equal(bool, None, higher, u32_0).unwrap();
+                            let lower_bits =
+                                self.emit().i_add(u32, None, lower_bits, u32_32).unwrap();
+                            self.emit()
+                                .select(u32, None, use_lower, lower_bits, higher_bits)
+                                .unwrap()
+                        } else {
+                            let use_higher = self.emit().i_equal(bool, None, lower, u32_0).unwrap();
+                            let higher_bits =
+                                self.emit().i_add(u32, None, higher_bits, u32_32).unwrap();
+                            self.emit()
+                                .select(u32, None, use_higher, higher_bits, lower_bits)
+                                .unwrap()
+                        }
+                    }
+                    _ => {
+                        let undef = self.undef(ty).def(self);
+                        self.zombie(undef, &format!(
+                            "counting leading / trailing zeros on unsupported {ty:?} bit integer type"
+                        ));
+                        undef
+                    }
+                };
+
                 if non_zero {
-                    find_xsb
+                    converted
                 } else {
+                    let int_0 = self.constant_int(ty, 0).def(self);
+                    let int_bits = self.constant_int(u32, bits as u128).def(self);
                     let is_0 = self
                         .emit()
-                        .i_equal(bool, None, arg.def(self), int_0.def(self))
+                        .i_equal(bool, None, arg.def(self), int_0)
                         .unwrap();
                     self.emit()
-                        .select(ty, None, is_0, int_bits, find_xsb)
+                        .select(u32, None, is_0, int_bits, converted)
                         .unwrap()
                 }
-                .with_type(ty)
+                .with_type(u32)
             }
             _ => self.fatal("counting leading / trailing zeros on a non-integer type"),
         }