Merge pull request #104 from JuliaParallel/vc/scale

add scale! for Vector and DMatrix

Author: vchuravy

src/linalg.jl

@@ -165,6 +165,23 @@ function Ac_mul_B!(α::Number, A::DMatrix, x::AbstractVector, β::Number, y::DVe
     return y
 end
 
+function Base.LinAlg.scale!(b::AbstractVector, DA::DMatrix)
+    s = verified_destination_serializer(procs(DA), size(DA.indexes)) do pididx
+        b[DA.indexes[pididx][1]]
+    end
+    map_localparts!(DA) do lDA
+        scale!(localpart(s), lDA)
+    end
+end
+
+function Base.LinAlg.scale!(DA::DMatrix, b::AbstractVector)
+    s = verified_destination_serializer(procs(DA), size(DA.indexes)) do pididx
+        b[DA.indexes[pididx][2]]
+    end
+    map_localparts!(DA) do lDA
+        scale!(lDA, localpart(s))
+    end
+end
 
 # Level 3
 function _matmatmul!(α::Number, A::DMatrix, B::AbstractMatrix, β::Number, C::DMatrix, tA)

test/darray.jl

@@ -198,6 +198,21 @@ end
 
 check_leaks()
 
+@testset "test scale!(b, A)" begin
+    A = randn(100, 100)
+    b = randn(100)
+    DA = distribute(A)
+    @test scale!(b, A) == scale!(b, DA)
+    close(DA)
+    A = randn(100, 100)
+    b = randn(100)
+    DA = distribute(A)
+    @test scale!(A, b) == scale!(DA, b)
+    close(DA)
+end
+
+check_leaks()
+
 @testset "test mapreduce on DArrays" begin
     for _ = 1:25, f = [x -> Int128(2x), x -> Int128(x^2), x -> Int128(x^2 + 2x - 1)], opt = [+, *]
         A = rand(1:5, rand(2:30))