Core: Use a scalar tail for inner-unrolled partial-packet assignments

libeigen/eigen!2866

Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com>

Eigen/src/Core/AssignEvaluator.h

@@ -132,12 +132,14 @@ struct copy_using_evaluator_traits {
   static constexpr int ActualPacketSize = Vectorized ? unpacket_traits<PacketType>::size : 1;
   static constexpr int UnrollingLimit = EIGEN_UNROLLING_LIMIT * ActualPacketSize;
   static constexpr int CoeffReadCost = int(DstEvaluator::CoeffReadCost) + int(SrcEvaluator::CoeffReadCost);
-  static constexpr bool MayUnrollCompletely =
-      (SizeAtCompileTime != Dynamic) && (SizeAtCompileTime * CoeffReadCost <= UnrollingLimit);
   static constexpr bool MayUnrollInner =
       (InnerSizeAtCompileTime != Dynamic) && (InnerSizeAtCompileTime * CoeffReadCost <= UnrollingLimit);
 
  public:
+  // True when the whole assignment is a fixed-size kernel cheap enough to emit as straight-line
+  // code. Selects CompleteUnrolling, and gates the scalar tail in the SliceVectorized loop below.
+  static constexpr bool MayUnrollCompletely =
+      (SizeAtCompileTime != Dynamic) && (SizeAtCompileTime * CoeffReadCost <= UnrollingLimit);
   static constexpr int Unrolling =
       (Traversal == InnerVectorizedTraversal || Traversal == DefaultTraversal)
           ? (MayUnrollCompletely ? CompleteUnrolling
@@ -704,7 +706,14 @@ struct dense_assignment_loop_impl<Kernel, SliceVectorizedTraversal, InnerUnrolli
   static constexpr int PacketSize = unpacket_traits<PacketType>::size;
   static constexpr int InnerSize = Kernel::AssignmentTraits::InnerSizeAtCompileTime;
   static constexpr int VectorizableSize = numext::round_down(InnerSize, PacketSize);
-  static constexpr bool UsePacketSegment = Kernel::AssignmentTraits::UsePacketSegment;
+
+  // The tail length is a compile-time constant here, so it can be emitted as a masked packet
+  // segment or as scalars. Scalars win for the same fixed-size kernels LinearVectorizedTraversal
+  // already emits them for: their destination is a temporary the enclosing expression reloads by
+  // packet, and on AVX a masked store forwards poorly to those loads. Everything else keeps the
+  // segment, where one packet evaluation replaces up to PacketSize - 1 scalar ones.
+  static constexpr bool UsePacketSegment =
+      Kernel::AssignmentTraits::UsePacketSegment && !Kernel::AssignmentTraits::MayUnrollCompletely;
 
   using packet_loop = copy_using_evaluator_innervec_InnerUnrolling<Kernel, 0, VectorizableSize, Unaligned, Unaligned>;
   using packet_segment_loop = copy_using_evaluator_innervec_segment<Kernel, VectorizableSize, InnerSize, Unaligned,

benchmarks/Core/bench_segment_tail.cpp

@@ -21,7 +21,15 @@
 // ops where each result is consumed by the next. A partial-packet tail store
 // then collides with the consumer's packet load -- the store-to-load
 // forwarding hazard that makes a masked tail far more costly than a scalar
-// one. This is the cost Parts A and B cannot see in isolation.
+// one. This is the cost Parts A and B cannot see in isolation. Chained/Rotation
+// additionally covers the lazy-product shape, whose tail is reached through
+// SliceVectorizedTraversal with InnerUnrolling and through the product
+// evaluator's own partial loads, not through LinearVectorizedTraversal.
+//
+// Part D takes the same traversal as Chained/Rotation but streams over a runtime
+// outer size with nothing reloading the destination, so the segment amortizes a
+// source evaluation instead of stalling a consumer. It is the shape that decides
+// how narrowly the scalar tail may be applied.
 //
 // The active packet size and the value of has_packet_segment per scalar type
 // are emitted as Google Benchmark custom context, so a captured run is
@@ -418,6 +426,66 @@ void BM_Chained_Block23(benchmark::State& state) {
   }
 }
 
+// #3083 Example A: the matrix algebra of one IMU integration step. Its 3x3
+// products are lazy, so unlike the kernels above the destination assignment
+// takes SliceVectorizedTraversal with InnerUnrolling, and the product evaluator
+// itself loads a partial packet. Neither path was covered by the kernels above,
+// which is why this one kept a masked tail after !2581.
+void BM_Chained_Rotation(benchmark::State& state) {
+  Matrix3d S = Matrix3d::Random(), R = Matrix3d::Random(), Xi = Matrix3d::Random();
+  Vector3d a = Vector3d::Random(), v = Vector3d::Zero();
+  const double dt = 0.005;
+  for (auto _ : state) {
+    benchmark::DoNotOptimize(S);
+    benchmark::DoNotOptimize(R);
+    benchmark::DoNotOptimize(a);
+    const Matrix3d S2 = S * S;
+    Xi = dt * Matrix3d::Identity() + 0.5 * dt * dt * S + (dt * dt * dt / 6.0) * S2;
+    v += R * Xi * a;
+    benchmark::DoNotOptimize(Xi);
+    benchmark::DoNotOptimize(v);
+    benchmark::ClobberMemory();
+  }
+}
+
+// ===========================================================================
+// Part D : streaming slices (compile-time inner size, runtime outer size)
+//
+// SliceVectorizedTraversal with InnerUnrolling, as in Chained/Rotation, but the
+// tail is paid once per column of a long runtime outer loop and no consumer
+// reloads the destination. Sqrt makes one packet evaluation of the source far
+// cheaper than PacketSize scalar ones, so unlike Part C this shape wants the
+// masked segment; the copy variant is the cheap-source control.
+// ===========================================================================
+
+template <typename T, int Rows>
+void BM_SliceCopy(benchmark::State& state) {
+  const Index cols = state.range(0);
+  using Mat = Matrix<T, Dynamic, Dynamic>;
+  Mat src = Mat::Random(Rows + 9, cols), dst = Mat::Zero(Rows + 9, cols);
+  for (auto _ : state) {
+    benchmark::DoNotOptimize(src.data());
+    dst.template topRows<Rows>() = src.template topRows<Rows>();
+    benchmark::DoNotOptimize(dst.data());
+    benchmark::ClobberMemory();
+  }
+  state.SetBytesProcessed(state.iterations() * Rows * cols * static_cast<int64_t>(sizeof(T)) * 2);
+}
+
+template <typename T, int Rows>
+void BM_SliceSqrt(benchmark::State& state) {
+  const Index cols = state.range(0);
+  using Mat = Matrix<T, Dynamic, Dynamic>;
+  Mat src = Mat::Random(Rows + 9, cols).cwiseAbs(), dst = Mat::Zero(Rows + 9, cols);
+  for (auto _ : state) {
+    benchmark::DoNotOptimize(src.data());
+    dst.template topRows<Rows>() = src.template topRows<Rows>().array().sqrt().matrix();
+    benchmark::DoNotOptimize(dst.data());
+    benchmark::ClobberMemory();
+  }
+  state.SetBytesProcessed(state.iterations() * Rows * cols * static_cast<int64_t>(sizeof(T)) * 2);
+}
+
 // ===========================================================================
 // Registration
 //
@@ -528,6 +596,18 @@ void add_all_prim_ct() {
   add_prim_ct<T, 7>();
 }
 
+// Part D : one fixed inner size that is not a packet multiple for any supported
+// packet width, swept over outer sizes that span the cache hierarchy.
+template <typename T>
+void add_slice() {
+  const std::string tag = std::string("/") + type_tag<T>() + "/7";
+  for (auto entry : {std::make_pair("Slice/copy" + tag, &BM_SliceCopy<T, 7>),
+                     std::make_pair("Slice/sqrt" + tag, &BM_SliceSqrt<T, 7>)}) {
+    auto* b = benchmark::RegisterBenchmark(entry.first, entry.second);
+    for (int c : {64, 1024, 4096}) b->Arg(c);
+  }
+}
+
 int RegisterAll() {
   add_trait_context<float>("f32");
   add_trait_context<double>("f64");
@@ -554,6 +634,10 @@ int RegisterAll() {
   benchmark::RegisterBenchmark("Chained/Inverse3x3", &BM_Chained_Inverse3x3);
   benchmark::RegisterBenchmark("Chained/Camera", &BM_Chained_Camera);
   benchmark::RegisterBenchmark("Chained/Block23", &BM_Chained_Block23);
+  benchmark::RegisterBenchmark("Chained/Rotation", &BM_Chained_Rotation);
+
+  add_slice<float>();
+  add_slice<double>();
   return 0;
 }