juliagpu / xla.jl Goto Github PK

View Code? Open in Web Editor NEW

224.0 19.0 20.0 8.13 MB

Julia on TPUs

License: Other

Julia 99.74% TeX 0.26%

julia-language tpu going-faster xla peanut-butter machine-learning deep-learning

xla.jl's Introduction

XLA.jl - Compiling Julia to XLA

NOTE: We're in the process of adding better instructions. Check back in a bit.

Getting started on TPUs

Running on Colab

Google currently offers free access to Cloud TPUs through its Colab notebook service. Colab does not officially support julia at the moment, but it is possible to install julia by manually installing it into the runtime (though this has to be done every time the runtime gets reset). By this mechanism, you can get access to TPUs through the notebook interface. Start with the installation notebook in docs/colab/InstallJuliaXLA.ipynb:

Open in Colab
Run in Google Colab

Afterwards, any JuliaTPU notebook should work if opened without resetting the runtime in between.

Running on GCP

The process for setting up this repository to run against TPUs is much the same as the process for setting up the repository locally. However, since there is additional steps involved in launching the actual TPU, we are providing a tutorial to walk you through all the steps. It is recommended for those new to Julia and/or TPUs. If you're already familiar with both, you may skip the tutorial and just use the setup guide below. The tutorial will open in Google Cloud Shell, by clicking the button below:

Getting started (CPU/GPU backend)

Grab julia on branch kf/tpu3 (Prebuilt Linux x86_64 binaries with TPU support are available here)
Instantiate this repo
julia> using TensorFlow
Get yourself an xrt_server (either running it locally via run(`$(joinpath(dirname(pathof(TensorFlow)),"..","deps","downloads","bin","xrt_server"))`) or by spinning up a Google Cloud TPU and starting an SSH tunnel to expose its port to the world) and connect it to localhost:8470
Run the script in examples/vgg_forward.jl

xla.jl's People

Contributors

Stargazers

Watchers

xla.jl's Issues

Include dumped_computation_to_text in tensorflow download

Add a test for the OneHotMatrix off by one

Error when activating XLA in Google Colab

Hi, I tried running the Julia TPU notebook after running the installation one but I found the below error:

ERROR: LoadError: LoadError: syntax: type declarations on global variables are not yet supported
Stacktrace:
[1] eval(::Module, ::Any) at ./boot.jl:319
[2] top-level scope at none:0
[3] include at ./boot.jl:317 [inlined]
[4] include_relative(::Module, ::String) at ./loading.jl:1038
[5] include at ./sysimg.jl:29 [inlined]
[6] include(::String) at /content/XLA.jl/src/XLA.jl:1
[7] top-level scope at none:0
[8] include at ./boot.jl:317 [inlined]
[9] include_relative(::Module, ::String) at ./loading.jl:1038
[10] include(::Module, ::String) at ./sysimg.jl:29
[11] top-level scope at none:2
[12] eval at ./boot.jl:319 [inlined]
[13] eval(::Expr) at ./client.jl:389
[14] top-level scope at ./none:3
in expression starting at /content/XLA.jl/src/linalg.jl:108
in expression starting at /content/XLA.jl/src/XLA.jl:25
Failed to precompile XLA [1ae4bca4-de81-11e8-0eca-6d3e4e7c4181] to /root/.julia/compiled/v1.1/XLA/bZBiw.ji.

This is my info of my notebook:
versioninfo()

Julia Version 1.1.0-DEV.576
Commit fff87381740 (2018-11-19 09:29 UTC)
Platform Info:
OS: Linux (x86_64-linux-gnu)
CPU: Intel(R) Xeon(R) CPU @ 2.30GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-6.0.1 (ORCJIT, haswell)

Julia install on Google Colab doesn't work

InstallJuliaXLA.ipynb fails with the following error:

Updating registry at `~/.julia/registries/General`
  Updating git-repo `https://github.com/JuliaRegistries/General.git`
 Resolving package versions...
  Updating `~/.julia/environments/v1.1/Project.toml`
 [no changes]
  Updating `~/.julia/environments/v1.1/Manifest.toml`
 [no changes]
Already up to date.
  Updating registry at `~/.julia/registries/General`
  Updating git-repo `https://github.com/JuliaRegistries/General.git`
ERROR: AssertionError: haskey(hashes, uuid)
Stacktrace:
 [1] macro expansion at ./logging.jl:308 [inlined]
 [2] version_data!(::Pkg.Types.Context, ::Array{Pkg.Types.PackageSpec,1}) at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Pkg/src/Operations.jl:406
 [3] #instantiate#59(::Nothing, ::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, ::Function, ::Pkg.Types.Context) at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Pkg/src/API.jl:559
 [4] instantiate at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Pkg/src/API.jl:528 [inlined]
 [5] do_instantiate!(::Dict{Symbol,Any}, ::Array{String,1}, ::Dict{Symbol,Any}) at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Pkg/src/REPLMode.jl:671
 [6] #invokelatest#1(::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, ::Function, ::Any, ::Any, ::Vararg{Any,N} where N) at ./essentials.jl:697
 [7] invokelatest(::Any, ::Any, ::Vararg{Any,N} where N) at ./essentials.jl:696
 [8] do_cmd!(::Pkg.REPLMode.PkgCommand, ::Pkg.REPLMode.MiniREPL) at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Pkg/src/REPLMode.jl:603
 [9] #do_cmd#33(::Bool, ::Function, ::Pkg.REPLMode.MiniREPL, ::String) at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Pkg/src/REPLMode.jl:577
 [10] (::getfield(Pkg.REPLMode, Symbol("#kw##do_cmd")))(::NamedTuple{(:do_rethrow,),Tuple{Bool}}, ::typeof(Pkg.REPLMode.do_cmd), ::Pkg.REPLMode.MiniREPL, ::String) at ./none:0
 [11] top-level scope at none:0

Where are the prebuilt TPU binaries for julia 1.3.x

Possible TPU crash

Base.@pure calc_rounds(L; exponent=3) = ceil(UInt32, exponent * log(L) / log(typemax(UInt32)))
function shuffle(x::XRTArray)
    rounds = calc_rounds(length(x))
    round_idx = XRTArray(rounds)
    while round_idx > XRTArray(0)
        keys = rand(XRTArray{UInt32}, length(x))
        x = sort(x, keys=keys)
        # Commenting this out crashes the TPU
        #round_idx -= XRTArray(1)
    end
    return x
end

XRTArray(XRTArray(1)) doesn't work

Currently it throws the opaque error:

AssertionError: isexpr(cond_stmt, :invoke)

Stacktrace:
 [1] outline_if!(::Core.Compiler.IRCode, ::Core.Compiler.OptimizationState, ::Core.Compiler.DomTree, ::Int64, ::Int64) at /home/sabae/.julia/dev/XLA/src/outlining.jl:468
 [2] outline_control_flow!(::Core.Compiler.IRCode, ::Core.Compiler.OptimizationState) at /home/sabae/.julia/dev/XLA/src/outlining.jl:503
 [3] _compile_to_xla!(::Array{XLA.xla.HloComputationProto,1}, ::XLA.xla.HloComputationProto, ::Core.Compiler.IRCode, ::Core.Compiler.OptimizationState) at /home/sabae/.julia/dev/XLA/src/compiler.jl:57
 [4] #compile_to_xla#220(::Nothing, ::Function, ::Core.Compiler.IRCode, ::Core.Compiler.OptimizationState) at /home/sabae/.julia/dev/XLA/src/compiler.jl:342
 [5] (::getfield(XLA, Symbol("#kw##compile_to_xla")))(::NamedTuple{(:replica_device_coords,),Tuple{Nothing}}, ::typeof(XLA.compile_to_xla), ::Core.Compiler.IRCode, ::Core.Compiler.OptimizationState) at ./none:0
 [6] top-level scope at /home/sabae/.julia/dev/XLA/src/compiler_interface.jl:115
 [7] top-level scope at In[25]:8

This should just return XRTArray(1).

Colab notebook startup needs to link to the other JuliaTPU.ipynb notebook

After the InstallJuliaXLA.ipynb is run in the first cell, a link to JuliaTPU.ipynb is required to start the second Julia notebook in the container.

XLA.jl recompiles every time I start it

viralbshah@arctic1:~/XLA.jl$ ../julia/julia --project=.
               _
   _       _ _(_)_     |  Documentation: https://docs.julialang.org
  (_)     | (_) (_)    |
   _ _   _| |_  __ _   |  Type "?" for help, "]?" for Pkg help.
  | | | | | | |/ _` |  |
  | | |_| | | | (_| |  |  Version 1.1.0-DEV.567 (2018-11-01)
 _/ |\__'_|_|_|\__'_|  |  kf/tpu3/25d8567707 (fork: 13 commits, 6 days)
|__/                   |

julia> using TensorFlow

julia> using XLA
[ Info: Precompiling XLA [1ae4bca4-de81-11e8-0eca-6d3e4e7c4181]
┌ Warning: Package XLA does not have LinearAlgebra in its dependencies:
│ - If you have XLA checked out for development and have
│   added LinearAlgebra as a dependency but haven't updated your primary
│   environment's manifest file, try `Pkg.resolve()`.
│ - Otherwise you may need to report an issue with XLA
└ Loading LinearAlgebra into XLA from project dependency, future warnings for XLA are suppressed.
WARNING: could not import NNlib.cudata into Tracker
WARNING: Method definition _forward(Zygote.Context{T} where T<:Union{Nothing, Base.IdDict{Any, Any}}, typeof(Base.prod), Any) in module Zygote at /home/viralbshah/.julia/packages/Zygote/0lzma/src/lib/grad.jl:55 overwritten at /home/viralbshah/.julia/packages/Zygote/0lzma/src/lib/grad.jl:21.

julia> 
viralbshah@arctic1:~/XLA.jl$ ../julia/julia --project=.
               _
   _       _ _(_)_     |  Documentation: https://docs.julialang.org
  (_)     | (_) (_)    |
   _ _   _| |_  __ _   |  Type "?" for help, "]?" for Pkg help.
  | | | | | | |/ _` |  |
  | | |_| | | | (_| |  |  Version 1.1.0-DEV.567 (2018-11-01)
 _/ |\__'_|_|_|\__'_|  |  kf/tpu3/25d8567707 (fork: 13 commits, 6 days)
|__/                   |

julia> using TensorFlow

julia> using XLA
[ Info: Recompiling stale cache file /home/viralbshah/.julia/compiled/v1.1/XLA/bZBiw.ji for XLA [1ae4bca4-de81-11e8-0eca-6d3e4e7c4181]
┌ Warning: Package XLA does not have LinearAlgebra in its dependencies:
│ - If you have XLA checked out for development and have
│   added LinearAlgebra as a dependency but haven't updated your primary
│   environment's manifest file, try `Pkg.resolve()`.
│ - Otherwise you may need to report an issue with XLA
└ Loading LinearAlgebra into XLA from project dependency, future warnings for XLA are suppressed.

Running an XRT computation with 0 arguments does not work


julia> f() = XLA.HloRng(Float32, (5,5), 1)(XRTArray(0f0),XRTArray(1f0))
f (generic function with 1 method)

julia> @tpu_compile f()
XLA.XRTCompilation(Session(Ptr{Nothing} @0x00007f87b0290e50), nothing, XLA.xla.ProgramShape(XLA.xla.Shape[], XLA.xla.Shape(11, [5, 5], #undef, XLA.xla.Layout(1, [0, 1], #undef, 0, 0)), #undef), XRTArray{Float32,(5, 5),2}, 2205288722480676172)

julia> run(ans)
ERROR: MethodError: run(::XLA.XRTCompilation) is ambiguous. Candidates:
  run(::XLA.XRTCompilation, xrt::Union{Complex{Float32}, Bool, Float16, Float32, Float64, Int16, Int32, Int64, Int8, UInt16, UInt32, UInt64, UInt8, XRTArray}...) in XLA at /home/viralbshah/XLAHacks.jl/src/execute.jl:45
  run(::XLA.XRTCompilation, com::XRTAllocation...) in XLA at /home/viralbshah/XLAHacks.jl/src/xrt.jl:77
Possible fix, define
  run(::XLA.XRTCompilation)
Stacktrace:
 [1] top-level scope at none:0

ProcessExitedException() when calling train.minimize

I am unable to run the example code from TensorFlow.jl without getting a ProcessExitedException() when calling train.minimize() in the CoLab environment

Input Code:

using TensorFlow
using MLDatasets
import Random

mutable struct DataLoader
    cur_id::Int
    order::Vector{Int}
end

DataLoader() = DataLoader(1, Random.shuffle(1:60000))

function next_batch(loader::DataLoader, batch_size)
    x = zeros(Float32, batch_size, 784)
    y = zeros(Float32, batch_size, 10)
    for i in 1:batch_size
        data, label = MLDatasets.MNIST.traindata(loader.order[loader.cur_id])
        x[i, :] = reshape(data, (28*28))
        y[i, Int(label)+1] = 1.0
        loader.cur_id += 1
        if loader.cur_id > 60000
            loader.cur_id = 1
        end
    end
    x, y
end

function load_test_set(N=10000)
    x = zeros(Float32, N, 784)
    y = zeros(Float32, N, 10)
    for i in 1:N
        data, label = MLDatasets.MNIST.testdata(i)
        x[i, :] = reshape(data, (28*28))
        y[i, Int(label)+1] = 1.0
    end
    x,y
end


loader = DataLoader()

sess = Session(Graph(); target="grpc://"*ENV["COLAB_TPU_ADDR"])

x = placeholder(Float32)
y_ = placeholder(Float32)

W = Variable(zeros(Float32, 784, 10))
b = Variable(zeros(Float32, 10))

run(sess, global_variables_initializer())

y = nn.softmax(x*W + b)

cross_entropy = reduce_mean(-reduce_sum(y_ .* log(y), axis=[2]))
train_step = train.minimize(train.GradientDescentOptimizer(.00001), cross_entropy)

correct_prediction = argmax(y, 2) .== argmax(y_, 2)
accuracy=reduce_mean(cast(correct_prediction, Float32))

for i in 1:1000
    batch = next_batch(loader, 100)
    run(sess, train_step, Dict(x=>batch[1], y_=>batch[2]))
end

testx, testy = load_test_set()

println(run(sess, accuracy, Dict(x=>testx, y_=>testy)))

Output:

┌ Info: Precompiling MLDatasets [eb30cadb-4394-5ae3-aed4-317e484a6458]
└ @ Base loading.jl:1186
┌ Warning: Module Compat with build ID 1186987443319 is missing from the cache.
│ This may mean Compat [34da2185-b29b-5c13-b0c7-acf172513d20] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:947
┌ Warning: Module Compat with build ID 1186987443319 is missing from the cache.
│ This may mean Compat [34da2185-b29b-5c13-b0c7-acf172513d20] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:947
┌ Info: Precompiling BinDeps [9e28174c-4ba2-5203-b857-d8d62c4213ee]
└ @ Base loading.jl:1186
┌ Warning: Module Compat with build ID 1186987443319 is missing from the cache.
│ This may mean Compat [34da2185-b29b-5c13-b0c7-acf172513d20] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:947
┌ Info: Precompiling URIParser [30578b45-9adc-5946-b283-645ec420af67]
└ @ Base loading.jl:1186
┌ Info: Precompiling DataDeps [124859b0-ceae-595e-8997-d05f6a7a8dfe]
└ @ Base loading.jl:1186
┌ Info: Precompiling GZip [92fee26a-97fe-5a0c-ad85-20a5f3185b63]
└ @ Base loading.jl:1186
ProcessExitedException()

Stacktrace:
 [1] worker_from_id(::Distributed.ProcessGroup, ::Int64) at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Distributed/src/cluster.jl:975
 [2] worker_from_id at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Distributed/src/cluster.jl:972 [inlined]
 [3] #remotecall_wait#157(::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, ::Function, ::Function, ::Int64) at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Distributed/src/remotecall.jl:433
 [4] remotecall_wait(::Function, ::Int64) at /workspace/srcdir/julia/usr/share/julia/stdlib/v1.1/Distributed/src/remotecall.jl:433
 [5] top-level scope at /root/.julia/packages/TensorFlow/YWnga/src/TensorFlow.jl:187
 [6] eval at ./boot.jl:319 [inlined]
 [7] eval at ./sysimg.jl:68 [inlined]
 [8] add_gradients_py(::Tensor{Float32}, ::Array{Any,1}, ::Nothing) at /root/.julia/packages/TensorFlow/YWnga/src/core.jl:1548
 [9] gradients at /root/.julia/packages/TensorFlow/YWnga/src/core.jl:1536 [inlined] (repeats 2 times)
 [10] compute_gradients(::TensorFlow.train.GradientDescentOptimizer, ::Tensor{Float32}, ::Nothing) at /root/.julia/packages/TensorFlow/YWnga/src/train.jl:49
 [11] #minimize#1(::Nothing, ::Nothing, ::Nothing, ::Function, ::TensorFlow.train.GradientDescentOptimizer, ::Tensor{Float32}) at /root/.julia/packages/TensorFlow/YWnga/src/train.jl:41
 [12] minimize(::TensorFlow.train.GradientDescentOptimizer, ::Tensor{Float32}) at /root/.julia/packages/TensorFlow/YWnga/src/train.jl:38
 [13] top-level scope at In[14]:54

`setindex` is zero based

Because that's what XLA's semantics is. We should add the appropriate offset.

Warning compiling XLA

julia> using XLA
[ Info: Precompiling XLA [1ae4bca4-de81-11e8-0eca-6d3e4e7c4181]
┌ Warning: Package XLA does not have LinearAlgebra in its dependencies:
│ - If you have XLA checked out for development and have
│   added LinearAlgebra as a dependency but haven't updated your primary
│   environment's manifest file, try `Pkg.resolve()`.
│ - Otherwise you may need to report an issue with XLA
└ Loading LinearAlgebra into XLA from project dependency, future warnings for XLA are suppressed.

Elliot, add test for broken reshape!

@staticfloat

Unable to compile XLA in JuliaTPU.ipynb

When trying to use the provided demo after running the InstallJuliaXLA.jl script I get this error in JuliaTPU.ipynb:

┌ Info: Precompiling XLA [1ae4bca4-de81-11e8-0eca-6d3e4e7c4181]
└ @ Base loading.jl:1186
┌ Warning: Loading a new version of TensorFlow.jl for the first time. This initial load can take around 5 minutes as code is precompiled; subsequent usage will only take a few seconds.
└ @ TensorFlow ~/.julia/packages/TensorFlow/YWnga/src/TensorFlow.jl:3
WARNING: could not import xla.Shape into XLA
WARNING: could not import NNlib.cudata into Tracker
WARNING: Method definition _forward(Zygote.Context{T} where T<:Union{Nothing, Base.IdDict{Any, Any}}, typeof(Base.prod), Any) in module Zygote at /root/.julia/packages/Zygote/lAhbY/src/lib/grad.jl:55 overwritten at /root/.julia/packages/Zygote/lAhbY/src/lib/grad.jl:21.
ERROR: LoadError: LoadError: LoadError: @code_ir f(args...)
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] code_irm(::Expr) at /root/.julia/packages/Zygote/lAhbY/src/tools/reflection.jl:9
 [3] @adjoint(::LineNumberNode, ::Module, ::Any) at /root/.julia/packages/Zygote/lAhbY/src/compiler/reverse.jl:398
 [4] eval(::Module, ::Any) at ./boot.jl:319
 [5] top-level scope at /content/XLA.jl/src/grad.jl:30
 [6] include at ./boot.jl:317 [inlined]
 [7] include_relative(::Module, ::String) at ./loading.jl:1038
 [8] include at ./sysimg.jl:29 [inlined]
 [9] include(::String) at /content/XLA.jl/src/XLA.jl:1
 [10] top-level scope at none:0
 [11] include at ./boot.jl:317 [inlined]
 [12] include_relative(::Module, ::String) at ./loading.jl:1038
 [13] include(::Module, ::String) at ./sysimg.jl:29
 [14] top-level scope at none:2
 [15] eval at ./boot.jl:319 [inlined]
 [16] eval(::Expr) at ./client.jl:389
 [17] top-level scope at ./none:3
in expression starting at /content/XLA.jl/src/grad.jl:31
in expression starting at /content/XLA.jl/src/grad.jl:28
in expression starting at /content/XLA.jl/src/XLA.jl:31
Failed to precompile XLA [1ae4bca4-de81-11e8-0eca-6d3e4e7c4181] to /root/.julia/compiled/v1.1/XLA/bZBiw.ji.

Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] macro expansion at ./logging.jl:313 [inlined]
 [3] compilecache(::Base.PkgId, ::String) at ./loading.jl:1184
 [4] macro expansion at ./logging.jl:311 [inlined]
 [5] _require(::Base.PkgId) at ./loading.jl:947
 [6] require(::Base.PkgId) at ./loading.jl:858
 [7] macro expansion at ./logging.jl:311 [inlined]
 [8] require(::Module, ::Symbol) at ./loading.jl:840
 [9] top-level scope at In[3]:3

CI

We should have some. Ideally with TPUs.

Improve printing of topology information

Currently we say things like "<8 TPU chips in 2x2x2 topology>", which is informative, but doesn't match Google terminology. We could try printing something like "<TPUv3-8 in 2x2 topology>". We don't get the chip kind information from the topology, but we could for now just determine it from the memory information.

XRTReal

Flux also used to use TrackedScalar{T} = TrackedArray{T,0}, but it turns out 0D arrays are not treated consistently like scalars in a lot of places, and it's generally unhelpful for dispatch. Having a separate XRTReal <: Real object would simplify a lot of stuff (e.g. a lot more of Zygote would work for free), so it seems worth considering unless there's a specific reason not too.

Ignore errors from finalizers

E.g. somebody else may have reset the TPU underneath of us, which is fine since our allocations are dead anyway.

Allocation transfers too slow

The actual network transfer takes only about 2-3 secs for our VGG model. However, the whole transfer call takes about 30 secs. We should figure out where that time goes.

Matrix * scalar crashed the TPU

I was doing a bunch of other random stuff before, so I am not sure if this is reproducible.

julia> f(x) = sum(x*XRTArray(2))
f (generic function with 2 methods)

julia> c = @tpu_compile f(XRTArray(rand(Float32,10000,10000)))
ERROR: Tensorflow error: Status: Socket closed

Is this project abandoned?

Given that there seems to be no activity in issues or commits?

XRTRemoteStruct on multiple cores

Now that multicore support kind of works, we should integrate it with the allocation handles. We should even be able to distribute models using the TPU fabric to avoid having to distribute it over the data center network.

TensorFlow build issue

Building TensorFlow ──────→ ~/.julia/packages/TensorFlow/kl0uH/deps/build.log
┌ Error: Error building TensorFlow:
│ [ Info: Building TensorFlow.jl for CPU use only. To enable the GPU, set the TF_USE_GPU environment variable to 1 and rebuild TensorFlow.jl
│ ERROR: LoadError: Python TensorFlow not installed
│ Please either:
│ - Rebuild PyCall to use Conda, by running in the julia REPL:
│ - ENV["PYTHON"]=""; Pkg.build("PyCall"); Pkg.build("TensorFlow")
│ - Or install the python binding yourself, eg by running pip
│ - pip install tensorflow
│ - then rebuilding TensorFlow.jl via Pkg.build("TensorFlow") in the julia REPL
│ - make sure you run the right pip, for the instance of python that PyCall is looking at.
│
│ Stacktrace:
│ [1] error(::String) at ./error.jl:33
│ [2] top-level scope at /home/viralbshah/.julia/packages/TensorFlow/kl0uH/deps/build.jl:48
│ [3] include at ./boot.jl:317 [inlined]
│ [4] include_relative(::Module, ::String) at ./loading.jl:1038
│ [5] include(::Module, ::String) at ./sysimg.jl:29
│ [6] include(::String) at ./client.jl:388
│ [7] top-level scope at none:0
│ in expression starting at /home/viralbshah/.julia/packages/TensorFlow/kl0uH/deps/build.jl:39
└ @ Pkg.Operations ~/julia/usr/share/julia/stdlib/v1.1/Pkg/src/Operations.jl:1069

Upstream dependencies

Keeping track of all the things I needed to fix upstream, so I don't forget anything:

Zygote

More stuff... (fill me in)

Flux

More stuff... (fill me in)

Julia

Done

Zygote

FluxML/Zygote.jl#25

Flux

FluxML/Flux.jl#353

ForwardDiff

JuliaDiff/ForwardDiff.jl#216

Julia

TensorFlow version mismatch

I started with TF 1.12.

julia> f
f (generic function with 1 method)

julia> @tpu_compile f(XRTArray(rand(Float32,1000,1000)))
1 1 ─ %1 = invoke $(QuoteNode(XLA.HloMap{typeof(+)}()))(+::typeof(+), _2::XRTArray{Float32,(1000, 1000),2}, _2::Vararg{XRTArray{Float32,(1000, 1000),2},N} where N)::XRTArray{Float32,(1000, 1000),2}
  └──      return %1                                                                           │    

57 1 ─ %1 = invoke GenericHloOp{:add}(Float32, ())(_2::XRTArray{Float32,(),0}, _3::Vararg{XRTArray{Float32,(),0},N} where N)::XRTArray{Float32,(),0}
   └──      return %1                                                                                │

ERROR: Tensorflow error: Status: FetchOutputs XRTCompile_3:1: output index too large, must be < 1

Convolution test suite

The mapping of our convolution semantics to HLO are complicated. We should have a test suite to make sure it's correct.

"Run in Google Colab" notebook no longer works

"Run in Google Colab" notebook no longer works as the julia-tpu-binaries can no longer be downloaded from:

https://storage.googleapis.com/julia-tpu-binaries/julia.v1.1.0-kf.tpu3.x86_64-linux-gnu.tar.gz

You get the following error when trying to download:

<Error>
  <Code>UserProjectAccountProblem</Code>
  <Message>User project billing account not in good standing.</Message>
  <Details>
    The billing account for the owning project is disabled in state closed
  </Details>
</Error>

Flux's Chain is now immutable

We may be able to remove the ImmutableChain abstraction.

Random array creation did not work

julia> f(x) = XLA.HloRng(Float32, (5,5), 1)(XRTArray(0f0),XRTArray(1f0))
f (generic function with 2 methods)

julia> @tpu_compile f(XRTArray(1));

julia> run(ans, XLA.gethandle!(sess, XRTArray(1)))
5×5 XRTArray{Float32,(5, 5),2}:
 0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0

julia>

Add TPU profiler client

Having to install the cloud-tpu-profiler package is a pain, because it pulls in all of tensorflow (and takes a really long time to start up). It's also unnecessary, because all it does is issue one grpc request to the profiler service and then wait until it's done. We should be able to do that ourselves.

explain should include fixes for Core._apply()

I had those fixes and then I didn't. Here one moment, then gone the next, like grains of sand on the wind.

juliagpu / xla.jl Goto Github PK

xla.jl's Introduction

XLA.jl - Compiling Julia to XLA

Getting started on TPUs

Running on Colab

Running on GCP

Getting started (CPU/GPU backend)

xla.jl's People

Contributors

Stargazers

Watchers

Forkers

xla.jl's Issues

Zygote

Flux

Julia

Done

Zygote

Flux

ForwardDiff

Julia

Recommend Projects

Recommend Topics

Recommend Org