third_party/highway/hwy/contrib/algo/transform-inl.h - aom - Git at Google

 // Copyright 2022 Google LLC
 // SPDX-License-Identifier: Apache-2.0
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 // Per-target include guard
 #if defined(HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_) == \
     defined(HWY_TARGET_TOGGLE)
 #ifdef HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_
 #undef HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_
 #else
 #define HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_
 #endif

 #include <stddef.h>

 #include "third_party/highway/hwy/highway.h"

 HWY_BEFORE_NAMESPACE();
 namespace hwy {
 namespace HWY_NAMESPACE {

 // These functions avoid having to write a loop plus remainder handling in the
 // (unfortunately still common) case where arrays are not aligned/padded. If the
 // inputs are known to be aligned/padded, it is more efficient to write a single
 // loop using Load(). We do not provide a TransformAlignedPadded because it
 // would be more verbose than such a loop.
 //
 // Func is either a functor with a templated operator()(d, v[, v1[, v2]]), or a
 // generic lambda if using C++14. The d argument is the same as was passed to
 // the Generate etc. functions. Due to apparent limitations of Clang, it is
 // currently necessary to add HWY_ATTR before the opening { of the lambda to
 // avoid errors about "always_inline function .. requires target".
 //
 // We do not check HWY_MEM_OPS_MIGHT_FAULT because LoadN/StoreN do not fault.

 // Fills `out[0, count)` with the vectors returned by `func(d, index_vec)`,
 // where `index_vec` is `Vec<RebindToUnsigned<D>>`. On the first call to `func`,
 // the value of its lane i is i, and increases by `Lanes(d)` after every call.
 // Note that some of these indices may be `>= count`, but the elements that
 // `func` returns in those lanes will not be written to `out`.
 template <class D, class Func, typename T = TFromD<D>>
 void Generate(D d, T* HWY_RESTRICT out, size_t count, const Func& func) {
   const RebindToUnsigned<D> du;
   using TU = TFromD<decltype(du)>;
   const size_t N = Lanes(d);

   size_t idx = 0;
   Vec<decltype(du)> vidx = Iota(du, 0);
   if (count >= N) {
     for (; idx <= count - N; idx += N) {
       StoreU(func(d, vidx), d, out + idx);
       vidx = Add(vidx, Set(du, static_cast<TU>(N)));
     }
   }

   // `count` was a multiple of the vector length `N`: already done.
   if (HWY_UNLIKELY(idx == count)) return;

   const size_t remaining = count - idx;
   HWY_DASSERT(0 != remaining && remaining < N);
   StoreN(func(d, vidx), d, out + idx, remaining);
 }

 // Calls `func(d, v)` for each input vector; out of bound lanes with index i >=
 // `count` are instead taken from `no[i % Lanes(d)]`.
 template <class D, class Func, typename T = TFromD<D>>
 void Foreach(D d, const T* HWY_RESTRICT in, const size_t count, const Vec<D> no,
              const Func& func) {
   const size_t N = Lanes(d);

   size_t idx = 0;
   if (count >= N) {
     for (; idx <= count - N; idx += N) {
       const Vec<D> v = LoadU(d, in + idx);
       func(d, v);
     }
   }

   // `count` was a multiple of the vector length `N`: already done.
   if (HWY_UNLIKELY(idx == count)) return;

   const size_t remaining = count - idx;
   HWY_DASSERT(0 != remaining && remaining < N);
   const Vec<D> v = LoadNOr(no, d, in + idx, remaining);
   func(d, v);
 }

 // Replaces `inout[idx]` with `func(d, inout[idx])`. Example usage: multiplying
 // array elements by a constant.
 template <class D, class Func, typename T = TFromD<D>>
 void Transform(D d, T* HWY_RESTRICT inout, size_t count, const Func& func) {
   const size_t N = Lanes(d);

   size_t idx = 0;
   if (count >= N) {
     for (; idx <= count - N; idx += N) {
       const Vec<D> v = LoadU(d, inout + idx);
       StoreU(func(d, v), d, inout + idx);
     }
   }

   // `count` was a multiple of the vector length `N`: already done.
   if (HWY_UNLIKELY(idx == count)) return;

   const size_t remaining = count - idx;
   HWY_DASSERT(0 != remaining && remaining < N);
   const Vec<D> v = LoadN(d, inout + idx, remaining);
   StoreN(func(d, v), d, inout + idx, remaining);
 }

 // Replaces `inout[idx]` with `func(d, inout[idx], in1[idx])`. Example usage:
 // multiplying array elements by those of another array.
 template <class D, class Func, typename T = TFromD<D>>
 void Transform1(D d, T* HWY_RESTRICT inout, size_t count,
                 const T* HWY_RESTRICT in1, const Func& func) {
   const size_t N = Lanes(d);

   size_t idx = 0;
   if (count >= N) {
     for (; idx <= count - N; idx += N) {
       const Vec<D> v = LoadU(d, inout + idx);
       const Vec<D> v1 = LoadU(d, in1 + idx);
       StoreU(func(d, v, v1), d, inout + idx);
     }
   }

   // `count` was a multiple of the vector length `N`: already done.
   if (HWY_UNLIKELY(idx == count)) return;

   const size_t remaining = count - idx;
   HWY_DASSERT(0 != remaining && remaining < N);
   const Vec<D> v = LoadN(d, inout + idx, remaining);
   const Vec<D> v1 = LoadN(d, in1 + idx, remaining);
   StoreN(func(d, v, v1), d, inout + idx, remaining);
 }

 // Replaces `inout[idx]` with `func(d, inout[idx], in1[idx], in2[idx])`. Example
 // usage: FMA of elements from three arrays, stored into the first array.
 template <class D, class Func, typename T = TFromD<D>>
 void Transform2(D d, T* HWY_RESTRICT inout, size_t count,
                 const T* HWY_RESTRICT in1, const T* HWY_RESTRICT in2,
                 const Func& func) {
   const size_t N = Lanes(d);

   size_t idx = 0;
   if (count >= N) {
     for (; idx <= count - N; idx += N) {
       const Vec<D> v = LoadU(d, inout + idx);
       const Vec<D> v1 = LoadU(d, in1 + idx);
       const Vec<D> v2 = LoadU(d, in2 + idx);
       StoreU(func(d, v, v1, v2), d, inout + idx);
     }
   }

   // `count` was a multiple of the vector length `N`: already done.
   if (HWY_UNLIKELY(idx == count)) return;

   const size_t remaining = count - idx;
   HWY_DASSERT(0 != remaining && remaining < N);
   const Vec<D> v = LoadN(d, inout + idx, remaining);
   const Vec<D> v1 = LoadN(d, in1 + idx, remaining);
   const Vec<D> v2 = LoadN(d, in2 + idx, remaining);
   StoreN(func(d, v, v1, v2), d, inout + idx, remaining);
 }

 template <class D, typename T = TFromD<D>>
 void Replace(D d, T* HWY_RESTRICT inout, size_t count, T new_t, T old_t) {
   const size_t N = Lanes(d);
   const Vec<D> old_v = Set(d, old_t);
   const Vec<D> new_v = Set(d, new_t);

   size_t idx = 0;
   if (count >= N) {
     for (; idx <= count - N; idx += N) {
       Vec<D> v = LoadU(d, inout + idx);
       StoreU(IfThenElse(Eq(v, old_v), new_v, v), d, inout + idx);
     }
   }

   // `count` was a multiple of the vector length `N`: already done.
   if (HWY_UNLIKELY(idx == count)) return;

   const size_t remaining = count - idx;
   HWY_DASSERT(0 != remaining && remaining < N);
   const Vec<D> v = LoadN(d, inout + idx, remaining);
   StoreN(IfThenElse(Eq(v, old_v), new_v, v), d, inout + idx, remaining);
 }

 template <class D, class Func, typename T = TFromD<D>>
 void ReplaceIf(D d, T* HWY_RESTRICT inout, size_t count, T new_t,
                const Func& func) {
   const size_t N = Lanes(d);
   const Vec<D> new_v = Set(d, new_t);

   size_t idx = 0;
   if (count >= N) {
     for (; idx <= count - N; idx += N) {
       Vec<D> v = LoadU(d, inout + idx);
       StoreU(IfThenElse(func(d, v), new_v, v), d, inout + idx);
     }
   }

   // `count` was a multiple of the vector length `N`: already done.
   if (HWY_UNLIKELY(idx == count)) return;

   const size_t remaining = count - idx;
   HWY_DASSERT(0 != remaining && remaining < N);
   const Vec<D> v = LoadN(d, inout + idx, remaining);
   StoreN(IfThenElse(func(d, v), new_v, v), d, inout + idx, remaining);
 }

 // NOLINTNEXTLINE(google-readability-namespace-comments)
 }  // namespace HWY_NAMESPACE
 }  // namespace hwy
 HWY_AFTER_NAMESPACE();

 #endif  // HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_
	// Copyright 2022 Google LLC
	// SPDX-License-Identifier: Apache-2.0
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	// Per-target include guard
	#if defined(HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_) == \
	defined(HWY_TARGET_TOGGLE)
	#ifdef HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_
	#undef HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_
	#else
	#define HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_
	#endif

	#include <stddef.h>

	#include "third_party/highway/hwy/highway.h"

	HWY_BEFORE_NAMESPACE();
	namespace hwy {
	namespace HWY_NAMESPACE {

	// These functions avoid having to write a loop plus remainder handling in the
	// (unfortunately still common) case where arrays are not aligned/padded. If the
	// inputs are known to be aligned/padded, it is more efficient to write a single
	// loop using Load(). We do not provide a TransformAlignedPadded because it
	// would be more verbose than such a loop.
	//
	// Func is either a functor with a templated operator()(d, v[, v1[, v2]]), or a
	// generic lambda if using C++14. The d argument is the same as was passed to
	// the Generate etc. functions. Due to apparent limitations of Clang, it is
	// currently necessary to add HWY_ATTR before the opening { of the lambda to
	// avoid errors about "always_inline function .. requires target".
	//
	// We do not check HWY_MEM_OPS_MIGHT_FAULT because LoadN/StoreN do not fault.

	// Fills `out[0, count)` with the vectors returned by `func(d, index_vec)`,
	// where `index_vec` is `Vec<RebindToUnsigned<D>>`. On the first call to `func`,
	// the value of its lane i is i, and increases by `Lanes(d)` after every call.
	// Note that some of these indices may be `>= count`, but the elements that
	// `func` returns in those lanes will not be written to `out`.
	template <class D, class Func, typename T = TFromD<D>>
	void Generate(D d, T* HWY_RESTRICT out, size_t count, const Func& func) {
	const RebindToUnsigned<D> du;
	using TU = TFromD<decltype(du)>;
	const size_t N = Lanes(d);

	size_t idx = 0;
	Vec<decltype(du)> vidx = Iota(du, 0);
	if (count >= N) {
	for (; idx <= count - N; idx += N) {
	StoreU(func(d, vidx), d, out + idx);
	vidx = Add(vidx, Set(du, static_cast<TU>(N)));
	}
	}

	// `count` was a multiple of the vector length `N`: already done.
	if (HWY_UNLIKELY(idx == count)) return;

	const size_t remaining = count - idx;
	HWY_DASSERT(0 != remaining && remaining < N);
	StoreN(func(d, vidx), d, out + idx, remaining);
	}

	// Calls `func(d, v)` for each input vector; out of bound lanes with index i >=
	// `count` are instead taken from `no[i % Lanes(d)]`.
	template <class D, class Func, typename T = TFromD<D>>
	void Foreach(D d, const T* HWY_RESTRICT in, const size_t count, const Vec<D> no,
	const Func& func) {
	const size_t N = Lanes(d);

	size_t idx = 0;
	if (count >= N) {
	for (; idx <= count - N; idx += N) {
	const Vec<D> v = LoadU(d, in + idx);
	func(d, v);
	}
	}

	// `count` was a multiple of the vector length `N`: already done.
	if (HWY_UNLIKELY(idx == count)) return;

	const size_t remaining = count - idx;
	HWY_DASSERT(0 != remaining && remaining < N);
	const Vec<D> v = LoadNOr(no, d, in + idx, remaining);
	func(d, v);
	}

	// Replaces `inout[idx]` with `func(d, inout[idx])`. Example usage: multiplying
	// array elements by a constant.
	template <class D, class Func, typename T = TFromD<D>>
	void Transform(D d, T* HWY_RESTRICT inout, size_t count, const Func& func) {
	const size_t N = Lanes(d);

	size_t idx = 0;
	if (count >= N) {
	for (; idx <= count - N; idx += N) {
	const Vec<D> v = LoadU(d, inout + idx);
	StoreU(func(d, v), d, inout + idx);
	}
	}

	// `count` was a multiple of the vector length `N`: already done.
	if (HWY_UNLIKELY(idx == count)) return;

	const size_t remaining = count - idx;
	HWY_DASSERT(0 != remaining && remaining < N);
	const Vec<D> v = LoadN(d, inout + idx, remaining);
	StoreN(func(d, v), d, inout + idx, remaining);
	}

	// Replaces `inout[idx]` with `func(d, inout[idx], in1[idx])`. Example usage:
	// multiplying array elements by those of another array.
	template <class D, class Func, typename T = TFromD<D>>
	void Transform1(D d, T* HWY_RESTRICT inout, size_t count,
	const T* HWY_RESTRICT in1, const Func& func) {
	const size_t N = Lanes(d);

	size_t idx = 0;
	if (count >= N) {
	for (; idx <= count - N; idx += N) {
	const Vec<D> v = LoadU(d, inout + idx);
	const Vec<D> v1 = LoadU(d, in1 + idx);
	StoreU(func(d, v, v1), d, inout + idx);
	}
	}

	// `count` was a multiple of the vector length `N`: already done.
	if (HWY_UNLIKELY(idx == count)) return;

	const size_t remaining = count - idx;
	HWY_DASSERT(0 != remaining && remaining < N);
	const Vec<D> v = LoadN(d, inout + idx, remaining);
	const Vec<D> v1 = LoadN(d, in1 + idx, remaining);
	StoreN(func(d, v, v1), d, inout + idx, remaining);
	}

	// Replaces `inout[idx]` with `func(d, inout[idx], in1[idx], in2[idx])`. Example
	// usage: FMA of elements from three arrays, stored into the first array.
	template <class D, class Func, typename T = TFromD<D>>
	void Transform2(D d, T* HWY_RESTRICT inout, size_t count,
	const T* HWY_RESTRICT in1, const T* HWY_RESTRICT in2,
	const Func& func) {
	const size_t N = Lanes(d);

	size_t idx = 0;
	if (count >= N) {
	for (; idx <= count - N; idx += N) {
	const Vec<D> v = LoadU(d, inout + idx);
	const Vec<D> v1 = LoadU(d, in1 + idx);
	const Vec<D> v2 = LoadU(d, in2 + idx);
	StoreU(func(d, v, v1, v2), d, inout + idx);
	}
	}

	// `count` was a multiple of the vector length `N`: already done.
	if (HWY_UNLIKELY(idx == count)) return;

	const size_t remaining = count - idx;
	HWY_DASSERT(0 != remaining && remaining < N);
	const Vec<D> v = LoadN(d, inout + idx, remaining);
	const Vec<D> v1 = LoadN(d, in1 + idx, remaining);
	const Vec<D> v2 = LoadN(d, in2 + idx, remaining);
	StoreN(func(d, v, v1, v2), d, inout + idx, remaining);
	}

	template <class D, typename T = TFromD<D>>
	void Replace(D d, T* HWY_RESTRICT inout, size_t count, T new_t, T old_t) {
	const size_t N = Lanes(d);
	const Vec<D> old_v = Set(d, old_t);
	const Vec<D> new_v = Set(d, new_t);

	size_t idx = 0;
	if (count >= N) {
	for (; idx <= count - N; idx += N) {
	Vec<D> v = LoadU(d, inout + idx);
	StoreU(IfThenElse(Eq(v, old_v), new_v, v), d, inout + idx);
	}
	}

	// `count` was a multiple of the vector length `N`: already done.
	if (HWY_UNLIKELY(idx == count)) return;

	const size_t remaining = count - idx;
	HWY_DASSERT(0 != remaining && remaining < N);
	const Vec<D> v = LoadN(d, inout + idx, remaining);
	StoreN(IfThenElse(Eq(v, old_v), new_v, v), d, inout + idx, remaining);
	}

	template <class D, class Func, typename T = TFromD<D>>
	void ReplaceIf(D d, T* HWY_RESTRICT inout, size_t count, T new_t,
	const Func& func) {
	const size_t N = Lanes(d);
	const Vec<D> new_v = Set(d, new_t);

	size_t idx = 0;
	if (count >= N) {
	for (; idx <= count - N; idx += N) {
	Vec<D> v = LoadU(d, inout + idx);
	StoreU(IfThenElse(func(d, v), new_v, v), d, inout + idx);
	}
	}

	// `count` was a multiple of the vector length `N`: already done.
	if (HWY_UNLIKELY(idx == count)) return;

	const size_t remaining = count - idx;
	HWY_DASSERT(0 != remaining && remaining < N);
	const Vec<D> v = LoadN(d, inout + idx, remaining);
	StoreN(IfThenElse(func(d, v), new_v, v), d, inout + idx, remaining);
	}

	// NOLINTNEXTLINE(google-readability-namespace-comments)
	} // namespace HWY_NAMESPACE
	} // namespace hwy
	HWY_AFTER_NAMESPACE();

	#endif // HIGHWAY_HWY_CONTRIB_ALGO_TRANSFORM_INL_H_