VectorizedArray: Implement a compare_and_apply_mask function

diff --git a/include/deal.II/base/vectorization.h b/include/deal.II/base/vectorization.h
index 9c754284653c2d568ba973979491598d8f968963..543efe5ed870d5c3830f0cbda103148bc4f55b27 100644 (file)
--- a/include/deal.II/base/vectorization.h
+++ b/include/deal.II/base/vectorization.h
@@ -3957,9 +3957,210 @@ operator<<(std::ostream &out, const VectorizedArray<Number, width> &p)
   return out;
 }
 
+//@}
+
+/**
+ * @name Ternary operations on VectorizedArray
+ */
+//@{
+
+
+/**
+ * TODO
+ */
+enum SIMDComparison : int
+{
+  equal                 = _CMP_EQ_OQ,
+  not_equal             = _CMP_NEQ_OQ,
+  less_than             = _CMP_LT_OQ,
+  less_than_or_equal    = _CMP_LE_OQ,
+  greater_than          = _CMP_GT_OQ,
+  greater_than_or_equal = _CMP_GE_OQ
+};
+
+
+/**
+ * TODO
+ */
+template <SIMDComparison predicate, typename Number>
+DEAL_II_ALWAYS_INLINE inline Number
+compare_and_apply_mask(const Number &left,
+                       const Number &right,
+                       const Number &true_value,
+                       const Number &false_value)
+{
+  bool mask;
+
+  switch (predicate)
+    {
+      case SIMDComparison::equal:
+        mask = left == right;
+        break;
+      case SIMDComparison::not_equal:
+        mask = left != right;
+        break;
+      case SIMDComparison::less_than:
+        mask = left < right;
+        break;
+      case SIMDComparison::less_than_or_equal:
+        mask = left <= right;
+        break;
+      case SIMDComparison::greater_than:
+        mask = left > right;
+        break;
+      case SIMDComparison::greater_than_or_equal:
+        mask = left >= right;
+        break;
+    }
+
+  return mask ? true_value : false_value;
+}
+
+
+/**
+ * TODO
+ */
+template <SIMDComparison predicate, typename Number>
+DEAL_II_ALWAYS_INLINE inline VectorizedArray<Number, 1>
+compare_and_apply_mask(const VectorizedArray<Number, 1> &left,
+                       const VectorizedArray<Number, 1> &right,
+                       const VectorizedArray<Number, 1> &true_value,
+                       const VectorizedArray<Number, 1> &false_value)
+{
+  bool mask;
+
+  switch (predicate)
+    {
+      case SIMDComparison::equal:
+        mask = left.data == right.data;
+        break;
+      case SIMDComparison::not_equal:
+        mask = left.data != right.data;
+        break;
+      case SIMDComparison::less_than:
+        mask = left.data < right.data;
+        break;
+      case SIMDComparison::less_than_or_equal:
+        mask = left.data <= right.data;
+        break;
+      case SIMDComparison::greater_than:
+        mask = left.data > right.data;
+        break;
+      case SIMDComparison::greater_than_or_equal:
+        mask = left.data >= right.data;
+        break;
+    }
+
+  return mask ? true_value.data : false_value.data;
+}
 
 //@}
 
+#ifndef DOXYGEN
+#  if DEAL_II_COMPILER_VECTORIZATION_LEVEL >= 3 && defined(__AVX512F__)
+
+template <int predicate>
+DEAL_II_ALWAYS_INLINE inline VectorizedArray<float, 16>
+compare_and_apply_mask(const VectorizedArray<float, 16> &left,
+                       const VectorizedArray<float, 16> &right,
+                       const VectorizedArray<float, 16> &true_values,
+                       const VectorizedArray<float, 16> &false_values)
+{
+  const auto mask = _mm512_cmp_ps(left.data, right.data, predicate);
+  VectorizedArray<float, 16> result;
+  result.data = _mm512_or_ps(_mm512_and_ps(mask, true_values.data),
+                             _mm512_andnot_ps(mask, false_values.data));
+  return result;
+}
+
+
+
+template <int predicate>
+DEAL_II_ALWAYS_INLINE inline VectorizedArray<double, 8>
+compare_and_apply_mask(const VectorizedArray<double, 8> &left,
+                       const VectorizedArray<double, 8> &right,
+                       const VectorizedArray<double, 8> &true_values,
+                       const VectorizedArray<double, 8> &false_values)
+{
+  const auto mask = _mm512_cmp_pd(left.data, right.data, predicate);
+  VectorizedArray<double, 8> result;
+  result.data = _mm256_or_ps(_mm512_and_ps(mask, true_values.data),
+                             _mm512_andnot_ps(mask, false_values.data));
+  return result;
+}
+
+#  endif
+
+#  if DEAL_II_COMPILER_VECTORIZATION_LEVEL >= 2 && defined(__AVX__)
+
+template <int predicate>
+DEAL_II_ALWAYS_INLINE inline VectorizedArray<float, 8>
+compare_and_apply_mask(const VectorizedArray<float, 8> &left,
+                       const VectorizedArray<float, 8> &right,
+                       const VectorizedArray<float, 8> &true_values,
+                       const VectorizedArray<float, 8> &false_values)
+{
+  const auto mask = _mm256_cmp_ps(left.data, right.data, predicate);
+
+  VectorizedArray<float, 8> result;
+  result.data = _mm256_or_ps(_mm256_and_ps(mask, true_values.data),
+                             _mm256_andnot_ps(mask, false_values.data));
+  return result;
+}
+
+
+template <int predicate>
+DEAL_II_ALWAYS_INLINE inline VectorizedArray<double, 4>
+compare_and_apply_mask(const VectorizedArray<double, 4> &left,
+                       const VectorizedArray<double, 4> &right,
+                       const VectorizedArray<double, 4> &true_values,
+                       const VectorizedArray<double, 4> &false_values)
+{
+  const auto mask = _mm256_cmp_pd(left.data, right.data, predicate);
+
+  VectorizedArray<double, 4> result;
+  result.data = _mm256_or_ps(_mm256_and_ps(mask, true_values.data),
+                             _mm256_andnot_ps(mask, false_values.data));
+  return result;
+}
+
+#  endif
+
+#  if DEAL_II_COMPILER_VECTORIZATION_LEVEL >= 1 && defined(__SSE2__)
+
+template <int predicate>
+DEAL_II_ALWAYS_INLINE inline VectorizedArray<float, 4>
+compare_and_apply_mask(const VectorizedArray<float, 4> &left,
+                       const VectorizedArray<float, 4> &right,
+                       const VectorizedArray<float, 4> &true_values,
+                       const VectorizedArray<float, 4> &false_values)
+{
+  const auto mask = _mm_cmp_ps(left.data, right.data, predicate);
+
+  VectorizedArray<float, 4> result;
+  result.data = _mm_or_ps(_mm_and_ps(mask, true_values.data),
+                          _mm_andnot_ps(mask, false_values.data));
+  return result;
+}
+
+
+template <int predicate>
+DEAL_II_ALWAYS_INLINE inline VectorizedArray<double, 2>
+compare_and_apply_mask(const VectorizedArray<double, 2> &left,
+                       const VectorizedArray<double, 2> &right,
+                       const VectorizedArray<double, 2> &true_values,
+                       const VectorizedArray<double, 2> &false_values)
+{
+  const auto mask = _mm_cmp_pd(left.data, right.data, predicate);
+
+  VectorizedArray<double, 2> result;
+  result.data = _mm_or_ps(_mm_and_ps(mask, true_values.data),
+                          _mm_andnot_ps(mask, false_values.data));
+  return result;
+}
+
+#  endif
+#endif // DOXYGEN
 
 
 DEAL_II_NAMESPACE_CLOSE