leading_zeros: add support for leading_zeros and trailing_zeros, limited to u32/i32

Author: Firestar99

crates/rustc_codegen_spirv/src/builder/ext_inst.rs

@@ -1,8 +1,8 @@
 use super::Builder;
 use crate::builder_spirv::{SpirvValue, SpirvValueExt};
 use crate::custom_insts;
+use rspirv::dr::Operand;
 use rspirv::spirv::{GLOp, Word};
-use rspirv::{dr::Operand, spirv::Capability};
 
 const GLSL_STD_450: &str = "GLSL.std.450";
 
@@ -13,7 +13,6 @@ pub struct ExtInst {
     custom: Option<Word>,
 
     glsl: Option<Word>,
-    integer_functions_2_intel: bool,
 }
 
 impl ExtInst {
@@ -38,32 +37,11 @@ impl ExtInst {
             id
         }
     }
-
-    pub fn require_integer_functions_2_intel(&mut self, bx: &Builder<'_, '_>, to_zombie: Word) {
-        if !self.integer_functions_2_intel {
-            self.integer_functions_2_intel = true;
-            if !bx
-                .builder
-                .has_capability(Capability::IntegerFunctions2INTEL)
-            {
-                bx.zombie(to_zombie, "capability IntegerFunctions2INTEL is required");
-            }
-            if !bx
-                .builder
-                .has_extension(bx.sym.spv_intel_shader_integer_functions2)
-            {
-                bx.zombie(
-                    to_zombie,
-                    "extension SPV_INTEL_shader_integer_functions2 is required",
-                );
-            }
-        }
-    }
 }
 
 impl<'a, 'tcx> Builder<'a, 'tcx> {
     pub fn custom_inst(
-        &mut self,
+        &self,
         result_type: Word,
         inst: custom_insts::CustomInst<Operand>,
     ) -> SpirvValue {
@@ -80,12 +58,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             .with_type(result_type)
     }
 
-    pub fn gl_op(
-        &mut self,
-        op: GLOp,
-        result_type: Word,
-        args: impl AsRef<[SpirvValue]>,
-    ) -> SpirvValue {
+    pub fn gl_op(&self, op: GLOp, result_type: Word, args: impl AsRef<[SpirvValue]>) -> SpirvValue {
         let args = args.as_ref();
         let glsl = self.ext_inst.borrow_mut().import_glsl(self);
         self.emit()

crates/rustc_codegen_spirv/src/builder/intrinsics.rs

@@ -211,35 +211,12 @@ impl<'a, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'a, 'tcx> {
                 self.rotate(val, shift, is_left)
             }
 
-            // TODO: Do we want to manually implement these instead of using intel instructions?
-            sym::ctlz | sym::ctlz_nonzero => {
-                let result = self
-                    .emit()
-                    .u_count_leading_zeros_intel(
-                        args[0].immediate().ty,
-                        None,
-                        args[0].immediate().def(self),
-                    )
-                    .unwrap();
-                self.ext_inst
-                    .borrow_mut()
-                    .require_integer_functions_2_intel(self, result);
-                result.with_type(args[0].immediate().ty)
-            }
-            sym::cttz | sym::cttz_nonzero => {
-                let result = self
-                    .emit()
-                    .u_count_trailing_zeros_intel(
-                        args[0].immediate().ty,
-                        None,
-                        args[0].immediate().def(self),
-                    )
-                    .unwrap();
-                self.ext_inst
-                    .borrow_mut()
-                    .require_integer_functions_2_intel(self, result);
-                result.with_type(args[0].immediate().ty)
+            sym::ctlz => self.count_leading_trailing_zeros(args[0].immediate(), false, false),
+            sym::ctlz_nonzero => {
+                self.count_leading_trailing_zeros(args[0].immediate(), false, true)
             }
+            sym::cttz => self.count_leading_trailing_zeros(args[0].immediate(), true, false),
+            sym::cttz_nonzero => self.count_leading_trailing_zeros(args[0].immediate(), true, true),
 
             sym::ctpop => self
                 .emit()
@@ -398,6 +375,50 @@ impl<'a, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'a, 'tcx> {
 }
 
 impl Builder<'_, '_> {
+    pub fn count_leading_trailing_zeros(
+        &self,
+        arg: SpirvValue,
+        trailing: bool,
+        non_zero: bool,
+    ) -> 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);
+                let bool = SpirvType::Bool.def(self.span(), self);
+
+                let gl_op = if trailing {
+                    // rust is always unsigned
+                    GLOp::FindILsb
+                } else {
+                    GLOp::FindUMsb
+                };
+
+                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();
+                if non_zero {
+                    find_xsb
+                } else {
+                    let is_0 = self
+                        .emit()
+                        .i_equal(bool, None, arg.def(self), int_0.def(self))
+                        .unwrap();
+                    self.emit()
+                        .select(ty, None, is_0, int_bits, find_xsb)
+                        .unwrap()
+                }
+                .with_type(ty)
+            }
+            _ => self.fatal("counting leading / trailing zeros on a non-integer type"),
+        }
+    }
+
     pub fn abort_with_kind_and_message_debug_printf(
         &mut self,
         kind: &str,

crates/rustc_codegen_spirv/src/symbols.rs

@@ -21,7 +21,6 @@ pub struct Symbols {
     pub spirv: Symbol,
     pub libm: Symbol,
     pub entry_point_name: Symbol,
-    pub spv_intel_shader_integer_functions2: Symbol,
     pub spv_khr_vulkan_memory_model: Symbol,
 
     descriptor_set: Symbol,
@@ -411,9 +410,6 @@ impl Symbols {
             spirv: Symbol::intern("spirv"),
             libm: Symbol::intern("libm"),
             entry_point_name: Symbol::intern("entry_point_name"),
-            spv_intel_shader_integer_functions2: Symbol::intern(
-                "SPV_INTEL_shader_integer_functions2",
-            ),
             spv_khr_vulkan_memory_model: Symbol::intern("SPV_KHR_vulkan_memory_model"),
 
             descriptor_set: Symbol::intern("descriptor_set"),

tests/ui/lang/u32/trailing_leading_bits.rs

@@ -0,0 +1,37 @@
+// Test all trailing and leading zeros. No need to test ones, they just call the zero variant with !value
+
+// build-pass
+
+use spirv_std::spirv;
+
+#[spirv(fragment)]
+pub fn leading_zeros_u32(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buffer: &u32,
+    out: &mut u32,
+) {
+    *out = u32::leading_zeros(*buffer);
+}
+
+#[spirv(fragment)]
+pub fn trailing_zeros_u32(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buffer: &u32,
+    out: &mut u32,
+) {
+    *out = u32::trailing_zeros(*buffer);
+}
+
+#[spirv(fragment)]
+pub fn leading_zeros_i32(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buffer: &i32,
+    out: &mut u32,
+) {
+    *out = i32::leading_zeros(*buffer);
+}
+
+#[spirv(fragment)]
+pub fn trailing_zeros_i32(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buffer: &i32,
+    out: &mut u32,
+) {
+    *out = i32::trailing_zeros(*buffer);
+}