enzymead / reactant.jl Goto Github PK

julia> f(a′, b′)
2048×2048 Tensor{Float64, 2, Reactant.ConcreteRArray{Float64, (2048, 2048), 2}}:
   0.0    233.947  239.041  242.471  244.843  233.684  234.742  242.731  241.478  235.834  246.115  228.468  162.586   153.2     …  241.457   238.066   238.748   231.795   250.899   Error showing value of type Tensor{Float64, 2, Reactant.ConcreteRArray{Float64, (2048, 2048), 2}}:
ERROR: ArgumentError: can't repeat a string -1 times
Stacktrace:
  [1] repeat(s::String, r::Int64)
    @ Base ./strings/substring.jl:263
  [2] repeat
    @ Base ./strings/substring.jl:260 [inlined]
  [3] print_matrix_row(io::IOContext{Base.TTY}, X::AbstractVecOrMat, A::Vector{Tuple{Int64, Int64}}, i::Int64, cols::Vector{Int64}, sep::String, idxlast::Int64)
    @ Base ./arrayshow.jl:118
  [4] _print_matrix(io::IOContext{Base.TTY}, X::AbstractVecOrMat, pre::String, sep::String, post::String, hdots::String, vdots::String, ddots::String, hmod::Int64, vmod::Int64, rowsA::UnitRange{Int64}, colsA::UnitRange{Int64})
    @ Base ./arrayshow.jl:254
  [5] print_matrix(io::IOContext{Base.TTY}, X::Tensor{Float64, 2, Reactant.ConcreteRArray{Float64, (2048, 2048), 2}}, pre::String, sep::String, post::String, hdots::String, vdots::String, ddots::String, hmod::Int64, vmod::Int64)
    @ Base ./arrayshow.jl:171
  [6] print_matrix(io::IO, X::AbstractVecOrMat, pre::AbstractString, sep::AbstractString, post::AbstractString, hdots::AbstractString, vdots::AbstractString, ddots::AbstractString, hmod::Integer, vmod::Integer)
    @ Base ./arrayshow.jl:171 [inlined]
  [7] print_array
    @ ./arrayshow.jl:358 [inlined]
  [8] show(io::IOContext{Base.TTY}, ::MIME{Symbol("text/plain")}, X::Tensor{Float64, 2, Reactant.ConcreteRArray{Float64, (2048, 2048), 2}})
    @ Base ./arrayshow.jl:399
  [9] (::REPL.var"#55#56"{REPL.REPLDisplay{REPL.LineEditREPL}, MIME{Symbol("text/plain")}, Base.RefValue{Any}})(io::Any)
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:273
 [10] with_repl_linfo(f::Any, repl::REPL.LineEditREPL)
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:569
 [11] display(d::REPL.REPLDisplay, mime::MIME{Symbol("text/plain")}, x::Any)
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:259
 [12] display(d::REPL.REPLDisplay, x::Any)
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:278
 [13] display(x::Any)
    @ Base.Multimedia ./multimedia.jl:340
 [14] #invokelatest#2
    @ ./essentials.jl:887 [inlined]
 [15] invokelatest
    @ ./essentials.jl:884 [inlined]
 [16] print_response(errio::IO, response::Any, show_value::Bool, have_color::Bool, specialdisplay::Union{Nothing, AbstractDisplay})
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:315
 [17] (::REPL.var"#57#58"{REPL.LineEditREPL, Pair{Any, Bool}, Bool, Bool})(io::Any)
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:284
 [18] with_repl_linfo(f::Any, repl::REPL.LineEditREPL)
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:569
 [19] print_response(repl::REPL.AbstractREPL, response::Any, show_value::Bool, have_color::Bool)
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:282
 [20] (::REPL.var"#do_respond#80"{Bool, Bool, REPL.var"#93#103"{REPL.LineEditREPL, REPL.REPLHistoryProvider}, REPL.LineEditREPL, REPL.LineEdit.Prompt})(s::REPL.LineEdit.MIState, buf::Any, ok::Bool)
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:911
 [21] (::VSCodeServer.var"#103#106"{REPL.var"#do_respond#80"{Bool, Bool, REPL.var"#93#103"{REPL.LineEditREPL, REPL.REPLHistoryProvider}, REPL.LineEditREPL, REPL.LineEdit.Prompt}})(mi::REPL.LineEdit.MIState, buf::IOBuffer, ok::Bool)
    @ VSCodeServer ~/.vscode-server/extensions/julialang.language-julia-1.79.2/scripts/packages/VSCodeServer/src/repl.jl:122
 [22] #invokelatest#2
    @ Base ./essentials.jl:887 [inlined]
 [23] invokelatest
    @ Base ./essentials.jl:884 [inlined]
 [24] run_interface(terminal::REPL.Terminals.TextTerminal, m::REPL.LineEdit.ModalInterface, s::REPL.LineEdit.MIState)
    @ REPL.LineEdit /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/LineEdit.jl:2656
 [25] run_frontend(repl::REPL.LineEditREPL, backend::REPL.REPLBackendRef)
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:1312
 [26] (::REPL.var"#62#68"{REPL.LineEditREPL, REPL.REPLBackendRef})()
    @ REPL /gpfs/apps/MN5/GPP/JULIA/1.10.0/INTEL/share/julia/stdlib/v1.10/REPL/src/REPL.jl:386

using Reactant
using Cassette

struct Tensor{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
           data::A
           inds::Vector{Symbol}
       end

Tensor(data::A, inds::AbstractVector{Symbol}) where {T,N,A<:AbstractArray{T,N}} = Tensor{T,N,A}(data, inds)

Base.parent(x::Tensor) = x.data
Base.size(t::Tensor) = size(parent(t))
Base.@propagate_inbounds Base.getindex(t::Tensor, i...) = getindex(parent(t), i...)

n = 2048
a = Tensor(rand(n,n), [:i, :j]);
b = Tensor(rand(n,n), [:j, :k])

a′ = Tensor(Reactant.ConcreteRArray(a.data), a.inds);
b′ = Tensor(Reactant.ConcreteRArray(b.data), b.inds);

contract(a,b) = a.data * b.data
function contract(a::Tensor{Ta,Na,Aa}, b::Tensor{Tb,Nb,Ab}) where {Ta,Na,Aa<:Reactant.TracedRArray,Tb,Nb,Ab<:Reactant.TracedRArray}
	ia = collect(a.inds)
	ib = collect(b.inds)
	i = ∩(ia, ib)
	
	ic::Vector{Symbol} = setdiff(ia ∪ ib, i)
	
	T = Base.promote_eltype(a, b)
	mlirty = Reactant.MLIR.IR.Type(T)
	
	op_a = parent(a).mlir_data
	op_b = parent(b).mlir_data
	rsize = (size(a.data, 1), size(b.data, 2))
	result_0 = Reactant.MLIR.IR.TensorType(rsize, mlirty)
	einsum_config = Reactant.MLIR.IR.Attribute("$(join(ia)),$(join(ib))->$(join(ic))")
	
	result = Reactant.MLIR.IR.result(Reactant.MLIR.Dialects.stablehlo.einsum(op_a, op_b; result_0, einsum_config))
	
	data = Reactant.TracedRArray{T,rsize,length(ic)}((), result)
	_res = Tensor(data, ic)
	return _res
end

f = Reactant.compile(contract, (a′,b′))

f(a′,b′)

using Reactant
using Cassette

n = 2048
a = rand(n, n);
b = rand(n, n);

a′ = Reactant.ConcreteRArray(a);
b′ = Reactant.ConcreteRArray(b);

matmul(a,b) = a * b

function matmul(a::Reactant.TracedRArray{Ta,Sa,Na}, b::Reactant.TracedRArray{Tb,Sb,Nb}) where {Ta,Tb,Sa,Sb,Na,Nb}
	T = Base.promote_type(Ta,Tb)
	mlirty = Reactant.MLIR.IR.Type(T)
	rsize = (Sa[1], Sb[2])
	result_0 = Reactant.MLIR.IR.TensorType(rsize, mlirty)
	einsum_config = Reactant.MLIR.IR.Attribute("ij,jk->ik")
	
	result = Reactant.MLIR.IR.result(Reactant.MLIR.Dialects.stablehlo.einsum(a.mlir_data, b.mlir_data; result_0, einsum_config))
	
	return Reactant.TracedRArray{T,rsize,2}((), result)
end

Cassette.overdub(ctx::Reactant.TraceCtx, f::typeof(matmul), args...; kwargs...) = f(args...; kwargs...)

f = Reactant.compile(matmul, (a′,b′))

f(a′,b′)

julia> using Reactant

julia> Reactant.ConcreteRArray(rand(10, 3))
10×3 Reactant.ConcreteRArray{Float64, (10, 3), 2}:
 0.241899   0.892     0.344609
 0.658603   0.978326  0.0684158
 0.154474   0.506356  0.737901
 0.665978   0.635999  0.312166
 0.955006   0.613347  0.834144
 0.261397   0.361902  0.758361
 0.838547   0.939787  0.191731
 0.336063   0.649798  0.305722
 0.0539672  0.754429  0.444449
 0.964148   0.524124  0.628569

julia> Reactant.ConcreteRArray(rand(10, 3)) .+ 1
10×3 Matrix{Float64}:
 1.60667  1.2144   1.15662
 1.79119  1.18538  1.81092
 1.01788  1.73736  1.10406
 1.88946  1.51043  1.7367
 1.11927  1.46719  1.82236
 1.18602  1.32522  1.6286
 1.12319  1.79284  1.85603
 1.42769  1.33347  1.96626
 1.40952  1.39544  1.52497
 1.23689  1.28314  1.1096

Precompiling project...
  ✗ Reactant → ReactantNNlibExt
  1 dependency successfully precompiled in 7 seconds. 127 already precompiled.
  1 dependency had output during precompilation:
┌ Reactant
│  error in running finalizer: TypeError(func=Symbol("ccall: first argument not a pointer or valid constant expression"), context="", expected=Ptr{T} where T, got=(:PjRtBufferFree, 0x000000008fc4a150))
│  ijl_type_error_rt at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-HL2F7YQ3XH.0/build/default-honeycrisp-HL2F7YQ3XH-0/julialang/julia-release-1-dot-10/src/rtutils.c:119
│  ijl_type_error at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-HL2F7YQ3XH.0/build/default-honeycrisp-HL2F7YQ3XH-0/julialang/julia-release-1-dot-10/src/rtutils.c:127
│  #5 at /Users/mofeing/Developer/Reactant.jl/src/XLA.jl:114
│  unknown function (ip: 0x105c080e3)
│  _jl_invoke at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-HL2F7YQ3XH.0/build/default-honeycrisp-HL2F7YQ3XH-0/julialang/julia-release-1-dot-10/src/gf.c:0 [inlined]
│  ijl_apply_generic at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-HL2F7YQ3XH.0/build/default-honeycrisp-HL2F7YQ3XH-0/julialang/julia-release-1-dot-10/src/gf.c:3077
│  run_finalizer at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-HL2F7YQ3XH.0/build/default-honeycrisp-HL2F7YQ3XH-0/julialang/julia-release-1-dot-10/src/gc.c:318
│  jl_gc_run_finalizers_in_list at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-HL2F7YQ3XH.0/build/default-honeycrisp-HL2F7YQ3XH-0/julialang/julia-release-1-dot-10/src/gc.c:408
│  run_finalizers at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-HL2F7YQ3XH.0/build/default-honeycrisp-HL2F7YQ3XH-0/julialang/julia-release-1-dot-10/src/gc.c:454
│  ijl_atexit_hook at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-HL2F7YQ3XH.0/build/default-honeycrisp-HL2F7YQ3XH-0/julialang/julia-release-1-dot-10/src/init.c:299
│  jl_repl_entrypoint at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-HL2F7YQ3XH.0/build/default-honeycrisp-HL2F7YQ3XH-0/julialang/julia-release-1-dot-10/src/jlapi.c:732
└  
     Testing Running tests...
libc++abi: terminating due to uncaught exception of type xla::XlaRuntimeError: NOT_FOUND: Could not find registered platform with name: "cuda". Available platform names are: 

[42937] signal (6): Abort trap: 6
in expression starting at /Users/mofeing/Developer/Reactant.jl/test/runtests.jl:1
__pthread_kill at /usr/lib/system/libsystem_kernel.dylib (unknown line)
Allocations: 2908 (Pool: 2900; Big: 8); GC: 0
ERROR: Package Reactant errored during testing (received signal: 6)

using Reactant, Enzyme

function simple_reduce(w, x)
    y = sum(w * x; dims=1)
    return sum(y)
end

w = Reactant.ConcreteRArray(randn(Float32, 10, 3))
x = Reactant.ConcreteRArray(randn(Float32, 3, 5))

simple_reduce_xla = Reactant.compile(simple_reduce, (w, x))

simple_reduce_xla(w, x) # Works

function simple_reduce_grad(w, x)
    dw = Enzyme.make_zero(w)
    dx = Enzyme.make_zero(x)
    Enzyme.autodiff(
        Enzyme.Reverse, simple_reduce, Active, Duplicated(w, dw), Duplicated(x, dx))
    return dw, dx
end

simple_reduce_grad_xla = Reactant.compile(simple_reduce_grad, (w, x))

simple_reduce_grad_xla(w, x)

function simple_reduce(w, x)
    return sum(w * x)
end

using Reactant, Lux, Random
using Test

# Generate some data for the XOR problem: vectors of length 2, as columns of a matrix:
noisy = rand(Float32, 2, 1000)                                        # 2×1000 Matrix{Float32}
truth = [xor(col[1] > 0.5, col[2] > 0.5) for col in eachcol(noisy)]   # 1000-element Vector{Bool}

# Define our model, a multi-layer perceptron with one hidden layer of size 3:
model = Chain(Dense(2 => 3, tanh),   # activation function inside layer
    BatchNorm(3), Dense(3 => 2), softmax)
ps, st = Lux.setup(Xoshiro(123), model)

using BenchmarkTools

origout, _ = model(noisy, ps, st)
@show origout[3]
@btime model($noisy, $ps, $st)

cmodel = Reactant.make_tracer(IdDict(), model, (), Reactant.ArrayToConcrete, nothing)
cps = Reactant.make_tracer(IdDict(), ps, (), Reactant.ArrayToConcrete, nothing)
cst = Reactant.make_tracer(IdDict(), st, (), Reactant.ArrayToConcrete, nothing)
cnoisy = Reactant.ConcreteRArray(noisy)

f = Reactant.compile((a, b, c, d) -> a(b, c, d), (cmodel, cnoisy, cps, CST))

using Reactant

noisy = rand(Float32, 2, 1000)
cnoisy = Reactant.ConcreteRArray(noisy)
similar(cnoisy)  # Matrix