Skip to content

Commit

Permalink
Merge pull request #17 from brenhinkeller/withmallocarray
Browse files Browse the repository at this point in the history 
Support "withmallocarray"-type do-block syntax for most `MallocArray` constructors
  • Loading branch information
@brenhinkeller
brenhinkeller authored Jul 25, 2022
2 parents 3c236e4 + b42c65a commit 796d1ea
Show file tree
Hide file tree
Showing 4 changed files with 194 additions and 19 deletions.
114 changes: 112 additions & 2 deletions src/mallocarray.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,11 @@

"""
```julia
MallocArray(T, dims) do A
...
end
```
```julia
MallocArray{T}(undef, dims)
MallocArray{T,N}(undef, dims)
```
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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)

Expand All @@ -145,6 +168,11 @@
# Other custom constructors
"""
```julia
mzeros([T], dims) do A
...
end
```
```julia
mzeros([T=Float64,] dims::Tuple)
mzeros([T=Float64,] dims...)
```
Expand All @@ -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...)
```
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
21 changes: 21 additions & 0 deletions test/scripts/withmallocarray.jl
View file
Original file line number Diff line number Diff line change
@@ -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}}), "./")
28 changes: 28 additions & 0 deletions test/testmallocarray.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
50 changes: 33 additions & 17 deletions test/teststaticcompiler.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand All @@ -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")
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -134,7 +152,7 @@ end
end

let
# Attempt to compile...
# Compile...
status = -1
try
isfile("loopvec_matrix") && rm("loopvec_matrix")
Expand Down Expand Up @@ -165,7 +183,7 @@ let
end

let
# Attempt to compile...
# Compile...
status = -1
try
isfile("loopvec_matrix_stack") && rm("loopvec_matrix_stack")
Expand Down Expand Up @@ -196,7 +214,7 @@ end
## --- Test string handling

let
# Attempt to compile...
# Compile...
status = -1
try
isfile("print_args") && rm("print_args")
Expand Down Expand Up @@ -224,5 +242,3 @@ end
## --- Clean up

cd(testpath)

## ---

2 comments on commit 796d1ea

@brenhinkeller
Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

Release notes:

  • Add do-block syntax to most MallocArray constructors so users can avoid having to ever manually call free (i.e., the "withmallocarray" trick)
  • Update documentation

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/64973

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.6.1 -m "<description of version>" 796d1ea66c8adfd3a65554ce3027d36ec6b5181e
git push origin v0.6.1

Please sign in to comment.