Merge pull request #17 from brenhinkeller/withmallocarray

Support "withmallocarray"-type do-block syntax for most `MallocArray` constructors

Author: brenhinkeller

src/mallocarray.jl

@@ -38,6 +38,11 @@
 
     """
     ```julia
+    MallocArray(T, dims) do A
+        ...
+    end
+    ```
+    ```julia
     MallocArray{T}(undef, dims)
     MallocArray{T,N}(undef, dims)
     ```
@@ -82,6 +87,17 @@
     julia> free(A)
     0
     ```
+    To avoid having to manually `free` allocated memory, it is recommended to
+    use the following supported do-block syntax whenever possible, i.e.
+    ```julia
+    julia> MallocArray(Float64, 2, 2) do A
+               A .= 0
+               printf(A)
+           end
+    0.000000e+00    0.000000e+00
+    0.000000e+00    0.000000e+00
+    0
+    ```
     """
     @inline function MallocArray{T,N}(::UndefInitializer, length::Int, dims::Dims{N}) where {T,N}
         @assert Base.allocatedinline(T)
@@ -99,6 +115,12 @@
     @inline MallocArray{T}(x::PointerOrInitializer, dims::Dims{N}) where {T,N} = MallocArray{T,N}(x, prod(dims), dims)
     @inline MallocArray{T,N}(x::PointerOrInitializer, dims::Vararg{Int}) where {T,N} = MallocArray{T,N}(x, prod(dims), dims)
     @inline MallocArray{T}(x::PointerOrInitializer, dims::Vararg{Int}) where {T} = MallocArray{T}(x, dims)
+    @inline function MallocArray(f::Function, T::Type, dims...)
+        M = MallocArray{T}(undef, dims...)
+        y = f(M)
+        free(M)
+        y
+    end
 
     """
     ```julia
@@ -127,6 +149,7 @@
     """
     @inline MallocArray(x::AbstractArray{T,N}) where {T,N} = copyto!(MallocArray{T,N}(undef, length(x), size(x)), x)
 
+
     # Destructor:
     @inline free(a::MallocArray) = free(a.pointer)
 
@@ -145,6 +168,11 @@
     # Other custom constructors
     """
     ```julia
+    mzeros([T], dims) do A
+        ...
+    end
+    ```
+    ```julia
     mzeros([T=Float64,] dims::Tuple)
     mzeros([T=Float64,] dims...)
     ```
@@ -160,14 +188,35 @@
      0  0
      0  0
     ```
+    To avoid having to manually `free` allocated memory, it is recommended to
+    use the following supported do-block syntax whenever possible, i.e.
+    ```julia
+    julia> mzeros(2,2) do A
+               printf(A)
+           end
+    0.000000e+00    0.000000e+00
+    0.000000e+00    0.000000e+00
+    0
+    ```
     """
     @inline mzeros(dims::Vararg{Int}) = mzeros(dims)
     @inline mzeros(dims::Dims{N}) where {N} = MallocArray{Float64,N}(zeros, dims)
     @inline mzeros(T::Type, dims::Vararg{Int}) = mzeros(T, dims)
     @inline mzeros(::Type{T}, dims::Dims{N}) where {T,N} = MallocArray{T,N}(zeros, dims)
+    @inline function mzeros(f::Function, args...)
+        M = mzeros(args...)
+        y = f(M)
+        free(M)
+        y
+    end
 
     """
     ```julia
+    mones([T=Float64,] dims) do A
+        ...
+    end
+    ```
+    ```julia
     mones([T=Float64,] dims::Tuple)
     mones([T=Float64,] dims...)
     ```
@@ -178,28 +227,66 @@
 
     ## Examples
     ```julia
-    julia> mones(Int32, 2,2)
+    julia> A = mones(Int32, 2,2)
     2×2 MallocMatrix{Int32}:
      1  1
      1  1
+
+    julia> free(A)
+    0
+    ```
+    To avoid having to manually `free` allocated memory, it is recommended to
+    use the following supported do-block syntax whenever possible, i.e.
+    ```julia
+    julia> mones(2,2) do A
+               printf(A)
+           end
+    1.000000e+00    1.000000e+00
+    1.000000e+00    1.000000e+00
+    0
     ```
     """
     @inline mones(dims...) = mones(Float64, dims...)
     @inline mones(::Type{T}, dims...) where {T} = mfill(one(T), dims...)
+    @inline function mones(f::Function, args...)
+        M = mones(args...)
+        y = f(M)
+        free(M)
+        y
+    end
 
     """
     ```julia
+    meye([T=Float64,] dims) do A
+        ...
+    end
+    ```
+    ```julia
     meye([T=Float64,] dim::Int)
     ```
     Create a `MallocArray{T}` containing an identity matrix of type `T`,
     of size `dim` x `dim`.
 
     ## Examples
     ```julia
-    julia> meye(Int32, 2)
+    julia> A = meye(Int32, 2)
+
     2×2 MallocMatrix{Int32}:
      1  0
      0  1
+
+    julia> free(A)
+    0
+    ```
+    To avoid having to manually `free` allocated memory, it is recommended to
+    use the following supported do-block syntax whenever possible, i.e.
+    ```julia
+    julia> meye(2) do A
+               printf(A)
+           end
+    1.000000e+00    0.000000e+00
+    0.000000e+00    1.000000e+00
+    0
     ```
     """
     @inline meye(dim::Int) = meye(Float64, dim)
@@ -211,6 +298,13 @@
         end
         return A
     end
+    @inline function meye(f::Function, args...)
+        M = meye(args...)
+        y = f(M)
+        free(M)
+        y
+    end
+
 
     """
     ```julia
@@ -229,9 +323,25 @@
      3  3
      3  3
     ```
+    To avoid having to manually `free` allocated memory, it is recommended to
+    use the following supported do-block syntax whenever possible, i.e.
+    ```julia
+    julia> mfill(1.5, 2,2) do A
+               printf(A)
+           end
+    1.500000e+00    1.500000e+00
+    1.500000e+00    1.500000e+00
+    0
+    ```
     """
     @inline mfill(x, dims::Vararg{Int}) = mfill(x, dims)
     @inline function mfill(x::T, dims::Dims{N}) where {T,N}
         A = MallocArray{T,N}(undef, prod(dims), dims)
         fill!(A, x)
     end
+    @inline function mfill(f::Function, args...)
+        M = mfill(args...)
+        y = f(M)
+        free(M)
+        y
+    end

test/scripts/withmallocarray.jl

@@ -0,0 +1,21 @@
+using StaticCompiler
+using StaticTools
+
+function withmallocarray(argc::Int, argv::Ptr{Ptr{UInt8}})
+    argc == 3 || return printf(stderrp(), c"Incorrect number of command-line arguments\n")
+    rows = parse(Int64, argv, 2)            # First command-line argument
+    cols = parse(Int64, argv, 3)            # Second command-line argument
+
+    mzeros(rows, cols) do A
+        printf(A)
+    end
+    mones(Int, rows, cols) do A
+        printf(A)
+    end
+    mfill(3.141592, rows, cols) do A
+        printf(A)
+    end
+end
+
+# Attempt to compile
+path = compile_executable(withmallocarray, (Int64, Ptr{Ptr{UInt8}}), "./")

test/testmallocarray.jl

@@ -203,3 +203,31 @@
     @test A == B
     free(A)
     free(B)
+
+## --- test "withmallocarray" pattern
+
+    s = MallocArray(Int64, 2, 2) do A
+            A .= 1
+            sum(A)
+        end
+    @test s === 4
+
+    s = mones(2, 2) do A
+            sum(A)
+        end
+    @test s === 4.0
+
+    s = mzeros(2, 2) do A
+            sum(A)
+        end
+    @test s === 0.0
+
+    s = meye(2) do A
+            sum(A)
+        end
+    @test s === 2.0
+
+    s = mfill(3.14, 2, 2) do A
+            sum(A)
+        end
+    @test s === 12.56

test/teststaticcompiler.jl

@@ -4,7 +4,6 @@ scratch = tempdir()
 cd(scratch)
 
 ## --- Times table, file IO, mallocarray
-
 let
     # Attempt to compile
     # We have to start a new Julia process to get around the fact that Pkg.test
@@ -38,13 +37,36 @@ let
     @test fread!(szeros(Int, 5,5), c"table.b") == (1:5)*(1:5)'
 end
 
-## --- Random number generation
+## --- "withmallocarray"-type do-block pattern
+let
+    # Compile...
+    status = -1
+    try
+        isfile("withmallocarray") && rm("withmallocarray")
+        status = run(`julia --compile=min $testpath/scripts/withmallocarray.jl`)
+    catch e
+        @warn "Could not compile $testpath/scripts/withmallocarray.jl"
+        println(e)
+    end
+    @test isa(status, Base.Process)
+    @test isa(status, Base.Process) && status.exitcode == 0
+
+    # Run...
+    println("5x5 uniform random matrix:")
+    status = -1
+    try
+        status = run(`./withmallocarray 5 5`)
+    catch e
+        @warn "Could not run $(scratch)/withmallocarray"
+        println(e)
+    end
+    @test isa(status, Base.Process)
+    @test isa(status, Base.Process) && status.exitcode == 0
+end
 
+## --- Random number generation
 let
-    # Attempt to compile...
-    # We have to start a new Julia process to get around the fact that Pkg.test
-    # disables `@inbounds`, but ironically we can use `--compile=min` to make that
-    # faster.
+    # Compile...
     status = -1
     try
         isfile("rand_matrix") && rm("rand_matrix")
@@ -70,10 +92,7 @@ let
 end
 
 let
-    # Attempt to compile...
-    # We have to start a new Julia process to get around the fact that Pkg.test
-    # disables `@inbounds`, but ironically we can use `--compile=min` to make that
-    # faster.
+    # Compile...
     status = -1
     try
         isfile("randn_matrix") && rm("randn_matrix")
@@ -103,10 +122,9 @@ let
 end
 
 ## --- Test LoopVectorization integration
-
 @static if LoopVectorization.VectorizationBase.has_feature(Val{:x86_64_avx2})
     let
-        # Attempt to compile...
+        # Compile...
         status = -1
         try
             isfile("loopvec_product") && rm("loopvec_product")
@@ -134,7 +152,7 @@ end
 end
 
 let
-    # Attempt to compile...
+    # Compile...
     status = -1
     try
         isfile("loopvec_matrix") && rm("loopvec_matrix")
@@ -165,7 +183,7 @@ let
 end
 
 let
-    # Attempt to compile...
+    # Compile...
     status = -1
     try
         isfile("loopvec_matrix_stack") && rm("loopvec_matrix_stack")
@@ -196,7 +214,7 @@ end
 ## --- Test string handling
 
 let
-    # Attempt to compile...
+    # Compile...
     status = -1
     try
         isfile("print_args") && rm("print_args")
@@ -224,5 +242,3 @@ end
 ## --- Clean up
 
 cd(testpath)
-
-## ---