From f27c5ad652818c011ea6f5d57e746844267be9c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 28 Oct 2026 14:07:57 +0100
Subject: [PATCH] Support buffered dot product for any element type

`buffer_for(LinearAlgebra.dot, ...)` and `buffered_operate_to!(buffer, output,
LinearAlgebra.dot, x, y)` were only implemented for `BigFloat` vectors. For
other element types, the buffer of the scalar `add_dot` is now returned and
reused for every entry so that computing several dot products with the same
buffer and output does not allocate. `operate_to!(output, LinearAlgebra.dot,
x, y)` is also implemented.
---
 src/implementations/LinearAlgebra.jl | 42 ++++++++++++++++
 test/test_big.jl                     | 75 ++++++++++++++++++++++++++++
 2 files changed, 117 insertions(+)

diff --git a/src/implementations/LinearAlgebra.jl b/src/implementations/LinearAlgebra.jl
index f1f3852..3a214f8 100644
--- a/src/implementations/LinearAlgebra.jl
+++ b/src/implementations/LinearAlgebra.jl
@@ -558,3 +558,45 @@ function operate(
 )
     return fused_map_reduce(add_dot, x, y)
 end
+
+# The buffer of the scalar `add_dot` operation can be reused for every entry, so
+# computing several dot products with the same buffer allocates nothing.
+function buffer_for(
+    ::typeof(LinearAlgebra.dot),
+    ::Type{<:AbstractArray{S}},
+    ::Type{<:AbstractArray{T}},
+) where {S,T}
+    U = promote_operation(LinearAlgebra.dot, S, T)
+    return buffer_for(add_dot, U, S, T)
+end
+
+function _buffered_dot_to!(buffer, output, x, y)
+    _check_same_length(x, y)
+    operate!(zero, output)
+    for (xi, yi) in zip(x, y)
+        buffered_operate!(buffer, add_dot, output, xi, yi)
+    end
+    return output
+end
+
+function operate_to_fallback!(
+    ::IsMutable,
+    output,
+    ::typeof(LinearAlgebra.dot),
+    x::AbstractArray,
+    y::AbstractArray,
+)
+    buffer = buffer_for(LinearAlgebra.dot, typeof(x), typeof(y))
+    return _buffered_dot_to!(buffer, output, x, y)
+end
+
+function buffered_operate_to_fallback!(
+    ::IsMutable,
+    buffer,
+    output,
+    ::typeof(LinearAlgebra.dot),
+    x::AbstractArray,
+    y::AbstractArray,
+)
+    return _buffered_dot_to!(buffer, output, x, y)
+end
diff --git a/test/test_big.jl b/test/test_big.jl
index 9cd8672..f3fb2ae 100644
--- a/test/test_big.jl
+++ b/test/test_big.jl
@@ -317,6 +317,81 @@ function test_alloc()
     return
 end
 
+function _buffered_dot_helper(buffer, output, x, y)
+    let b = buffer, o = output, x = x, y = y
+        () -> MA.buffered_operate_to!!(b, o, LinearAlgebra.dot, x, y)
+    end
+end
+
+function test_buffered_dot()
+    for (S, T) in (
+        (BigInt, BigInt),
+        (BigInt, Int),
+        (Rational{BigInt}, Rational{BigInt}),
+        (Rational{BigInt}, BigInt),
+        (BigFloat, Float64),
+    )
+        x = S[1, 2, -3]
+        y = T[4, -5, 6]
+        expected = LinearAlgebra.dot(x, y)
+        U = typeof(expected)
+        buffer = MA.buffer_for(LinearAlgebra.dot, typeof(x), typeof(y))
+        output = U(7)
+        @test MA.buffered_operate_to!(
+            buffer,
+            output,
+            LinearAlgebra.dot,
+            x,
+            y,
+        ) === output
+        @test output == expected
+        # Called a second time to check that the output is reset to zero
+        @test MA.buffered_operate_to!!(
+            buffer,
+            output,
+            LinearAlgebra.dot,
+            x,
+            y,
+        ) === output
+        @test output == expected
+        output = U(7)
+        @test MA.operate_to!(output, LinearAlgebra.dot, x, y) === output
+        @test output == expected
+        @test x == S[1, 2, -3]
+        @test y == T[4, -5, 6]
+        if S == T
+            # There is no buffer for the scalar `add_mul` of mixed types
+            alloc_test(_buffered_dot_helper(buffer, output, x, y), 0)
+        end
+        @test_throws DimensionMismatch MA.buffered_operate_to!(
+            buffer,
+            output,
+            LinearAlgebra.dot,
+            x,
+            T[1, 2],
+        )
+    end
+    # Types without a buffer and types that cannot be mutated
+    @test MA.buffer_for(LinearAlgebra.dot, Vector{Int}, Vector{Int}) === nothing
+    @test MA.buffered_operate_to!!(
+        nothing,
+        0,
+        LinearAlgebra.dot,
+        [1, 2],
+        [3, 4],
+    ) == 11
+    output = big(0)
+    @test MA.buffered_operate_to!(
+        nothing,
+        output,
+        LinearAlgebra.dot,
+        big.([1, 2]),
+        [3, 4],
+    ) === output
+    @test output == 11
+    return
+end
+
 function test_buffered_BigInt_mixed_integers()
     for op in (MA.add_mul, MA.sub_mul),
         args in ((big(2), 3, big(4)), (2, big(3), 4), (big(2), 3, big(4), 5))