[wgsl-in] Allow sign() to take int argument

fornwall · Sep 5, 2023 · c69e5c1 · c69e5c1
1 parent 5329aa2
commit c69e5c1
Show file tree

Hide file tree

Showing 50 changed files with 706 additions and 105 deletions.
diff --git a/src/back/msl/keywords.rs b/src/back/msl/keywords.rs
@@ -216,4 +216,5 @@ pub const RESERVED: &[&str] = &[
     "clamped_lod_e",
     super::writer::FREXP_FUNCTION,
     super::writer::MODF_FUNCTION,
+    super::writer::NAGA_ISIGN_FUNCTION,
 ];
diff --git a/src/back/msl/writer.rs b/src/back/msl/writer.rs
@@ -34,6 +34,7 @@ const RAY_QUERY_FUN_MAP_INTERSECTION: &str = "_map_intersection_type";
 
 pub(crate) const MODF_FUNCTION: &str = "naga_modf";
 pub(crate) const FREXP_FUNCTION: &str = "naga_frexp";
+pub(crate) const NAGA_ISIGN_FUNCTION: &str = "naga_isign";
 
 /// Write the Metal name for a Naga numeric type: scalar, vector, or matrix.
 ///
@@ -1647,8 +1648,9 @@ impl<W: Write> Writer<W> {
             } => {
                 use crate::MathFunction as Mf;
 
-                let scalar_argument = match *context.resolve_type(arg) {
-                    crate::TypeInner::Scalar { .. } => true,
+                let arg_type = context.resolve_type(arg);
+                let scalar_argument = match arg_type {
+                    &crate::TypeInner::Scalar { .. } => true,
                     _ => false,
                 };
 
@@ -1713,7 +1715,17 @@ impl<W: Write> Writer<W> {
                     Mf::Reflect => "reflect",
                     Mf::Refract => "refract",
                     // computational
-                    Mf::Sign => "sign",
+                    Mf::Sign => match arg_type {
+                        &crate::TypeInner::Scalar {
+                            kind: crate::ScalarKind::Sint,
+                            ..
+                        }
+                        | crate::TypeInner::Vector {
+                            kind: crate::ScalarKind::Sint,
+                            ..
+                        } => NAGA_ISIGN_FUNCTION,
+                        _ => "sign",
+                    },
                     Mf::Fma => "fma",
                     Mf::Mix => "mix",
                     Mf::Step => "step",
@@ -1816,7 +1828,7 @@ impl<W: Write> Writer<W> {
                     write!(self.out, "((")?;
                     self.put_expression(arg, context, false)?;
                     write!(self.out, ") * 57.295779513082322865)")?;
-                } else if fun == Mf::Modf || fun == Mf::Frexp {
+                } else if fun == Mf::Modf || fun == Mf::Frexp || fun_name == NAGA_ISIGN_FUNCTION {
                     write!(self.out, "{fun_name}")?;
                     self.put_call_parameters(iter::once(arg), context)?;
                 } else {
@@ -3091,6 +3103,7 @@ impl<W: Write> Writer<W> {
 
         self.write_type_defs(module)?;
         self.write_global_constants(module)?;
+        self.write_polyfills()?;
         self.write_functions(module, info, options, pipeline_options)
     }
 
@@ -4103,6 +4116,26 @@ impl<W: Write> Writer<W> {
         Ok(info)
     }
 
+    fn write_polyfills(&mut self) -> Result<(), Error> {
+        writeln!(self.out)?;
+        for size in 1..5 {
+            let tmp;
+            let type_name = if size == 1 {
+                "int"
+            } else {
+                tmp = format!("int{size}");
+                &tmp
+            };
+            writeln!(
+                self.out,
+                "{type_name} {NAGA_ISIGN_FUNCTION}({type_name} arg) {{
+    return {NAMESPACE}::select({NAMESPACE}::select({type_name}(-1), {type_name}(1), (arg > 0)), 0, (arg == 0));
+}}"
+            )?;
+        }
+        Ok(())
+    }
+
     fn write_barrier(&mut self, flags: crate::Barrier, level: back::Level) -> BackendResult {
         // Note: OR-ring bitflags requires `__HAVE_MEMFLAG_OPERATORS__`,
         // so we try to avoid it here.

diff --git a/src/valid/expression.rs b/src/valid/expression.rs
@@ -976,7 +976,6 @@ impl super::Validator {
                     | Mf::Log
                     | Mf::Log2
                     | Mf::Length
-                    | Mf::Sign
                     | Mf::Sqrt
                     | Mf::InverseSqrt => {
                         if arg1_ty.is_some() | arg2_ty.is_some() | arg3_ty.is_some() {
@@ -992,6 +991,22 @@ impl super::Validator {
                             _ => return Err(ExpressionError::InvalidArgumentType(fun, 0, arg)),
                         }
                     }
+                    Mf::Sign => {
+                        if arg1_ty.is_some() | arg2_ty.is_some() | arg3_ty.is_some() {
+                            return Err(ExpressionError::WrongArgumentCount(fun));
+                        }
+                        match *arg_ty {
+                            Ti::Scalar {
+                                kind: Sk::Float | Sk::Sint,
+                                ..
+                            }
+                            | Ti::Vector {
+                                kind: Sk::Float | Sk::Sint,
+                                ..
+                            } => {}
+                            _ => return Err(ExpressionError::InvalidArgumentType(fun, 0, arg)),
+                        }
+                    }
                     Mf::Atan2 | Mf::Pow | Mf::Distance | Mf::Step => {
                         let arg1_ty = match (arg1_ty, arg2_ty, arg3_ty) {
                             (Some(ty1), None, None) => ty1,

diff --git a/tests/in/math-functions.wgsl b/tests/in/math-functions.wgsl
@@ -8,6 +8,7 @@ fn main() {
     let d = radians(v);
     let e = saturate(v);
     let g = refract(v, v, f);
+    let h = sign(-1);
     let const_dot = dot(vec2<i32>(), vec2<i32>());
     let first_leading_bit_abs = firstLeadingBit(abs(0u));
     let flb_a = firstLeadingBit(-1);

diff --git a/tests/out/glsl/math-functions.main.Fragment.glsl b/tests/out/glsl/math-functions.main.Fragment.glsl
@@ -62,6 +62,7 @@ void main() {
     vec4 d = radians(v);
     vec4 e = clamp(v, vec4(0.0), vec4(1.0));
     vec4 g = refract(v, v, 1.0);
+    int h = sign(-1);
     int const_dot = ( + ivec2(0).x * ivec2(0).x + ivec2(0).y * ivec2(0).y);
     uint first_leading_bit_abs = uint(findMSB(uint(abs(int(0u)))));
     int flb_a = findMSB(-1);
@@ -81,8 +82,8 @@ void main() {
     ivec2 ctz_h = ivec2(min(uvec2(findLSB(ivec2(1))), uvec2(32u)));
     int clz_a = (-1 < 0 ? 0 : 31 - findMSB(-1));
     uint clz_b = uint(31 - findMSB(1u));
-    ivec2 _e58 = ivec2(-1);
-    ivec2 clz_c = mix(ivec2(31) - findMSB(_e58), ivec2(0), lessThan(_e58, ivec2(0)));
+    ivec2 _e60 = ivec2(-1);
+    ivec2 clz_c = mix(ivec2(31) - findMSB(_e60), ivec2(0), lessThan(_e60, ivec2(0)));
     uvec2 clz_d = uvec2(ivec2(31) - findMSB(uvec2(1u)));
     float lde_a = ldexp(1.0, 2);
     vec2 lde_b = ldexp(vec2(1.0, 2.0), ivec2(3, 4));

diff --git a/tests/out/hlsl/math-functions.hlsl b/tests/out/hlsl/math-functions.hlsl
@@ -72,6 +72,7 @@ void main()
     float4 d = radians(v);
     float4 e = saturate(v);
     float4 g = refract(v, v, 1.0);
+    int h = sign(-1);
     int const_dot = dot((int2)0, (int2)0);
     uint first_leading_bit_abs = firstbithigh(abs(0u));
     int flb_a = asint(firstbithigh(-1));
@@ -91,8 +92,8 @@ void main()
     int2 ctz_h = asint(min((32u).xx, firstbitlow((1).xx)));
     int clz_a = (-1 < 0 ? 0 : 31 - asint(firstbithigh(-1)));
     uint clz_b = (31u - firstbithigh(1u));
-    int2 _expr58 = (-1).xx;
-    int2 clz_c = (_expr58 < (0).xx ? (0).xx : (31).xx - asint(firstbithigh(_expr58)));
+    int2 _expr60 = (-1).xx;
+    int2 clz_c = (_expr60 < (0).xx ? (0).xx : (31).xx - asint(firstbithigh(_expr60)));
     uint2 clz_d = ((31u).xx - firstbithigh((1u).xx));
     float lde_a = ldexp(1.0, 2);
     float2 lde_b = ldexp(float2(1.0, 2.0), int2(3, 4));

diff --git a/tests/out/msl/access.msl b/tests/out/msl/access.msl
@@ -58,6 +58,19 @@ struct type_26 {
     metal::float4 inner[2];
 };
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 void test_matrix_within_struct_accesses(
     constant Baz& baz
 ) {

diff --git a/tests/out/msl/array-in-ctor.msl b/tests/out/msl/array-in-ctor.msl
@@ -11,6 +11,19 @@ struct Ah {
     type_1 inner;
 };
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 kernel void cs_main(
   device Ah const& ah [[user(fake0)]]
 ) {

diff --git a/tests/out/msl/array-in-function-return-type.msl b/tests/out/msl/array-in-function-return-type.msl
@@ -8,6 +8,19 @@ struct type_1 {
     float inner[2];
 };
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 type_1 ret_array(
 ) {
     return type_1 {1.0, 2.0};

diff --git a/tests/out/msl/atomicOps.msl b/tests/out/msl/atomicOps.msl
@@ -12,6 +12,19 @@ struct Struct {
     type_2 atomic_arr;
 };
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 struct cs_mainInput {
 };
 kernel void cs_main(

diff --git a/tests/out/msl/binding-arrays.msl b/tests/out/msl/binding-arrays.msl
@@ -17,6 +17,19 @@ struct FragmentIn {
     uint index;
 };
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 struct main_Input {
     uint index [[user(loc0), flat]];
 };

diff --git a/tests/out/msl/bitcast.msl b/tests/out/msl/bitcast.msl
@@ -5,6 +5,19 @@
 using metal::uint;
 
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 kernel void main_(
 ) {
     metal::int2 i2_ = {};

diff --git a/tests/out/msl/bits.msl b/tests/out/msl/bits.msl
@@ -5,6 +5,19 @@
 using metal::uint;
 
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 kernel void main_(
 ) {
     int i = {};

diff --git a/tests/out/msl/boids.msl b/tests/out/msl/boids.msl
@@ -28,6 +28,19 @@ struct Particles {
 };
 constant uint NUM_PARTICLES = 1500u;
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 struct main_Input {
 };
 kernel void main_(

diff --git a/tests/out/msl/bounds-check-image-restrict.msl b/tests/out/msl/bounds-check-image-restrict.msl
@@ -5,6 +5,19 @@
 using metal::uint;
 
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 metal::float4 test_textureLoad_1d(
     int coords,
     int level,

diff --git a/tests/out/msl/bounds-check-image-rzsw.msl b/tests/out/msl/bounds-check-image-rzsw.msl
@@ -11,6 +11,19 @@ struct DefaultConstructible {
 };
 
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 metal::float4 test_textureLoad_1d(
     int coords,
     int level,

diff --git a/tests/out/msl/bounds-check-restrict.msl b/tests/out/msl/bounds-check-restrict.msl
@@ -20,6 +20,19 @@ struct Globals {
     type_4 d;
 };
 
+int naga_isign(int arg) {
+    return metal::select(metal::select(int(-1), int(1), (arg > 0)), 0, (arg == 0));
+}
+int2 naga_isign(int2 arg) {
+    return metal::select(metal::select(int2(-1), int2(1), (arg > 0)), 0, (arg == 0));
+}
+int3 naga_isign(int3 arg) {
+    return metal::select(metal::select(int3(-1), int3(1), (arg > 0)), 0, (arg == 0));
+}
+int4 naga_isign(int4 arg) {
+    return metal::select(metal::select(int4(-1), int4(1), (arg > 0)), 0, (arg == 0));
+}
+
 float index_array(
     int i,
     device Globals const& globals,