[msl-out] Emit sign(i32) code inline

fornwall · Sep 6, 2023 · e209bd9 · e209bd9
1 parent f6a2a91
commit e209bd9
Show file tree

Hide file tree

Showing 44 changed files with 42 additions and 185 deletions.
diff --git a/src/back/msl/keywords.rs b/src/back/msl/keywords.rs
@@ -216,5 +216,4 @@ pub const RESERVED: &[&str] = &[
     "clamped_lod_e",
     super::writer::FREXP_FUNCTION,
     super::writer::MODF_FUNCTION,
-    super::writer::ISIGN_FUNCTION,
 ];
diff --git a/src/back/msl/writer.rs b/src/back/msl/writer.rs
@@ -34,7 +34,6 @@ 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 ISIGN_FUNCTION: &str = "naga_isign";
 
 /// Write the Metal name for a Naga numeric type: scalar, vector, or matrix.
 ///
@@ -1185,6 +1184,31 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
+    /// Emit code for the sign(i32) expression.
+    ///
+    fn put_isign(
+        &mut self,
+        arg: Handle<crate::Expression>,
+        context: &ExpressionContext,
+    ) -> BackendResult {
+        write!(self.out, "{NAMESPACE}::select({NAMESPACE}::select(")?;
+        match context.resolve_type(arg) {
+            &crate::TypeInner::Vector { size, .. } => {
+                let size = back::vector_size_str(size);
+                write!(self.out, "int{size}(-1), int{size}(1)")?;
+            }
+            _ => {
+                write!(self.out, "-1, 1")?;
+            }
+        }
+        write!(self.out, ", (")?;
+        self.put_expression(arg, context, true)?;
+        write!(self.out, " > 0)), 0, (")?;
+        self.put_expression(arg, context, true)?;
+        write!(self.out, " == 0))")?;
+        Ok(())
+    }
+
     fn put_const_expression(
         &mut self,
         expr_handle: Handle<crate::Expression>,
@@ -1716,7 +1740,9 @@ impl<W: Write> Writer<W> {
                     Mf::Refract => "refract",
                     // computational
                     Mf::Sign => match arg_type.scalar_kind() {
-                        Some(crate::ScalarKind::Sint) => ISIGN_FUNCTION,
+                        Some(crate::ScalarKind::Sint) => {
+                            return self.put_isign(arg, context);
+                        }
                         _ => "sign",
                     },
                     Mf::Fma => "fma",
@@ -1821,7 +1847,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 || fun_name == ISIGN_FUNCTION {
+                } else if fun == Mf::Modf || fun == Mf::Frexp {
                     write!(self.out, "{fun_name}")?;
                     self.put_call_parameters(iter::once(arg), context)?;
                 } else {
@@ -2428,6 +2454,16 @@ impl<W: Write> Writer<W> {
                     crate::MathFunction::FindMsb => {
                         self.need_bake_expressions.insert(arg);
                     }
+                    crate::MathFunction::Sign => {
+                        // WGSL's `sign` function works also on signed ints, but Metal's only
+                        // works on floating-point vectors, so we emit inline code for
+                        // integer `sign` calls. But that code uses each argument 2
+                        // times, so to avoid duplicated evaluation, we must bake the argument.
+                        let inner = context.resolve_type(expr_handle);
+                        if inner.scalar_kind() == Some(crate::ScalarKind::Sint) {
+                            self.need_bake_expressions.insert(arg);
+                        }
+                    }
                     _ => {}
                 }
             }
@@ -3096,7 +3132,6 @@ 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)
     }
 
@@ -4109,16 +4144,6 @@ impl<W: Write> Writer<W> {
         Ok(info)
     }
 
-    fn write_polyfills(&mut self) -> Result<(), Error> {
-        writeln!(
-            self.out,
-            "\ntemplate <typename T> inline T {ISIGN_FUNCTION}(T arg) {{
-    return {NAMESPACE}::select({NAMESPACE}::select(T(-1), T(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/tests/out/msl/access.msl b/tests/out/msl/access.msl
@@ -58,10 +58,6 @@ struct type_26 {
     metal::float4 inner[2];
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +11,6 @@ struct Ah {
     type_1 inner;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +8,6 @@ struct type_1 {
     float inner[2];
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +12,6 @@ struct Struct {
     type_2 atomic_arr;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +17,6 @@ struct FragmentIn {
     uint index;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +5,6 @@
 using metal::uint;
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +5,6 @@
 using metal::uint;
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +28,6 @@ struct Particles {
 };
 constant uint NUM_PARTICLES = 1500u;
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +5,6 @@
 using metal::uint;
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +11,6 @@ struct DefaultConstructible {
 };
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(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,10 +20,6 @@ struct Globals {
     type_4 d;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 float index_array(
     int i,
     device Globals const& globals,

diff --git a/tests/out/msl/bounds-check-zero-atomic.msl b/tests/out/msl/bounds-check-zero-atomic.msl
@@ -24,10 +24,6 @@ struct Globals {
     type_2 c;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 uint fetch_add_atomic(
     device Globals& globals,
     constant _mslBufferSizes& _buffer_sizes

diff --git a/tests/out/msl/bounds-check-zero.msl b/tests/out/msl/bounds-check-zero.msl
@@ -26,10 +26,6 @@ struct Globals {
     type_4 d;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 float index_array(
     int i,
     device Globals const& globals,

diff --git a/tests/out/msl/break-if.msl b/tests/out/msl/break-if.msl
@@ -5,10 +5,6 @@
 using metal::uint;
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 void breakIfEmpty(
 ) {
     bool loop_init = true;

diff --git a/tests/out/msl/collatz.msl b/tests/out/msl/collatz.msl
@@ -13,10 +13,6 @@ struct PrimeIndices {
     type_1 data;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 uint collatz_iterations(
     uint n_base
 ) {

diff --git a/tests/out/msl/constructors.msl b/tests/out/msl/constructors.msl
@@ -30,10 +30,6 @@ constant type_10 cz6_ = type_10 {};
 constant Foo cz7_ = Foo {};
 constant type_11 cp3_ = type_11 {0, 1, 2, 3};
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 kernel void main_(
 ) {
     Foo foo = {};

diff --git a/tests/out/msl/control-flow.msl b/tests/out/msl/control-flow.msl
@@ -5,10 +5,6 @@
 using metal::uint;
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 void switch_default_break(
     int i
 ) {

diff --git a/tests/out/msl/do-while.msl b/tests/out/msl/do-while.msl
@@ -5,10 +5,6 @@
 using metal::uint;
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 void fb1_(
     thread bool& cond
 ) {

diff --git a/tests/out/msl/dualsource.msl b/tests/out/msl/dualsource.msl
@@ -9,10 +9,6 @@ struct FragmentOutput {
     metal::float4 mask;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 struct main_Input {
 };
 struct main_Output {

diff --git a/tests/out/msl/empty-global-name.msl b/tests/out/msl/empty-global-name.msl
@@ -8,10 +8,6 @@ struct type_1 {
     int member;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 void function(
     device type_1& unnamed
 ) {

diff --git a/tests/out/msl/empty.msl b/tests/out/msl/empty.msl
@@ -5,10 +5,6 @@
 using metal::uint;
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 kernel void main_(
 ) {
     return;

diff --git a/tests/out/msl/extra.msl b/tests/out/msl/extra.msl
@@ -14,10 +14,6 @@ struct FragmentIn {
     uint primitive_index;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 struct main_Input {
     metal::float4 color [[user(loc0), center_perspective]];
 };

diff --git a/tests/out/msl/fragment-output.msl b/tests/out/msl/fragment-output.msl
@@ -21,10 +21,6 @@ struct FragmentOutputVec2Scalar {
     uint scalaru;
 };
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 struct main_vec4vec3_Output {
     metal::float4 vec4f [[color(0)]];
     metal::int4 vec4i [[color(1)]];

diff --git a/tests/out/msl/functions.msl b/tests/out/msl/functions.msl
@@ -5,10 +5,6 @@
 using metal::uint;
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 metal::float2 test_fma(
 ) {
     metal::float2 a = metal::float2(2.0, 2.0);

diff --git a/tests/out/msl/globals.msl b/tests/out/msl/globals.msl
@@ -33,10 +33,6 @@ struct type_15 {
 };
 constant bool Foo_1 = true;
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 void test_msl_packed_vec3_as_arg(
     metal::float3 arg
 ) {

diff --git a/tests/out/msl/image.msl b/tests/out/msl/image.msl
@@ -5,10 +5,6 @@
 using metal::uint;
 
 
-template <typename T> inline T naga_isign(T arg) {
-    return metal::select(metal::select(T(-1), T(1), (arg > 0)), 0, (arg == 0));
-}
-
 struct main_Input {
 };
 kernel void main_(