feat(const_eval): impl. firstLeadingBit

ErichDonGubler · ErichDonGubler · commit 7d758fe2e9e1 · 2024-07-23T12:46:48.000-04:00
diff --git a/naga/src/proc/constant_evaluator.rs b/naga/src/proc/constant_evaluator.rs
@@ -1233,6 +1233,9 @@ impl<'a> ConstantEvaluator<'a> {
             crate::MathFunction::ReverseBits => {
                 component_wise_concrete_int!(self, span, [arg], |e| { Ok([e.reverse_bits()]) })
             }
+            crate::MathFunction::FirstLeadingBit => {
+                component_wise_concrete_int(self, span, [arg], |ci| Ok(first_leading_bit(ci)))
+            }
 
             fun => Err(ConstantEvaluatorError::NotImplemented(format!(
                 "{fun:?} built-in function"
@@ -2098,6 +2101,102 @@ impl<'a> ConstantEvaluator<'a> {
     }
 }
 
+fn first_leading_bit(concrete_int: ConcreteInt<1>) -> ConcreteInt<1> {
+    // NOTE: Bit indices for this built-in start at 0 at the "right" (or LSB). For example, 1 means
+    // the least significant bit is set. Therefore, an input of 1 would return a right-to-left bit
+    // index of 0.
+    let highest_rtl_bit_index = 32 /* bit */ - 1 /* count to zero index */;
+    match concrete_int {
+        ConcreteInt::I32([e]) => ConcreteInt::I32([{
+            if matches!(e, 0 | -1) {
+                -1
+            } else {
+                let rtl_bit_index = if e.is_positive() {
+                    e.leading_zeros()
+                } else {
+                    // must be negative
+                    e.leading_ones()
+                };
+                (highest_rtl_bit_index - rtl_bit_index)
+                    .try_into()
+                    .expect("bit count outside of 0..32 (!?)")
+            }
+        }]),
+        ConcreteInt::U32([e]) => ConcreteInt::U32([{
+            if e == 0 {
+                u32::MAX
+            } else {
+                let rtl_bit_index = e.leading_zeros();
+                highest_rtl_bit_index - rtl_bit_index
+            }
+        }]),
+    }
+}
+
+#[test]
+fn first_leading_bit_smoke() {
+    assert_eq!(
+        first_leading_bit(ConcreteInt::I32([-1])),
+        ConcreteInt::I32([-1])
+    );
+    assert_eq!(
+        first_leading_bit(ConcreteInt::I32([0])),
+        ConcreteInt::I32([-1])
+    );
+    assert_eq!(
+        first_leading_bit(ConcreteInt::I32([1])),
+        ConcreteInt::I32([0])
+    );
+    assert_eq!(
+        first_leading_bit(ConcreteInt::I32([-2])),
+        ConcreteInt::I32([0])
+    );
+    assert_eq!(
+        first_leading_bit(ConcreteInt::I32([1234 + 4567])),
+        ConcreteInt::I32([12])
+    );
+    assert_eq!(
+        first_leading_bit(ConcreteInt::I32([i32::MAX])),
+        ConcreteInt::I32([30])
+    );
+    assert_eq!(
+        first_leading_bit(ConcreteInt::I32([i32::MIN])),
+        ConcreteInt::I32([30])
+    );
+    // NOTE: Ignore the sign bit, which is a separate (above) case.
+    for idx in 0..(32 - 1) {
+        assert_eq!(
+            first_leading_bit(ConcreteInt::I32([1 << idx])),
+            ConcreteInt::I32([idx])
+        );
+    }
+    for idx in 1..(32 - 1) {
+        assert_eq!(
+            first_leading_bit(ConcreteInt::I32([(1 << idx) * -1])),
+            ConcreteInt::I32([idx - 1])
+        );
+    }
+
+    assert_eq!(
+        first_leading_bit(ConcreteInt::U32([0])),
+        ConcreteInt::U32([u32::MAX])
+    );
+    assert_eq!(
+        first_leading_bit(ConcreteInt::U32([1])),
+        ConcreteInt::U32([0])
+    );
+    assert_eq!(
+        first_leading_bit(ConcreteInt::U32([u32::MAX])),
+        ConcreteInt::U32([31])
+    );
+    for idx in 0..32 {
+        assert_eq!(
+            first_leading_bit(ConcreteInt::U32([1 << idx])),
+            ConcreteInt::U32([idx])
+        )
+    }
+}
+
 /// Trait for conversions of abstract values to concrete types.
 trait TryFromAbstract<T>: Sized {
     /// Convert an abstract literal `value` to `Self`.
diff --git a/naga/tests/out/glsl/math-functions.main.Fragment.glsl b/naga/tests/out/glsl/math-functions.main.Fragment.glsl
@@ -65,10 +65,8 @@ void main() {
     ivec4 sign_b = ivec4(-1, -1, -1, -1);
     vec4 sign_d = vec4(-1.0, -1.0, -1.0, -1.0);
     int const_dot = ( + ivec2(0).x * ivec2(0).x + ivec2(0).y * ivec2(0).y);
-    uint first_leading_bit_abs = uint(findMSB(0u));
-    int flb_a = findMSB(-1);
-    ivec2 flb_b = findMSB(ivec2(-1));
-    uvec2 flb_c = uvec2(findMSB(uvec2(1u)));
+    ivec2 flb_b = ivec2(-1, -1);
+    uvec2 flb_c = uvec2(0u, 0u);
     int ftb_a = findLSB(-1);
     uint ftb_b = uint(findLSB(1u));
     ivec2 ftb_c = findLSB(ivec2(-1));
diff --git a/naga/tests/out/hlsl/math-functions.hlsl b/naga/tests/out/hlsl/math-functions.hlsl
@@ -79,10 +79,8 @@ void main()
     int4 sign_b = int4(-1, -1, -1, -1);
     float4 sign_d = float4(-1.0, -1.0, -1.0, -1.0);
     int const_dot = dot(ZeroValueint2(), ZeroValueint2());
-    uint first_leading_bit_abs = firstbithigh(0u);
-    int flb_a = asint(firstbithigh(-1));
-    int2 flb_b = asint(firstbithigh((-1).xx));
-    uint2 flb_c = firstbithigh((1u).xx);
+    int2 flb_b = int2(-1, -1);
+    uint2 flb_c = uint2(0u, 0u);
     int ftb_a = asint(firstbitlow(-1));
     uint ftb_b = firstbitlow(1u);
     int2 ftb_c = asint(firstbitlow((-1).xx));
diff --git a/naga/tests/out/msl/math-functions.msl b/naga/tests/out/msl/math-functions.msl
@@ -67,12 +67,8 @@ fragment void main_(
     metal::int4 sign_b = metal::int4(-1, -1, -1, -1);
     metal::float4 sign_d = metal::float4(-1.0, -1.0, -1.0, -1.0);
     int const_dot = ( + metal::int2 {}.x * metal::int2 {}.x + metal::int2 {}.y * metal::int2 {}.y);
-    uint first_leading_bit_abs = metal::select(31 - metal::clz(0u), uint(-1), 0u == 0 || 0u == -1);
-    int flb_a = metal::select(31 - metal::clz(metal::select(-1, ~-1, -1 < 0)), int(-1), -1 == 0 || -1 == -1);
-    metal::int2 _e29 = metal::int2(-1);
-    metal::int2 flb_b = metal::select(31 - metal::clz(metal::select(_e29, ~_e29, _e29 < 0)), int2(-1), _e29 == 0 || _e29 == -1);
-    metal::uint2 _e32 = metal::uint2(1u);
-    metal::uint2 flb_c = metal::select(31 - metal::clz(_e32), uint2(-1), _e32 == 0 || _e32 == -1);
+    metal::int2 flb_b = metal::int2(-1, -1);
+    metal::uint2 flb_c = metal::uint2(0u, 0u);
     int ftb_a = (((metal::ctz(-1) + 1) % 33) - 1);
     uint ftb_b = (((metal::ctz(1u) + 1) % 33) - 1);
     metal::int2 ftb_c = (((metal::ctz(metal::int2(-1)) + 1) % 33) - 1);
diff --git a/naga/tests/out/spv/math-functions.spvasm b/naga/tests/out/spv/math-functions.spvasm
@@ -1,7 +1,7 @@
 ; SPIR-V
 ; Version: 1.1
 ; Generator: rspirv
-; Bound: 96
+; Bound: 94
 OpCapability Shader
 %1 = OpExtInstImport "GLSL.std.450"
 OpMemoryModel Logical GLSL450
@@ -40,77 +40,75 @@ OpMemberDecorate %15 1 Offset 16
 %24 = OpConstant  %4  -1.0
 %25 = OpConstantComposite  %3  %24 %24 %24 %24
 %26 = OpConstantNull  %7
-%27 = OpConstant  %9  0
+%27 = OpConstant  %9  4294967295
 %28 = OpConstantComposite  %7  %22 %22
-%29 = OpConstant  %9  1
+%29 = OpConstant  %9  0
 %30 = OpConstantComposite  %8  %29 %29
-%31 = OpConstant  %9  32
-%32 = OpConstant  %6  32
-%33 = OpConstant  %6  0
-%34 = OpConstantComposite  %8  %31 %31
-%35 = OpConstantComposite  %7  %32 %32
-%36 = OpConstantComposite  %8  %27 %27
-%37 = OpConstantComposite  %7  %33 %33
-%38 = OpConstant  %9  31
-%39 = OpConstantComposite  %8  %38 %38
-%40 = OpConstant  %6  2
-%41 = OpConstant  %4  2.0
-%42 = OpConstantComposite  %10  %19 %41
-%43 = OpConstant  %6  3
-%44 = OpConstant  %6  4
-%45 = OpConstantComposite  %7  %43 %44
-%46 = OpConstant  %4  1.5
-%47 = OpConstantComposite  %10  %46 %46
-%48 = OpConstantComposite  %3  %46 %46 %46 %46
-%55 = OpConstantComposite  %3  %19 %19 %19 %19
-%58 = OpConstantNull  %6
+%31 = OpConstant  %9  1
+%32 = OpConstantComposite  %7  %22 %22
+%33 = OpConstantComposite  %8  %31 %31
+%34 = OpConstant  %9  32
+%35 = OpConstant  %6  32
+%36 = OpConstant  %6  0
+%37 = OpConstantComposite  %8  %34 %34
+%38 = OpConstantComposite  %7  %35 %35
+%39 = OpConstantComposite  %7  %36 %36
+%40 = OpConstant  %9  31
+%41 = OpConstantComposite  %8  %40 %40
+%42 = OpConstant  %6  2
+%43 = OpConstant  %4  2.0
+%44 = OpConstantComposite  %10  %19 %43
+%45 = OpConstant  %6  3
+%46 = OpConstant  %6  4
+%47 = OpConstantComposite  %7  %45 %46
+%48 = OpConstant  %4  1.5
+%49 = OpConstantComposite  %10  %48 %48
+%50 = OpConstantComposite  %3  %48 %48 %48 %48
+%57 = OpConstantComposite  %3  %19 %19 %19 %19
+%60 = OpConstantNull  %6
 %17 = OpFunction  %2  None %18
 %16 = OpLabel
-OpBranch %49
-%49 = OpLabel
-%50 = OpExtInst  %4  %1 Degrees %19
-%51 = OpExtInst  %4  %1 Radians %19
-%52 = OpExtInst  %3  %1 Degrees %21
-%53 = OpExtInst  %3  %1 Radians %21
-%54 = OpExtInst  %3  %1 FClamp %21 %21 %55
-%56 = OpExtInst  %3  %1 Refract %21 %21 %19
-%59 = OpCompositeExtract  %6  %26 0
-%60 = OpCompositeExtract  %6  %26 0
-%61 = OpIMul  %6  %59 %60
-%62 = OpIAdd  %6  %58 %61
-%63 = OpCompositeExtract  %6  %26 1
-%64 = OpCompositeExtract  %6  %26 1
-%65 = OpIMul  %6  %63 %64
-%57 = OpIAdd  %6  %62 %65
-%66 = OpExtInst  %9  %1 FindUMsb %27
-%67 = OpExtInst  %6  %1 FindSMsb %22
-%68 = OpExtInst  %7  %1 FindSMsb %28
-%69 = OpExtInst  %8  %1 FindUMsb %30
-%70 = OpExtInst  %6  %1 FindILsb %22
-%71 = OpExtInst  %9  %1 FindILsb %29
-%72 = OpExtInst  %7  %1 FindILsb %28
-%73 = OpExtInst  %8  %1 FindILsb %30
-%74 = OpExtInst  %4  %1 Ldexp %19 %40
-%75 = OpExtInst  %10  %1 Ldexp %42 %45
-%76 = OpExtInst  %11  %1 ModfStruct %46
-%77 = OpExtInst  %11  %1 ModfStruct %46
-%78 = OpCompositeExtract  %4  %77 0
-%79 = OpExtInst  %11  %1 ModfStruct %46
-%80 = OpCompositeExtract  %4  %79 1
-%81 = OpExtInst  %12  %1 ModfStruct %47
-%82 = OpExtInst  %13  %1 ModfStruct %48
-%83 = OpCompositeExtract  %3  %82 1
-%84 = OpCompositeExtract  %4  %83 0
-%85 = OpExtInst  %12  %1 ModfStruct %47
-%86 = OpCompositeExtract  %10  %85 0
-%87 = OpCompositeExtract  %4  %86 1
-%88 = OpExtInst  %14  %1 FrexpStruct %46
-%89 = OpExtInst  %14  %1 FrexpStruct %46
-%90 = OpCompositeExtract  %4  %89 0
-%91 = OpExtInst  %14  %1 FrexpStruct %46
-%92 = OpCompositeExtract  %6  %91 1
-%93 = OpExtInst  %15  %1 FrexpStruct %48
-%94 = OpCompositeExtract  %5  %93 1
-%95 = OpCompositeExtract  %6  %94 0
+OpBranch %51
+%51 = OpLabel
+%52 = OpExtInst  %4  %1 Degrees %19
+%53 = OpExtInst  %4  %1 Radians %19
+%54 = OpExtInst  %3  %1 Degrees %21
+%55 = OpExtInst  %3  %1 Radians %21
+%56 = OpExtInst  %3  %1 FClamp %21 %21 %57
+%58 = OpExtInst  %3  %1 Refract %21 %21 %19
+%61 = OpCompositeExtract  %6  %26 0
+%62 = OpCompositeExtract  %6  %26 0
+%63 = OpIMul  %6  %61 %62
+%64 = OpIAdd  %6  %60 %63
+%65 = OpCompositeExtract  %6  %26 1
+%66 = OpCompositeExtract  %6  %26 1
+%67 = OpIMul  %6  %65 %66
+%59 = OpIAdd  %6  %64 %67
+%68 = OpExtInst  %6  %1 FindILsb %22
+%69 = OpExtInst  %9  %1 FindILsb %31
+%70 = OpExtInst  %7  %1 FindILsb %32
+%71 = OpExtInst  %8  %1 FindILsb %33
+%72 = OpExtInst  %4  %1 Ldexp %19 %42
+%73 = OpExtInst  %10  %1 Ldexp %44 %47
+%74 = OpExtInst  %11  %1 ModfStruct %48
+%75 = OpExtInst  %11  %1 ModfStruct %48
+%76 = OpCompositeExtract  %4  %75 0
+%77 = OpExtInst  %11  %1 ModfStruct %48
+%78 = OpCompositeExtract  %4  %77 1
+%79 = OpExtInst  %12  %1 ModfStruct %49
+%80 = OpExtInst  %13  %1 ModfStruct %50
+%81 = OpCompositeExtract  %3  %80 1
+%82 = OpCompositeExtract  %4  %81 0
+%83 = OpExtInst  %12  %1 ModfStruct %49
+%84 = OpCompositeExtract  %10  %83 0
+%85 = OpCompositeExtract  %4  %84 1
+%86 = OpExtInst  %14  %1 FrexpStruct %48
+%87 = OpExtInst  %14  %1 FrexpStruct %48
+%88 = OpCompositeExtract  %4  %87 0
+%89 = OpExtInst  %14  %1 FrexpStruct %48
+%90 = OpCompositeExtract  %6  %89 1
+%91 = OpExtInst  %15  %1 FrexpStruct %50
+%92 = OpCompositeExtract  %5  %91 1
+%93 = OpCompositeExtract  %6  %92 0
 OpReturn
 OpFunctionEnd
diff --git a/naga/tests/out/wgsl/math-functions.wgsl b/naga/tests/out/wgsl/math-functions.wgsl
@@ -10,10 +10,8 @@ fn main() {
     let sign_b = vec4<i32>(-1i, -1i, -1i, -1i);
     let sign_d = vec4<f32>(-1f, -1f, -1f, -1f);
     let const_dot = dot(vec2<i32>(), vec2<i32>());
-    let first_leading_bit_abs = firstLeadingBit(0u);
-    let flb_a = firstLeadingBit(-1i);
-    let flb_b = firstLeadingBit(vec2(-1i));
-    let flb_c = firstLeadingBit(vec2(1u));
+    let flb_b = vec2<i32>(-1i, -1i);
+    let flb_c = vec2<u32>(0u, 0u);
     let ftb_a = firstTrailingBit(-1i);
     let ftb_b = firstTrailingBit(1u);
     let ftb_c = firstTrailingBit(vec2(-1i));
