From 3bb242b6f6d4656776aa27373f4f328e51ee4c02 Mon Sep 17 00:00:00 2001 From: Iblis Lin Date: Sat, 21 Dec 2019 18:06:31 +0000 Subject: [PATCH 1/9] julia: porting `current_context` - And introduce a new macro for changing default context `@with_context` --- julia/src/context.jl | 65 +++++++++++++++++++++++++++++++++++++- julia/src/ndarray/array.jl | 18 ++++++----- julia/src/ndarray/type.jl | 2 +- 3 files changed, 75 insertions(+), 10 deletions(-) diff --git a/julia/src/context.jl b/julia/src/context.jl index 68e69138e10e..c48bb4051701 100644 --- a/julia/src/context.jl +++ b/julia/src/context.jl @@ -31,11 +31,50 @@ struct Context Context(dev_type::CONTEXT_TYPE, dev_id::Integer = 0) = new(dev_type, dev_id) end +const _default_ctx = Ref{Context}(Context(CPU, 0)) + Context(dev_type::Integer, dev_id::Integer = 0) = Context(convert(CONTEXT_TYPE, dev_type), dev_id) Base.show(io::IO, ctx::Context) = - print(io, "$(ctx.device_type)$(ctx.device_id)") + print(io, lowercase("$(ctx.device_type)$(ctx.device_id)")) + +function _with_context(dev_type::Expr, dev_id::Integer, e::Expr) + global _default_ctx + quote + ctx = current_context() + ctx′ = Context($dev_type, $dev_id) + $_default_ctx[] = ctx′ + try + return $e + finally + $_default_ctx[] = ctx + end + end +end + +""" + @with_context device_type [device_id] expr + +Change the default context in the following expression. + +# Examples +```jl-repl +julia> mx.@with_context mx.GPU begin + mx.zeros(2, 3) + end +2×3 NDArray{Float32,2} @ gpu0: + 0.0f0 0.0f0 0.0f0 + 0.0f0 0.0f0 0.0f0 +``` +""" +macro with_context(dev_type::Expr, e::Expr) + _with_context(dev_type, 0, e) +end + +macro with_context(dev_type::Expr, dev_id::Integer, e::Expr) + _with_context(dev_type, dev_id, e) +end """ cpu(dev_id) @@ -86,3 +125,27 @@ function gpu_memory_info(dev_id = 0) @mxcall :MXGetGPUMemoryInformation64 (Cint, Ref{UInt64}, Ref{UInt64}) dev_id free n free[], n[] end + +""" + current_context() + +Return the current context. + +By default, `mx.cpu()` is used for all the computations +and it can be overridden by using the `@with_context` macro. + +# Examples +```jl-repl +julia> mx.current_context() +cpu0 + +julia> mx.@with_context mx.GPU 1 begin # Context changed in the following code block + mx.current_context() + end +gpu1 + +julia> mx.current_context() +cpu0 +``` +""" +current_context() = _default_ctx[] diff --git a/julia/src/ndarray/array.jl b/julia/src/ndarray/array.jl index b71e5ddf9397..2cd9c2e24f9c 100644 --- a/julia/src/ndarray/array.jl +++ b/julia/src/ndarray/array.jl @@ -28,13 +28,14 @@ Base.similar(x::NDArray{T,N}; writable = x.writable, ctx = context(x)) where {T, NDArray{T,N}(undef, size(x)...; writable = writable, ctx = ctx) """ - zeros([DType], dims, [ctx::Context = cpu()]) + zeros([DType], dims, ctx::Context = current_context()) zeros([DType], dims...) zeros(x::NDArray) Create zero-ed `NDArray` with specific shape and type. """ -function zeros(::Type{T}, dims::NTuple{N,Int}, ctx::Context = cpu()) where {N,T<:DType} +function zeros(::Type{T}, dims::NTuple{N,Int}, + ctx::Context = current_context()) where {N,T<:DType} x = NDArray{T}(undef, dims..., ctx = ctx) x[:] = zero(T) x @@ -42,7 +43,7 @@ end zeros(::Type{T}, dims::Int...) where {T<:DType} = zeros(T, dims) -zeros(dims::NTuple{N,Int}, ctx::Context = cpu()) where N = +zeros(dims::NTuple{N,Int}, ctx::Context = current_context()) where N = zeros(MX_float, dims, ctx) zeros(dims::Int...) = zeros(dims) @@ -50,13 +51,14 @@ zeros(x::NDArray)::typeof(x) = zeros_like(x) Base.zeros(x::NDArray)::typeof(x) = zeros_like(x) """ - ones([DType], dims, [ctx::Context = cpu()]) + ones([DType], dims, ctx::Context = current_context()) ones([DType], dims...) ones(x::NDArray) Create an `NDArray` with specific shape & type, and initialize with 1. """ -function ones(::Type{T}, dims::NTuple{N,Int}, ctx::Context = cpu()) where {N,T<:DType} +function ones(::Type{T}, dims::NTuple{N,Int}, + ctx::Context = current_context()) where {N,T<:DType} arr = NDArray{T}(undef, dims..., ctx = ctx) arr[:] = one(T) arr @@ -64,7 +66,7 @@ end ones(::Type{T}, dims::Int...) where T<:DType = ones(T, dims) -ones(dims::NTuple{N,Int}, ctx::Context = cpu()) where N = +ones(dims::NTuple{N,Int}, ctx::Context = current_context()) where N = ones(MX_float, dims, ctx) ones(dims::Int...) = ones(dims) @@ -458,12 +460,12 @@ function Base.fill!(arr::NDArray, x) end """ - fill(x, dims, ctx=cpu()) + fill(x, dims, ctx = current_context()) fill(x, dims...) Create an `NDArray` filled with the value `x`, like `Base.fill`. """ -function fill(x::T, dims::NTuple{N,Integer}, ctx::Context = cpu()) where {T,N} +function fill(x::T, dims::NTuple{N,Integer}, ctx::Context = current_context()) where {T,N} arr = NDArray{T}(undef, dims, ctx = ctx) arr[:] = x arr diff --git a/julia/src/ndarray/type.jl b/julia/src/ndarray/type.jl index 8d90d63f0360..e24c89291dcb 100644 --- a/julia/src/ndarray/type.jl +++ b/julia/src/ndarray/type.jl @@ -116,7 +116,7 @@ end # UndefInitializer constructors NDArray{T,N}(::UndefInitializer, dims::NTuple{N,Integer}; - writable = true, ctx::Context = cpu()) where {T,N} = + writable = true, ctx::Context = current_context()) where {T,N} = NDArray{T,N}(_ndarray_alloc(T, dims, ctx, false), writable) NDArray{T,N}(::UndefInitializer, dims::Vararg{Integer,N}; kw...) where {T,N} = NDArray{T,N}(undef, dims; kw...) From a4bec0bd00df7526e48fbe428dd3402f6cedbcc8 Mon Sep 17 00:00:00 2001 From: Iblis Lin Date: Sun, 22 Dec 2019 03:37:48 +0000 Subject: [PATCH 2/9] export --- julia/src/MXNet.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/julia/src/MXNet.jl b/julia/src/MXNet.jl index 86d7d06806fc..25b6ac1c486a 100644 --- a/julia/src/MXNet.jl +++ b/julia/src/MXNet.jl @@ -80,7 +80,9 @@ export Context, cpu, gpu, num_gpus, - gpu_memory_info + gpu_memory_info, + current_context, + @with_context # model.jl export AbstractModel, From 7a09e2e4a7404aabda2e4bd37891551a28e6bc5d Mon Sep 17 00:00:00 2001 From: Iblis Lin Date: Sun, 22 Dec 2019 03:40:14 +0000 Subject: [PATCH 3/9] docstring --- julia/src/context.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/julia/src/context.jl b/julia/src/context.jl index c48bb4051701..3b04b116ea23 100644 --- a/julia/src/context.jl +++ b/julia/src/context.jl @@ -66,6 +66,12 @@ julia> mx.@with_context mx.GPU begin 2×3 NDArray{Float32,2} @ gpu0: 0.0f0 0.0f0 0.0f0 0.0f0 0.0f0 0.0f0 + +julia> @with_context mx.GPU mx.zeros(3, 2) +3×2 NDArray{Float32,2} @ gpu0: + 0.0f0 0.0f0 + 0.0f0 0.0f0 + 0.0f0 0.0f0 ``` """ macro with_context(dev_type::Expr, e::Expr) From 2ae206a14474cd279ba9adcbf2a7930a397bafc1 Mon Sep 17 00:00:00 2001 From: Iblis Lin Date: Sun, 22 Dec 2019 06:40:38 +0000 Subject: [PATCH 4/9] add shorthand macros --- julia/src/MXNet.jl | 4 +++- julia/src/context.jl | 43 +++++++++++++++++++++++++++++++++++++++---- 2 files changed, 42 insertions(+), 5 deletions(-) diff --git a/julia/src/MXNet.jl b/julia/src/MXNet.jl index 25b6ac1c486a..c450ff1b4d70 100644 --- a/julia/src/MXNet.jl +++ b/julia/src/MXNet.jl @@ -82,7 +82,9 @@ export Context, num_gpus, gpu_memory_info, current_context, - @with_context + @with_context, + @with_cpu, + @with_gpu # model.jl export AbstractModel, diff --git a/julia/src/context.jl b/julia/src/context.jl index 3b04b116ea23..6dc0d14b6f21 100644 --- a/julia/src/context.jl +++ b/julia/src/context.jl @@ -39,14 +39,14 @@ Context(dev_type::Integer, dev_id::Integer = 0) = Base.show(io::IO, ctx::Context) = print(io, lowercase("$(ctx.device_type)$(ctx.device_id)")) -function _with_context(dev_type::Expr, dev_id::Integer, e::Expr) +function _with_context(dev_type::Union{Symbol,Expr}, dev_id::Integer, e::Expr) global _default_ctx quote ctx = current_context() ctx′ = Context($dev_type, $dev_id) $_default_ctx[] = ctx′ try - return $e + return $(esc(e)) finally $_default_ctx[] = ctx end @@ -74,14 +74,49 @@ julia> @with_context mx.GPU mx.zeros(3, 2) 0.0f0 0.0f0 ``` """ -macro with_context(dev_type::Expr, e::Expr) +macro with_context(dev_type, e::Expr) _with_context(dev_type, 0, e) end -macro with_context(dev_type::Expr, dev_id::Integer, e::Expr) +macro with_context(dev_type, dev_id::Integer, e::Expr) _with_context(dev_type, dev_id, e) end +for dev ∈ [:cpu, :gpu] + ctx = QuoteNode(Symbol(uppercase(string(dev)))) + fname = Symbol("with_", dev) + docstring = """ + @$fname [device_id] expr + + A shorthand for `@with_context mx.GPU`. + + # Examples + ```jl-repl + julia> mx.@with_gpu mx.zeros(2, 3) + 2×3 NDArray{Float32,2} @ gpu0: + 0.0f0 0.0f0 0.0f0 + 0.0f0 0.0f0 0.0f0 + ``` + """ + @eval begin + @doc $docstring -> + macro $fname(e::Expr) + ctx = $ctx + quote + @with_context $ctx $(esc(e)) + end + end + + macro $fname(dev_id::Integer, e::Expr) + ctx = $ctx + quote + @with_context $ctx $dev_id $(esc(e)) + end + end + end + +end + """ cpu(dev_id) From 48f885de1f49c93cc2b5984a724655d01ca94f12 Mon Sep 17 00:00:00 2001 From: Iblis Lin Date: Sun, 22 Dec 2019 06:53:54 +0000 Subject: [PATCH 5/9] add test cases --- julia/test/unittest/context.jl | 61 ++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/julia/test/unittest/context.jl b/julia/test/unittest/context.jl index 0a8f086a194a..182c6807d7a4 100644 --- a/julia/test/unittest/context.jl +++ b/julia/test/unittest/context.jl @@ -26,8 +26,69 @@ function test_num_gpus() @test num_gpus() >= 0 end +function test_with_context() + @info "Context::@with_context" + + @with_context mx.CPU 42 begin + ctx = mx.current_context() + @test ctx.device_type == mx.CPU + @test ctx.device_id == 42 + + @with_context mx.GPU 24 begin + ctx = mx.current_context() + @test ctx.device_type == mx.GPU + @test ctx.device_id == 42 + end + + ctx = mx.current_context() + @test ctx.device_type == mx.CPU + @test ctx.device_id == 42 + end + + function f() + ctx = mx.current_context() + @test ctx.device_type == mx.GPU + @test ctx.device_id == 123 + end + + @with_context mx.GPU 123 begin + f() + end + + @with_context mx.GPU begin + ctx = mx.current_context() + @test ctx.device_type == mx.GPU + @test ctx.device_id == 0 + end + + @with_context mx.CPU begin + ctx = mx.current_context() + @test ctx.device_type == mx.CPU + @test ctx.device_id == 0 + end + + @with_gpu 123 f() + @with_gpu begin + ctx = mx.current_context() + @test ctx.device_type == mx.GPU + @test ctx.device_id == 0 + end + + @with_cpu 123 begin + ctx = mx.current_context() + @test ctx.device_type == mx.CPU + @test ctx.device_id == 123 + end + @with_cpu begin + ctx = mx.current_context() + @test ctx.device_type == mx.CPU + @test ctx.device_id == 0 + end +end + @testset "Context Test" begin test_num_gpus() + test_with_context() end From 38cc431ac6605968ce9201efddd55090373e9080 Mon Sep 17 00:00:00 2001 From: Iblis Lin Date: Sun, 22 Dec 2019 07:48:12 +0000 Subject: [PATCH 6/9] fix tests --- julia/test/unittest/context.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/julia/test/unittest/context.jl b/julia/test/unittest/context.jl index 182c6807d7a4..1a3f1eb86bd0 100644 --- a/julia/test/unittest/context.jl +++ b/julia/test/unittest/context.jl @@ -37,7 +37,7 @@ function test_with_context() @with_context mx.GPU 24 begin ctx = mx.current_context() @test ctx.device_type == mx.GPU - @test ctx.device_id == 42 + @test ctx.device_id == 24 end ctx = mx.current_context() From 5a3557495e360404306ca95814b497d8596ce193 Mon Sep 17 00:00:00 2001 From: Iblis Lin Date: Sun, 22 Dec 2019 08:35:08 +0000 Subject: [PATCH 7/9] rename --- julia/src/MXNet.jl | 6 +++--- julia/src/context.jl | 29 ++++++++++++++--------------- julia/test/unittest/context.jl | 26 ++++++++++++++------------ 3 files changed, 31 insertions(+), 30 deletions(-) diff --git a/julia/src/MXNet.jl b/julia/src/MXNet.jl index c450ff1b4d70..4c50d7321158 100644 --- a/julia/src/MXNet.jl +++ b/julia/src/MXNet.jl @@ -82,9 +82,9 @@ export Context, num_gpus, gpu_memory_info, current_context, - @with_context, - @with_cpu, - @with_gpu + @context, + @cpu, + @gpu # model.jl export AbstractModel, diff --git a/julia/src/context.jl b/julia/src/context.jl index 6dc0d14b6f21..3f7128b0c2e5 100644 --- a/julia/src/context.jl +++ b/julia/src/context.jl @@ -54,41 +54,40 @@ function _with_context(dev_type::Union{Symbol,Expr}, dev_id::Integer, e::Expr) end """ - @with_context device_type [device_id] expr + @context device_type [device_id] expr Change the default context in the following expression. # Examples ```jl-repl -julia> mx.@with_context mx.GPU begin +julia> mx.@context mx.GPU begin mx.zeros(2, 3) end 2×3 NDArray{Float32,2} @ gpu0: 0.0f0 0.0f0 0.0f0 0.0f0 0.0f0 0.0f0 -julia> @with_context mx.GPU mx.zeros(3, 2) +julia> @context mx.GPU mx.zeros(3, 2) 3×2 NDArray{Float32,2} @ gpu0: 0.0f0 0.0f0 0.0f0 0.0f0 0.0f0 0.0f0 ``` """ -macro with_context(dev_type, e::Expr) +macro context(dev_type, e::Expr) _with_context(dev_type, 0, e) end -macro with_context(dev_type, dev_id::Integer, e::Expr) +macro context(dev_type, dev_id::Integer, e::Expr) _with_context(dev_type, dev_id, e) end for dev ∈ [:cpu, :gpu] ctx = QuoteNode(Symbol(uppercase(string(dev)))) - fname = Symbol("with_", dev) docstring = """ - @$fname [device_id] expr + @$dev [device_id] expr - A shorthand for `@with_context mx.GPU`. + A shorthand for `@context mx.GPU`. # Examples ```jl-repl @@ -100,17 +99,17 @@ for dev ∈ [:cpu, :gpu] """ @eval begin @doc $docstring -> - macro $fname(e::Expr) + macro $dev(e::Expr) ctx = $ctx quote - @with_context $ctx $(esc(e)) + @context $ctx $(esc(e)) end end - macro $fname(dev_id::Integer, e::Expr) + macro $dev(dev_id::Integer, e::Expr) ctx = $ctx quote - @with_context $ctx $dev_id $(esc(e)) + @context $ctx $dev_id $(esc(e)) end end end @@ -172,15 +171,15 @@ end Return the current context. -By default, `mx.cpu()` is used for all the computations -and it can be overridden by using the `@with_context` macro. +By default, `mx.cpu()` is used for all the computations +and it can be overridden by using the `@context` macro. # Examples ```jl-repl julia> mx.current_context() cpu0 -julia> mx.@with_context mx.GPU 1 begin # Context changed in the following code block +julia> mx.@context mx.GPU 1 begin # Context changed in the following code block mx.current_context() end gpu1 diff --git a/julia/test/unittest/context.jl b/julia/test/unittest/context.jl index 1a3f1eb86bd0..f65f1409f5b9 100644 --- a/julia/test/unittest/context.jl +++ b/julia/test/unittest/context.jl @@ -26,15 +26,15 @@ function test_num_gpus() @test num_gpus() >= 0 end -function test_with_context() - @info "Context::@with_context" +function test_context_macro() + @info "Context::@context" - @with_context mx.CPU 42 begin + @context mx.CPU 42 begin ctx = mx.current_context() @test ctx.device_type == mx.CPU @test ctx.device_id == 42 - @with_context mx.GPU 24 begin + @context mx.GPU 24 begin ctx = mx.current_context() @test ctx.device_type == mx.GPU @test ctx.device_id == 24 @@ -51,35 +51,37 @@ function test_with_context() @test ctx.device_id == 123 end - @with_context mx.GPU 123 begin + @context mx.GPU 123 begin f() end - @with_context mx.GPU begin + @context mx.GPU begin ctx = mx.current_context() @test ctx.device_type == mx.GPU @test ctx.device_id == 0 end - @with_context mx.CPU begin + @context mx.CPU begin ctx = mx.current_context() @test ctx.device_type == mx.CPU @test ctx.device_id == 0 end - @with_gpu 123 f() - @with_gpu begin + @info "Context::@gpu" + @gpu 123 f() + @gpu begin ctx = mx.current_context() @test ctx.device_type == mx.GPU @test ctx.device_id == 0 end - @with_cpu 123 begin + @info "Context::@cpu" + @cpu 123 begin ctx = mx.current_context() @test ctx.device_type == mx.CPU @test ctx.device_id == 123 end - @with_cpu begin + @cpu begin ctx = mx.current_context() @test ctx.device_type == mx.CPU @test ctx.device_id == 0 @@ -88,7 +90,7 @@ end @testset "Context Test" begin test_num_gpus() - test_with_context() + test_context_macro() end From 63c1e6b4916e91d3b62b3ca6382af2f9eecbf7c8 Mon Sep 17 00:00:00 2001 From: Iblis Lin Date: Sun, 22 Dec 2019 08:42:32 +0000 Subject: [PATCH 8/9] more tests --- julia/src/context.jl | 10 +++++----- julia/test/unittest/context.jl | 14 ++++++++++++++ julia/test/unittest/ndarray.jl | 2 +- 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/julia/src/context.jl b/julia/src/context.jl index 3f7128b0c2e5..5f8743f5715d 100644 --- a/julia/src/context.jl +++ b/julia/src/context.jl @@ -39,11 +39,11 @@ Context(dev_type::Integer, dev_id::Integer = 0) = Base.show(io::IO, ctx::Context) = print(io, lowercase("$(ctx.device_type)$(ctx.device_id)")) -function _with_context(dev_type::Union{Symbol,Expr}, dev_id::Integer, e::Expr) +function _with_context(dev_type::Union{Symbol,Expr}, dev_id, e::Expr) global _default_ctx quote ctx = current_context() - ctx′ = Context($dev_type, $dev_id) + ctx′ = Context($(esc(dev_type)), $(esc(dev_id))) $_default_ctx[] = ctx′ try return $(esc(e)) @@ -78,7 +78,7 @@ macro context(dev_type, e::Expr) _with_context(dev_type, 0, e) end -macro context(dev_type, dev_id::Integer, e::Expr) +macro context(dev_type, dev_id, e::Expr) _with_context(dev_type, dev_id, e) end @@ -106,10 +106,10 @@ for dev ∈ [:cpu, :gpu] end end - macro $dev(dev_id::Integer, e::Expr) + macro $dev(dev_id, e::Expr) ctx = $ctx quote - @context $ctx $dev_id $(esc(e)) + @context $ctx $(esc(dev_id)) $(esc(e)) end end end diff --git a/julia/test/unittest/context.jl b/julia/test/unittest/context.jl index f65f1409f5b9..e903f9212930 100644 --- a/julia/test/unittest/context.jl +++ b/julia/test/unittest/context.jl @@ -74,6 +74,13 @@ function test_context_macro() @test ctx.device_type == mx.GPU @test ctx.device_id == 0 end + let n = 321 + @gpu n begin + ctx = mx.current_context() + @test ctx.device_type == mx.GPU + @test ctx.device_id == 321 + end + end @info "Context::@cpu" @cpu 123 begin @@ -86,6 +93,13 @@ function test_context_macro() @test ctx.device_type == mx.CPU @test ctx.device_id == 0 end + let n = 321 + @cpu n begin + ctx = mx.current_context() + @test ctx.device_type == mx.CPU + @test ctx.device_id == 321 + end + end end @testset "Context Test" begin diff --git a/julia/test/unittest/ndarray.jl b/julia/test/unittest/ndarray.jl index 599b0a65bfc4..fb59b71edd60 100644 --- a/julia/test/unittest/ndarray.jl +++ b/julia/test/unittest/ndarray.jl @@ -1294,7 +1294,7 @@ function test_show() @test occursin("1×4", str) @test occursin("NDArray", str) @test occursin("Int64", str) - @test occursin("CPU", str) + @test occursin("cpu", str) @test match(r"1\s+2\s+3\s+4", str) != nothing end From 1b658872de43507c089f2365532e7aa8a72393a7 Mon Sep 17 00:00:00 2001 From: Iblis Lin Date: Sun, 22 Dec 2019 14:48:04 +0000 Subject: [PATCH 9/9] style --- julia/src/context.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/julia/src/context.jl b/julia/src/context.jl index 5f8743f5715d..bb80e0728724 100644 --- a/julia/src/context.jl +++ b/julia/src/context.jl @@ -113,8 +113,7 @@ for dev ∈ [:cpu, :gpu] end end end - -end +end # for dev ∈ [:cpu, :gpu] """ cpu(dev_id)