Merge pull request #3232 from KimYannn:feature/lookup64

PiperOrigin-RevId: 953262574

Author: copybara-github

g3doc/quick_reference.md

@@ -2645,15 +2645,17 @@ The following `ReverseN` must not be called if `Lanes(D()) < N`:
     `TwoTablesLookupLanes(a, b, indices)` on RVV/SVE if `Lanes(d) <
     Lanes(DFromV<V>())`.
 
-Each of the `Lookup8`, `Lookup16`, `Lookup32` (let $X denote the 8/16/32) ops
-below return `GatherIndex(D(), tbl, indices)`, but are much more efficient, and
-are limited to $X elements. Results are undefined if any indices are >= $X. They
-are implemented using `TableLookupLanes` or `TwoTablesLookupLanes`. Let `T`
-denote `TFromD<D>`. These ops are guaranteed to work if `D` is a full vector,
-`HWY_TARGET != HWY_SCALAR` and `HWY_MIN_BYTES / sizeof(T) >= $X/2`. Use the
-constexpr function `CanLookup$X(D())` to verify this. Even if it returns false,
-the ops are still safe to call if `Lanes(D())` >= $X/2. Note that `tbl` must be
-$X-element aligned!
+Each of the `Lookup8`, `Lookup16`, `Lookup32`, `Lookup64` (let $X denote the
+8/16/32/64) ops below return `GatherIndex(D(), tbl, indices)`, but are much more
+efficient, and are limited to $X elements. Results are undefined if any indices
+are >= $X. They are implemented using `TableLookupLanes` or
+`TwoTablesLookupLanes`. Let `T` denote `TFromD<D>`. These ops are guaranteed to
+work if `D` is a full vector, `HWY_TARGET != HWY_SCALAR` and
+`HWY_MIN_BYTES / sizeof(T) >= $X/2`. `Lookup64` is also guaranteed for 128-bit
+AArch64 NEON vectors if `T` is byte-sized. Use the constexpr function
+`CanLookup$X(D())` to verify this. Even if it returns false, the ops are still
+safe to call if `Lanes(D()) >= $X/2`. Note that `tbl` must be $X-element
+aligned!
 
 *   `D`: {u,i,f}{16,32,64} \
     <code>Vec&lt;D&gt; **Lookup8**(D, const TFromD&lt;D&gt;* tbl, VI
@@ -2667,6 +2669,10 @@ $X-element aligned!
     <code>Vec&lt;D&gt; **Lookup32**(D, const TFromD&lt;D&gt;* tbl, VI
     indices)</code>: as above, with $X = 32.
 
+*   `D`: {u,i}{8} \
+    <code>Vec&lt;D&gt; **Lookup64**(D, const TFromD&lt;D&gt;* tbl, VI
+    indices)</code>: as above, with $X = 64.
+
 *   <code>unspecified **IndicesFromVec**(D d, V idx)</code> prepares for
     `TableLookupLanes` or `TwoTablesLookupLanes` with integer indices in `idx`,
     which must be the same bit width as `TFromD<D>` and in the range `[0, 2 *

hwy/ops/arm_neon-inl.h

@@ -6428,6 +6428,40 @@ HWY_API Vec128<T> TwoTablesLookupLanes(Vec128<T> a, Vec128<T> b,
 #endif
 }
 
+// ------------------------------ Lookup64
+
+#if HWY_ARCH_ARM_A64
+
+// Per-target flag to prevent generic_ops-inl.h from defining Lookup64.
+#ifdef HWY_NATIVE_LOOKUP64
+#undef HWY_NATIVE_LOOKUP64
+#else
+#define HWY_NATIVE_LOOKUP64
+#endif
+
+template <class D, class VI, HWY_IF_UI8_D(D)>
+HWY_INLINE VFromD<D> Lookup64(D d, const TFromD<D>* HWY_RESTRICT table,
+                             VI indices) {
+  const DFromV<VI> di;
+  static_assert(sizeof(TFromD<D>) == sizeof(TFromD<decltype(di)>),
+                "Index/vector must have same lane size");
+  HWY_IF_CONSTEXPR(HWY_IS_DEBUG_BUILD) {
+    HWY_DASSERT(AllTrue(di, Lt(indices, Set(di, 64))));
+  }
+
+  const Full128<uint8_t> du8;
+  const auto t0 = Load(du8, reinterpret_cast<const uint8_t*>(table));
+  const auto t1 = Load(du8, reinterpret_cast<const uint8_t*>(table) + 16);
+  const auto t2 = Load(du8, reinterpret_cast<const uint8_t*>(table) + 32);
+  const auto t3 = Load(du8, reinterpret_cast<const uint8_t*>(table) + 48);
+  detail::Tuple4<uint8_t, 16> tables = {{{t0.raw, t1.raw, t2.raw, t3.raw}}};
+  const auto idx_u8 = ResizeBitCast(du8, indices);
+  return ResizeBitCast(
+      d, Vec128<uint8_t>{vqtbl4q_u8(tables.raw, idx_u8.raw)});
+}
+
+#endif  // HWY_ARCH_ARM_A64
+
 // ------------------------------ Reverse2 (CombineShiftRightBytes)
 
 // Per-target flag to prevent generic_ops-inl.h defining 8-bit Reverse2/4/8.

hwy/ops/generic_ops-inl.h

@@ -7286,6 +7286,52 @@ HWY_INLINE Vec<D> Lookup32(D d, const T* HWY_RESTRICT table, VI indices) {
   }
 }
 
+// ------------------------------ Lookup64
+
+#if (defined(HWY_NATIVE_LOOKUP64) == defined(HWY_TARGET_TOGGLE)) || HWY_IDE
+#ifdef HWY_NATIVE_LOOKUP64
+#undef HWY_NATIVE_LOOKUP64
+#else
+#define HWY_NATIVE_LOOKUP64
+#endif
+
+template <class D, typename T = TFromD<D>, class VI>
+HWY_INLINE Vec<D> Lookup64(D d, const T* HWY_RESTRICT table, VI indices) {
+  const DFromV<VI> di;
+  static_assert(sizeof(T) == 1, "Lookup64 requires 8-bit table lanes");
+  static_assert(sizeof(T) == sizeof(TFromD<decltype(di)>),
+                "Index/vector must have same lane size");
+  HWY_IF_CONSTEXPR(HWY_IS_DEBUG_BUILD) {
+    HWY_DASSERT(Lanes(d) >= 32);
+    HWY_DASSERT(AllTrue(di, Lt(indices, Set(di, 64))));
+  }
+
+  HWY_IF_CONSTEXPR(!HWY_HAVE_SCALABLE) {
+    HWY_IF_CONSTEXPR(MaxLanes(d) >= 64) {
+      const CappedTag<T, 64> d64;
+      const Vec<D> t0 = ZeroExtendResizeBitCast(d, d64, Load(d64, table));
+      return TableLookupLanes(t0, IndicesFromVec(d, indices));
+    }
+    HWY_IF_CONSTEXPR(MaxLanes(d) < 64) {
+      const Vec<D> t0 = Load(d, table);
+      const Vec<D> t1 = Load(d, table + 32);
+      return TwoTablesLookupLanes(d, t0, t1, IndicesFromVec(d, indices));
+    }
+  }
+
+  HWY_IF_CONSTEXPR(HWY_HAVE_SCALABLE) {
+    const Vec<D> t0 = LoadN(d, table, 32);
+    const Vec<D> t1 = LoadN(d, table + 32, 32);
+    const Mask<decltype(di)> ge_32 = Ge(indices, Set(di, 32));
+    const VI idx_for_t1 = Sub(indices, Set(di, 32));
+    const Vec<D> r0 = TableLookupLanes(t0, IndicesFromVec(d, indices));
+    const Vec<D> r1 = TableLookupLanes(t1, IndicesFromVec(d, idx_for_t1));
+    return IfThenElse(RebindMask(d, ge_32), r1, r0);
+  }
+}
+
+#endif  // HWY_NATIVE_LOOKUP64
+
 // ------------------------------ Reverse2, Reverse4, Reverse8 (8-bit)
 
 #if (defined(HWY_NATIVE_REVERSE2_8) == defined(HWY_TARGET_TOGGLE)) || HWY_IDE

hwy/ops/shared-inl.h

@@ -596,6 +596,14 @@ HWY_API constexpr bool CanLookup32(D d) {
          (HWY_HAVE_SCALABLE && detail::IsFull(d) && (sizeof(T) == 1));
 }
 
+// Returns whether `Lookup64` can definitely be used for vectors created from
+// tag `d`. May return a false negative for large scalable vectors.
+template <class D, typename T = TFromD<D>>
+HWY_API constexpr bool CanLookup64(D d) {
+  return (!HWY_HAVE_SCALABLE && MaxLanes(d) >= 32) ||
+         (HWY_ARCH_ARM_A64 && HWY_TARGET_IS_NEON && (sizeof(T) == 1));
+}
+
 // ------------------------------ Choosing overloads (SFINAE)
 
 // Same as base.h macros but with a Simd<T, N, kPow2> argument instead of T.

hwy/tests/table_test.cc

@@ -366,6 +366,42 @@ HWY_NOINLINE void TestAllLookup32() {
 #endif
 }
 
+struct TestLookup64 {
+  template <class T, class D>
+  HWY_NOINLINE void operator()(T /*unused*/, D d) {
+    using V = Vec<D>;
+    const RebindToUnsigned<D> du;
+    using TU = TFromD<decltype(du)>;
+
+    const size_t N = Lanes(d);
+    if (!CanLookup64(d)) return;
+
+    const size_t padded_N = HWY_MAX(N, 64);
+    auto tbl = AllocateAligned<T>(padded_N);
+    auto idx = AllocateAligned<TU>(padded_N);
+    auto expected = AllocateAligned<T>(padded_N);
+    HWY_ASSERT(tbl && idx && expected);
+
+    for (size_t i = 0; i < padded_N; ++i) {
+      tbl[i] = ConvertScalarTo<T>(i + static_cast<size_t>(Unpredictable1()));
+    }
+
+    HWY_ALIGN TU idx_source[16] = {0,  63, 1,  62, 15, 16, 31, 32,
+                                   47, 48, 7,  56, 23, 40, 55, 8};
+    for (size_t j = 0; j < padded_N; ++j) {
+      idx[j] = static_cast<TU>(idx_source[j & 15]);
+      expected[j] = ConvertScalarTo<T>(idx[j] + 1);
+    }
+
+    const V actual = Lookup64(d, tbl.get(), Load(du, idx.get()));
+    HWY_ASSERT_VEC_EQ(d, expected.get(), actual);
+  }
+};
+
+HWY_NOINLINE void TestAllLookup64() {
+  ForUI8(ForPartialVectors<TestLookup64>());
+}
+
 }  // namespace
 // NOLINTNEXTLINE(google-readability-namespace-comments)
 }  // namespace HWY_NAMESPACE
@@ -381,6 +417,7 @@ HWY_EXPORT_AND_TEST_P(HwyTableTest, TestAllTwoTablesLookupLanes);
 HWY_EXPORT_AND_TEST_P(HwyTableTest, TestAllLookup8);
 HWY_EXPORT_AND_TEST_P(HwyTableTest, TestAllLookup16);
 HWY_EXPORT_AND_TEST_P(HwyTableTest, TestAllLookup32);
+HWY_EXPORT_AND_TEST_P(HwyTableTest, TestAllLookup64);
 HWY_AFTER_TEST();
 }  // namespace
 }  // namespace hwy