Avoid dynamic allocation in WASM assembler for storing function types

PiperOrigin-RevId: 545204376

Author: akopich

src/xnnpack/array-helpers.h

@@ -14,19 +14,19 @@ void ArrayApplyImpl(std::array<T, N>&& args, F&& f,
 }
 
 template <typename T, size_t N, typename F,
-          typename Indx = std::make_index_sequence<N> >
+          typename Indx = std::make_index_sequence<N>>
 void ArrayApply(std::array<T, N>&& args, F&& f) {
   return ArrayApplyImpl(std::move(args), f, Indx{});
 }
 
 template <size_t... Is, typename V>
-std::array<V, sizeof...(Is)> MakeArrayImpl(
+constexpr std::array<V, sizeof...(Is)> MakeArrayImpl(
     V value, std::integer_sequence<size_t, Is...>) {
   return {((void)Is, value)...};
 }
 
 template <size_t N, typename V>
-std::array<V, N> MakeArray(V value) {
+constexpr std::array<V, N> MakeArray(V value) {
   return MakeArrayImpl(value, std::make_index_sequence<N>{});
 }
 
@@ -36,21 +36,24 @@ static constexpr T kDefault{};
 template <typename T, size_t max_size>
 class ArrayPrefix {
  public:
-  explicit ArrayPrefix(size_t size) : size_(size) { assert(size_ <= max_size); }
+  constexpr ArrayPrefix(size_t size, T t)
+      : size_(size), array_(MakeArray<max_size>(t)) {
+    assert(size_ <= max_size);
+  }
 
-  ArrayPrefix(size_t size, T t) : size_(size), array_(MakeArray<max_size>(t)) {
+  explicit constexpr ArrayPrefix(size_t size) : size_(size) {
     assert(size_ <= max_size);
   }
 
-  template <typename Array>
-  explicit ArrayPrefix(Array&& array) : ArrayPrefix({}) {
+  template <typename Array,
+            typename = std::enable_if_t<!std::is_integral_v<Array>>>
+  explicit constexpr ArrayPrefix(Array&& array) : ArrayPrefix({}) {
     for (const auto& v : array) {
-      array_[size_] = v;
-      size_++;
+      push_back(v);
     }
   }
 
-  ArrayPrefix(std::initializer_list<T> init)
+  constexpr ArrayPrefix(std::initializer_list<T> init)
       : ArrayPrefix(init.size(), kDefault<T>) {
     assert(size_ <= max_size);
     std::copy(init.begin(), init.end(), begin());
@@ -76,6 +79,10 @@ class ArrayPrefix {
     assert(index < size_);
     return array_[index];
   }
+  void push_back(const T& t) {
+    assert(size_ + 1 < max_size);
+    array_[size_++] = t;
+  }
   size_t size() const { return size_; }
 
  private:

src/xnnpack/wasm-assembler.h

@@ -79,7 +79,8 @@ static uint32_t VectorEncodingLength(
 }
 
 struct ResultType {
-  ResultType(std::initializer_list<ValType> codes) : type(kNoTypeCode) {
+  constexpr ResultType(std::initializer_list<ValType> codes)
+      : type(kNoTypeCode) {
     switch (codes.size()) {
       case 0:
         break;
@@ -99,20 +100,32 @@ struct ResultType {
   static constexpr byte kNoTypeCode = 0;
 };
 
+template <>
+static constexpr ResultType kDefault<ResultType>{};
+
 inline bool operator==(const ResultType& lhs, const ResultType& rhs) {
   return lhs.type == rhs.type;
 }
 
 static constexpr size_t kMaxParamsCount = 16;
-using Params = ArrayPrefix<ValType, kMaxParamsCount>;
+struct Params : ArrayPrefix<ValType, kMaxParamsCount> {
+  using ArrayPrefix::ArrayPrefix;
+};
+
+template <>
+static constexpr Params kDefault<Params>{0, kDefault<ValType>};
 
 struct FuncType {
-  FuncType(const Params& params, ResultType result)
+  constexpr FuncType(const Params& params, ResultType result)
       : params(params), result(result) {}
   Params params;
   ResultType result;
 };
 
+template <>
+static constexpr FuncType kDefault<FuncType>{kDefault<Params>,
+                                             kDefault<ResultType>};
+
 inline bool operator==(const FuncType& lhs, const FuncType& rhs) {
   return lhs.result == rhs.result &&
          std::equal(lhs.params.begin(), lhs.params.end(), rhs.params.begin(),
@@ -786,9 +799,12 @@ class WasmAssembler : public AssemblerBase, public internal::WasmOps {
     internal::StoreEncodedU32(n, [this](byte b) { emit8(b); });
   }
 
+  static constexpr size_t kMaxNumFuncTypes = 16;
+
   std::vector<Function> functions_;
-  std::vector<FuncType> func_types_;
   std::vector<Export> exports_;
+  internal::ArrayPrefix<FuncType, kMaxNumFuncTypes> func_types_{
+      0, internal::kDefault<FuncType>};
 };
 
 }  // namespace xnnpack