Interp2 (EnzymeAD#365)

* WIP: kernels

* more files

* fix

* wip

* wqtmp

* wip

* inc

* continuing

* wip

* more work

* inf rec

* fix

* overload working

* continuing

* continuing

* push

* fix `call_with_reactant_generator` for Julia 1.11 (EnzymeAD#359)

* conversion

* continuing

* Cleanup

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Delete test/cuda.jl

* fixup

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fix apply

* indep of change

* minor fix in name

* Update utils.jl

* Interp take 2

* continuing adentures

* delcode

* fix

* tmp

* make

* fix

* cleanup

* continuing

* more working

* further simplify

* fx

* more improvements

* minus show

* less prints

* even fewer

* confusion

* tmp

* force clean

* force oc

* clean

* Rewrite

* fixup

* fix

* fix

* fix

* fixup

* fix

* wip

* safe prints

* fix

* fix

* stackoverflow

* cleanup

* dyindex

* rt

* continue

* clean

* fix

* fix

* fix

* fix

* fixup

* fix

* fix

* capture oc

* compile perf

* v1.11 fix

* other way 'round

* formatting

---------

Co-authored-by: William Moses <[email protected]>
Co-authored-by: jumerckx <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: jumerckx <[email protected]>

Author: 4 people

Project.toml

@@ -41,14 +41,14 @@ Adapt = "4"
 ArrayInterface = "7.10"
 CEnum = "0.4, 0.5"
 Downloads = "1.6"
-Enzyme = "0.13.21"
+Enzyme = "0.13.22"
 EnzymeCore = "0.8.8"
 GPUArraysCore = "0.1.6, 0.2"
 LinearAlgebra = "1.10"
 NNlib = "0.9.26"
 OrderedCollections = "1"
 Preferences = "1.4"
-ReactantCore = "0.1.2"
+ReactantCore = "0.1.3"
 Reactant_jll = "0.0.26"
 Scratch = "1.2"
 Statistics = "1.10"

deps/ReactantExtra/API.cpp

@@ -376,6 +376,16 @@ extern "C" MlirModule ConvertLLVMToMLIR(LLVMModuleRef lmod, MlirContext cctx) {
     return wrap(res);
 }
 
+#include "llvm/IRReader/IRReader.h"
+extern "C" MlirModule ConvertLLVMStrToMLIR(const char* lmod, MlirContext cctx) {
+    LLVMContext Context;
+    SMDiagnostic Err;
+    auto llvmModule = llvm::parseIR(llvm::MemoryBufferRef(lmod, "conversion"), Err, Context);
+    mlir::MLIRContext &context = *unwrap(cctx);
+    auto res = mlir::translateLLVMIRToModule(std::move(llvmModule), &context, /*emitExpensiveWarnings*/false, /*dropDICompositeElements*/false).release();
+    return wrap(res);
+}
+
 
 /* Note that this */
 extern "C" xla::PjRtLoadedExecutable* ClientCompile(PjRtClient * client, MlirModule cmod) {

deps/ReactantExtra/BUILD

@@ -450,6 +450,8 @@ cc_library(
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:TransformDialect",
         "@llvm-project//mlir:Transforms",
+        
+        "@llvm-project//llvm:IRReader",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:AArch64AsmParser",
         "@llvm-project//llvm:AArch64CodeGen",

ext/ReactantNNlibExt.jl

@@ -2,16 +2,10 @@ module ReactantNNlibExt
 
 using NNlib
 using GPUArraysCore: @allowscalar
-using Reactant:
-    Reactant,
-    Ops,
-    TracedRArray,
-    AnyTracedRArray,
-    materialize_traced_array,
-    MLIR,
-    TracedRNumber,
-    get_mlir_data,
-    set_mlir_data!
+using Reactant: Reactant, Ops, TracedRArray, AnyTracedRArray, MLIR, TracedRNumber
+
+using Reactant.TracedUtils: materialize_traced_array, get_mlir_data, set_mlir_data!
+
 using ReactantCore: @trace
 using LinearAlgebra: LinearAlgebra, triu
 
@@ -238,9 +232,9 @@ function NNlib.batched_mul!(
     if size(x, 3) != size(y, 3)
         B = max(size(x, 3), size(y, 3))
         if size(x, 3) == 1
-            x = Reactant.broadcast_to_size(x, (size(x, 1), size(x, 2), B))
+            x = Reactant.TracedUtils.broadcast_to_size(x, (size(x, 1), size(x, 2), B))
         elseif size(y, 3) == 1
-            y = Reactant.broadcast_to_size(y, (size(y, 1), size(y, 2), B))
+            y = Reactant.TracedUtils.broadcast_to_size(y, (size(y, 1), size(y, 2), B))
         end
     end
 
@@ -250,9 +244,9 @@ function NNlib.batched_mul!(
     if size(x, 1) != size(y, 1)
         B = max(size(x, 1), size(y, 1))
         if size(x, 1) == 1
-            x = Reactant.broadcast_to_size(x, (B, size(x, 2), size(x, 3)))
+            x = Reactant.TracedUtils.broadcast_to_size(x, (B, size(x, 2), size(x, 3)))
         elseif size(y, 1) == 1
-            y = Reactant.broadcast_to_size(y, (B, size(y, 2), size(y, 3)))
+            y = Reactant.TracedUtils.broadcast_to_size(y, (B, size(y, 2), size(y, 3)))
         end
     end
 
@@ -270,7 +264,7 @@ end
 function NNlib.pad_constant(
     x::AnyTracedRArray{T,N}, pad::NTuple{N,Tuple{Int,Int}}, value
 ) where {T,N}
-    value = Reactant.promote_to(TracedRNumber{T}, value)
+    value = Reactant.TracedUtils.promote_to(TracedRNumber{T}, value)
     low = [i[1] for i in pad]
     high = [i[2] for i in pad]
     interior = [0 for i in pad]
@@ -329,7 +323,8 @@ function NNlib.gather!(dst::TracedRArray, src::AnyTracedRArray, idxs::AbstractAr
     start_sizes = ntuple(i -> size(src, i), dims)
     results = map(CartesianIndices(idxs)) do k
         res = @allowscalar src[colons..., Tuple(idxs[k])...]
-        res isa TracedRNumber && (res = Reactant.broadcast_to_size(res, (1,)))
+        res isa TracedRNumber &&
+            (res = Reactant.TracedUtils.broadcast_to_size(res, (1,)))
         return reshape(res, start_sizes..., :)
     end
     res = reshape(cat(results...; dims=(dims + 1)), size(dst))

ext/ReactantStatisticsExt.jl

@@ -1,6 +1,7 @@
 module ReactantStatisticsExt
 
-using Reactant: AnyTracedRArray, materialize_traced_array
+using Reactant: AnyTracedRArray
+using Reactant.TracedUtils: materialize_traced_array
 using Statistics: Statistics
 
 function Statistics.mean(A::AnyTracedRArray{T,N}; dims=:) where {T,N}

ext/ReactantYaoBlocksExt.jl

@@ -1,12 +1,13 @@
 module ReactantYaoBlocksExt
 
 using Reactant
+using Reactant.TracedUtils: broadcast_to_size
 using YaoBlocks
 
 function YaoBlocks.mat(
     ::Type{T}, R::RotationGate{D,Reactant.TracedRNumber{S},<:XGate}
 ) where {D,T,S}
-    M = Reactant.broadcast_to_size(zero(T), (2, 2))
+    M = broadcast_to_size(zero(T), (2, 2))
     c = cos(R.theta / 2)
     s = -im * sin(R.theta / 2)
     M[1, 1] = c
@@ -19,7 +20,7 @@ end
 function YaoBlocks.mat(
     ::Type{T}, R::RotationGate{D,Reactant.TracedRNumber{S},<:YGate}
 ) where {D,T,S}
-    M = Reactant.broadcast_to_size(zero(T), (2, 2))
+    M = broadcast_to_size(zero(T), (2, 2))
     c = cos(R.theta / 2)
     s = sin(R.theta / 2)
     M[1, 1] = c
@@ -32,7 +33,7 @@ end
 function YaoBlocks.mat(
     ::Type{T}, R::RotationGate{D,Reactant.TracedRNumber{S},<:ZGate}
 ) where {D,T,S}
-    M = Reactant.broadcast_to_size(zero(T), (2, 2))
+    M = broadcast_to_size(zero(T), (2, 2))
     x = exp(im * R.theta / 2)
     M[1, 1] = conj(x)
     M[2, 2] = x

lib/ReactantCore/Project.toml

@@ -1,7 +1,7 @@
 name = "ReactantCore"
 uuid = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>"]
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43"

lib/ReactantCore/src/ReactantCore.jl

@@ -153,15 +153,17 @@ function trace_for(mod, expr)
 
     all_syms = Expr(:tuple, counter, external_syms...)
     args_init = Expr(
-        :tuple, :(Reactant.promote_to(Reactant.TracedRNumber{Int}, 0)), external_syms...
+        :tuple,
+        :(Reactant.TracedUtils.promote_to(Reactant.TracedRNumber{Int}, 0)),
+        external_syms...,
     )
 
     reactant_code_block = quote
         let args = $(args_init)
             cond_fn =
                 $(all_syms) -> begin
                     local num_iters = div($limit - $start, $step, RoundDown)
-                    local num_iters = Reactant.promote_to(
+                    local num_iters = Reactant.TracedUtils.promote_to(
                         Reactant.TracedRNumber{Int64}, num_iters
                     )
                     $counter < num_iters + 1

src/Compiler.jl

@@ -292,7 +292,7 @@ function compile_mlir!(mod, f, args; optimize::Union{Bool,Symbol}=true)
     func2, traced_result, result, seen_args, ret, linear_args, in_tys,
     linear_results = MLIR.IR.mmodule!(mod) do
         MLIR.IR.block!(MLIR.IR.body(mod)) do
-            return Reactant.make_mlir_fn(f, args, (), "main", true)
+            return Reactant.TracedUtils.make_mlir_fn(f, args, (), "main", true)
         end
     end
 
@@ -779,6 +779,13 @@ function compile(f, args; client=nothing, optimize=true, sync=false)
     return register_thunk(fname, body)
 end
 
+# Compiling within a compile should return simply the original function
+Reactant.@reactant_override function Reactant.Compiler.compile(
+    f, args; client=nothing, optimize=true, sync=false
+)
+    return f
+end
+
 # inspired by RuntimeGeneratedFunction.jl
 const __thunk_body_cache = Dict{Symbol,Expr}()
 

src/ConcreteRArray.jl

@@ -99,7 +99,7 @@ end
 function Base.convert(
     ::Type{T}, X::WrappedConcreteRArray{ElType,N}
 ) where {T<:Array,ElType,N}
-    fn = compile(materialize_traced_array, (X,))
+    fn = compile(TracedUtils.materialize_traced_array, (X,))
     return convert(Array, fn(X))
 end
 Base.Array(x::AnyConcreteRArray) = convert(Array, x)
@@ -345,3 +345,11 @@ end
 
 buffer_on_cpu(::Any) = true
 buffer_on_cpu(x::ConcreteRArray) = XLA.BufferOnCPU(x.data.buffer)
+
+function Ops.constant(x::ConcreteRArray; kwargs...)
+    return Ops.constant(Base.convert(Array, x); kwargs...)
+end
+
+function Ops.constant(x::ConcreteRNumber{T}; kwargs...) where {T}
+    return Ops.constant(Base.convert(T, x); kwargs...)
+end