JuliaGPU/XLA.jl

Matrix * scalar crashed the TPU

ViralBShah opened this issue · 1 comments

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

Keno commented

Now returns a proper error, which is that we're trying to use Int64, rather than crashing:

ERROR: Tensorflow error: Status: While rewriting computation to not contain X64 element types, XLA encountered an HLO for which this rewriting is not implemented: %comp0_map10 = f32[10000,10000]{0,1} map(f32[10000,10000]{0,1} %comp0_parameter7, s64[10000,10000]{0,1} %comp0_broadcast9), dimensions={0,1}, to_apply=%comp1
	 [[{{node XRTCompile_3}}]]

Stacktrace:
 [1] check_status at /home/keno/.julia/packages/TensorFlow/YWnga/src/core.jl:374 [inlined]
 [2] #run#60(::Bool, ::Function, ::Session, ::Array{TensorFlow.Port,1}, ::Array{Any,1}, ::Array{TensorFlow.Port,1}, ::Array{Ptr{Nothing},1}) at /home/keno/.julia/packages/TensorFlow/YWnga/src/run.jl:117
 [3] run at /home/keno/.julia/packages/TensorFlow/YWnga/src/run.jl:90 [inlined]
 [4] #run#69(::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, ::Function, ::Session, ::Array{Tensor,1}, ::Dict{Tensor{String},String}) at /home/keno/.julia/packages/TensorFlow/YWnga/src/run.jl:233
 [5] run(::Session, ::Array{Tensor,1}, ::Dict{Tensor{String},String}) at /home/keno/.julia/packages/TensorFlow/YWnga/src/run.jl:210
 [6] compile(::Session, ::XLA.xrt.XLAComputation, ::Type) at /home/keno/XLAHacks.jl/src/xrt.jl:36
 [7] top-level scope at /home/keno/XLAHacks.jl/src/compiler_interface.jl:91

I'm gonna call this fixed on nightly. Obviously we should make this work also by fixing XLA.