Add basic WASM SIMD128 support to elementwise kernel generator.

PiperOrigin-RevId: 891805395

Author: vksnk

ynnpack/kernels/binary/BUILD

@@ -27,6 +27,7 @@ ARCHS = [
     "x86_avx512",
     "x86_avx512bf16",
     "arm_neon",
+    "wasm_simd128",
 ]
 
 [ynn_generate_src_hdr(

ynnpack/kernels/binary/generator.py

@@ -65,6 +65,17 @@ def main(argv: Sequence[str]) -> None:
           (max_fp32, (8, 1)),
           (min_fp32, (8, 1)),
       ],
+      "wasm_simd128": [
+          (add_fp32, (8, 1)),
+          (subtract_fp32, (8, 1)),
+          (multiply_fp32, (8, 1)),
+          # (multiply_int32_fp32, (8, 1)),
+          (divide_fp32, (8, 1)),
+          (copysign_fp32, (8, 1)),
+          (max_fp32, (8, 1)),
+          (min_fp32, (8, 1)),
+      ],
+
   }[target]
 
   generate_elementwise_kernels(output_src, output_inc, target, kernels)

ynnpack/kernels/binary/kernels.inc

@@ -20,3 +20,8 @@
 #include "ynnpack/kernels/binary/arm_neon.inc"
 #endif  // YNN_ARCH_ARM_NEON
 
+#ifdef YNN_ARCH_WASM_SIMD128
+#include "ynnpack/kernels/binary/wasm_simd128.inc"
+#endif  // YNN_ARCH_ARM_NEON
+
+

ynnpack/kernels/elementwise/BUILD

@@ -13,6 +13,7 @@ py_library(
         "compiler.py",
         "generator.py",
         "rules.py",
+        "wasm.py",
         "x86.py",
     ],
     tags = ["manual"],

ynnpack/kernels/elementwise/generator.py

@@ -5,6 +5,7 @@
 
 # pylint: disable=undefined-variable
 from ynnpack.kernels.elementwise.arm import *  # pylint: disable=wildcard-import
+from ynnpack.kernels.elementwise.wasm import *  # pylint: disable=wildcard-import
 from ynnpack.kernels.elementwise.x86 import *  # pylint: disable=wildcard-import
 
 arch_to_target = {
@@ -22,6 +23,7 @@
     "arm_neon": ARM(["NEON"]),
     "arm_neonfp16": ARM(["NEONFP16"]),
     "arm_neon_fma": ARM(["NEON", "FMA"]),
+    "wasm_simd128": WASM(["WASM_SIMD128"]),
 }
 
 

ynnpack/kernels/elementwise/wasm.py

@@ -0,0 +1,59 @@
+"""WASM SIMD target for elementwise kernels compiler."""
+
+# pylint: disable=undefined-variable
+from ynnpack.kernels.elementwise.compiler import *  # pylint: disable=wildcard-import
+from ynnpack.kernels.elementwise.rules import *  # pylint: disable=wildcard-import
+
+
+class WASM(Target):
+  """WASM SIMD target for elementwise kernels compiler."""
+
+  def update_for_simd128(self):
+    """Updates the target for WASM SIMD128 support."""
+    self.header += """
+namespace ynn {
+namespace {
+template <>
+YNN_INTRINSIC simd::vec<float, 4> select_greater_than(simd::vec<float, 4> a, simd::vec<float, 4> b, simd::vec<float, 4> c, simd::vec<float, 4> d) {
+  v128_t mask = wasm_f32x4_gt(a.v, b.v);
+  return simd::vec<float, 4>{wasm_v128_bitselect(c.v, d.v, mask)};
+}
+}
+} // namespace ynn
+"""
+
+  def __init__(self, features):
+    Target.__init__(self)
+    self.patterns += add_select_rules()
+    self.patterns += add_saturating_cast_rules()
+    self.patterns += add_shift_rules()
+
+    self.features = features
+    self.vector_bits = 128
+    self.tail_strategy = TailStrategy.VECTOR
+
+    # These are transitive.
+    implied_features = {
+        "WASM_SIMD128": [],
+    }
+    all_features = []
+    self.compute_all_features(features, implied_features, all_features)
+
+    known_features = ["WASM_SIMD128"]
+    for feature in all_features:
+      if feature not in known_features:
+        raise ValueError(f"Unknown feature: {feature}")
+
+    self.header += (
+        '#include "ynnpack/base/simd/wasm_simd128.h"\n'
+    )
+
+    if "WASM_SIMD128" in all_features:
+      self.update_for_simd128()
+
+  def arch_flags(self):
+    return "|".join(["arch_flag::" + i.lower() for i in self.features])
+
+  def arch_string(self):
+    features_str = "_".join([i.lower() for i in self.features])
+    return features_str