Insert call instructions (without any caching) #1276

jumerckx · 2025-05-12T20:17:40Z

No description provided.

…ure the latter can be removed?

….fill` instead of creating an uninitialized TracedRArray

…urse through the object. This is useful for constructing a new traced object without altering the paths.

…thing. This now handles that.

It checks the contents of the generated MLIR which is different when call operations are generated.

github-actions

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

JuliaFormatter

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Lines 1105 to 1113 in 68e7079

    
           traced_result = push_inst!(Expr( 
        
               :call, 
        
               GlobalRef(Core, :_apply_iterate), 
        
               Base.iterate, 
        
               call_epilogue, 
        
               push_inst!(Expr(:call, GlobalRef(Core, :tuple), fn_args[1], push_inst!(Expr(:call, GlobalRef(Core, :tuple), fn_args[2:end]...)))), 
        
               push_inst!(Expr(:call, GlobalRef(Core, :tuple), traced_result)), 
        
               push_inst!(Expr(:call, GlobalRef(Reactant, :get_args_from_finalize_function), finalize_function_result)), 
        
           ))

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Lines 1117 to 1122 in 68e7079

    
           push_inst!(Expr( 
        
               :call, 
        
               GlobalRef(Base, :setindex!), 
        
               GlobalRef(Reactant, :TRACE_CALLS), 
        
               should_trace_call, 
        
           ))

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Lines 1125 to 1130 in 68e7079

    
           push_inst!(Expr( 
        
               :call, 
        
               GlobalRef(Base, :setindex!), 
        
               GlobalRef(Reactant, :TRACE_CALLS), 
        
               TRACE_CALLS[], 
        
           ))

… call operation for blacklisted operations but not for the functions called within them.

This reverts commit 7c732dc.

jumerckx · 2025-05-14T14:56:22Z

Status:

Test Summary:           | Pass  Fail  Error  Broken  Total      Time
Reactant.jl Tests       | 1480     5     27       1   1513  45m55.8s

(https://github.com/EnzymeAD/Reactant.jl/actions/runs/15022467542/job/42214855196?pr=1276#step:10:6661)

…oadcastFunction

…en calling `make_tracer` on the `AnyTracedRArray`.

wsmoses · 2025-05-21T15:30:08Z

src/TracedRArray.jl

@@ -675,14 +675,14 @@ function Base.similar(
    ::Broadcasted{AbstractReactantArrayStyle{N}}, ::Type{T}, dims
 ) where {T<:Reactant.ReactantPrimitive,N}
    @assert N isa Int
-    return TracedRArray{T,length(dims)}((), nothing, map(length, dims))
+    return Ops.fill(zero(unwrapped_eltype(T)), dims)


minor, but could we make a smaller PR with some of these utilities to get that in first (and confirm everything works)

…ctant.apply`

…t call traced if they have no traced argument.

These tests inspect the generated MLIR code, which is different with call tracing.

github-actions · 2025-06-25T14:51:25Z

test/control_flow.jl

 @testset "call: caching for Julia operands" begin
    y = rand(3)
    y_ra = Reactant.to_rarray(y)


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

@testset "call: caching for Julia operands" begin

y = rand(3)

y_ra = Reactant.to_rarray(y)

@testset "call: caching for Julia operands" begin

y = rand(3)

y_ra = Reactant.to_rarray(y)

github-actions · 2025-06-25T14:51:27Z

test/control_flow.jl

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/control_flow.jl

Lines 721 to 725 in 5f91692

@testset "call: basic" begin

a = rand(2, 3)

b = rand(2, 3)

a_ra = Reactant.to_rarray(a)

b_ra = Reactant.to_rarray(b)

github-actions · 2025-06-25T14:51:28Z

test/control_flow.jl

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/control_flow.jl

Lines 738 to 742 in 5f91692

@test @jit(call1(a_ra, c_ra)) ≈ call1(a, c)

ir = @code_hlo optimize = false call1(a_ra, c_ra)

ops = [op for op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))]

@test length(ops) == 3

end

github-actions · 2025-06-25T14:51:29Z

test/control_flow.jl

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/control_flow.jl

Lines 776 to 779 in 5f91692

ir = @code_hlo optimize = false call3(y_ra)

ops = [op for op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))]

@test length(ops) == 5 # call3, .+, .*, _call3 (2X)

end

github-actions · 2025-06-25T14:51:30Z

test/control_flow.jl

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/control_flow.jl

Lines 797 to 806 in 5f91692

@testset "call: Caching struct arguments" begin

a = rand(10)

b = rand(10)

foo = Foo(Reactant.to_rarray(a))

foo2 = Foo(Reactant.to_rarray(b))

bar = Foo(Bar(Reactant.to_rarray(b))) # typeof(foo) == typeof(bar), but these don't match!

ir = @code_hlo optimize = false call4(foo, foo2, bar)

ops = [op for op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))]

@test length(ops) == 3 # call4, _call4 for {foo, foo2}, and _call4 for bar

end

jumerckx · 2025-06-25T15:39:58Z

lib/ReactantCore/src/ReactantCore.jl

                $(from_locals...)
+                Reactant.TRACE_CALLS[] = temp


This needs to be changed because we can't depend on Reactant here. I could add a macro @notrace_calls in ReactantCore?

jumerckx added 13 commits May 5, 2025 09:23

use args_in_result=:all instead of result_and_mutated. In the fut…

4dde0d3

…ure the latter can be removed?

fix implementation of Base.similar(::Broadcasted, ...) to call `Ops…

50f6d88

….fill` instead of creating an uninitialized TracedRArray

add make_tracer for Memory

a72a1a1

fix for TracedToType

454e280

allow path=nothing in make_tracer to keep path as is and just rec…

6a0fc38

…urse through the object. This is useful for constructing a new traced object without altering the paths.

make_mlir_fn fixes

885b09c

WIP automatic function call insertion

8d64d4d

add mutate_args

cc64c9e

TracedToTypes for TracedRArray should look at objectid.

caf0fe8

with automatic call tracing, it turns out that tobatch can be !isno…

fb60071

…thing. This now handles that.

automatic function call insertion but don't cache at all

f14c755

disable const dedup test for now.

3990b0e

It checks the contents of the generated MLIR which is different when call operations are generated.

Merge remote-tracking branch 'origin/main' into jm/make_mlir_fn

68e7079

github-actions bot reviewed May 12, 2025

View reviewed changes

jumerckx added 5 commits May 12, 2025 22:56

fix?

264d3dd

disable 1.10 CI until I fix it.

7c732dc

fix seen_args error

0eabdc7

change should_trace_call to trace_call_within to still generate a…

886b16c

… call operation for blacklisted operations but not for the functions called within them.

add try-finally in generated code

31a55b8

jumerckx force-pushed the jm/make_mlir_fn branch from 88a0db9 to 31a55b8 Compare May 13, 2025 12:53

jumerckx added 4 commits May 13, 2025 15:06

Merge branch 'main' into jm/make_mlir_fn

1735b43

set_mlir_data! for MissingTracedValue

0ad0017

Revert "disable 1.10 CI until I fix it."

02ca88f

This reverts commit 7c732dc.

fallback for set_mlir_data!

efee7ff

jumerckx force-pushed the jm/make_mlir_fn branch from c63dec4 to efee7ff Compare May 14, 2025 13:49

jumerckx added 4 commits May 17, 2025 16:25

don't trace through functions that capture values (sizeof != 0) or Br…

364c9b9

…oadcastFunction

also take into acount resarg

12b689e

fix

5fd1b4c

remove Dict typing in LinearAlgebra._diagm as this gives trouble wh…

899bf9f

…en calling `make_tracer` on the `AnyTracedRArray`.

jumerckx force-pushed the jm/make_mlir_fn branch from 810a76b to 899bf9f Compare May 17, 2025 20:22

jumerckx added 3 commits May 21, 2025 16:20

formatting

4e0af6f

Merge branch 'main' into jm/make_mlir_fn

9f15feb

guard make_tracer for Memory for v1.11

aeb068c

jumerckx marked this pull request as ready for review May 21, 2025 15:00

wsmoses reviewed May 21, 2025

View reviewed changes

jumerckx added 8 commits June 15, 2025 22:34

Merge branch 'main' into jm/make_mlir_fn

c680db1

bail is_traced when a field is undefined

0905b42

don't trace calls that have no traced arguments

c0448bb

insert function name in verify_arg_names when call is wrapped in `Rea…

9fd2246

…ctant.apply`

locally disable call tracing in Base._cat_t(..., ::TracedRArray)

02fe9b1

disable call tracing in from_locals expression in trace_while

6270b53

don't blacklist methods from Reactant modules now that methods are no…

49c1345

…t call traced if they have no traced argument.

disable for some control flow tests when call tracing is enabled.

5f91692

These tests inspect the generated MLIR code, which is different with call tracing.

github-actions bot reviewed Jun 25, 2025

View reviewed changes

jumerckx commented Jun 25, 2025

View reviewed changes

jumerckx added 3 commits July 16, 2025 09:37

Merge remote-tracking branch 'origin/main' into jm/make_mlir_fn

1936393

fix

6b15ca2

formatting

c288fce

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Insert call instructions (without any caching) #1276

Insert call instructions (without any caching) #1276

Uh oh!

jumerckx commented May 12, 2025

Uh oh!

github-actions bot left a comment

Uh oh!

jumerckx commented May 14, 2025 •

edited

Loading

Uh oh!

wsmoses May 21, 2025

Uh oh!

github-actions bot Jun 25, 2025

Uh oh!

github-actions bot Jun 25, 2025

Uh oh!

github-actions bot Jun 25, 2025

Uh oh!

github-actions bot Jun 25, 2025

Uh oh!

github-actions bot Jun 25, 2025

Uh oh!

jumerckx Jun 25, 2025

Uh oh!

Uh oh!

	traced_result = push_inst!(Expr(
	:call,
	GlobalRef(Core, :_apply_iterate),
	Base.iterate,
	call_epilogue,
	push_inst!(Expr(:call, GlobalRef(Core, :tuple), fn_args[1], push_inst!(Expr(:call, GlobalRef(Core, :tuple), fn_args[2:end]...)))),
	push_inst!(Expr(:call, GlobalRef(Core, :tuple), traced_result)),
	push_inst!(Expr(:call, GlobalRef(Reactant, :get_args_from_finalize_function), finalize_function_result)),
	))

	push_inst!(Expr(
	:call,
	GlobalRef(Base, :setindex!),
	GlobalRef(Reactant, :TRACE_CALLS),
	should_trace_call,
	))

	@testset "call: basic" begin
	a = rand(2, 3)
	b = rand(2, 3)
	a_ra = Reactant.to_rarray(a)
	b_ra = Reactant.to_rarray(b)

	@test @jit(call1(a_ra, c_ra)) ≈ call1(a, c)
	ir = @code_hlo optimize = false call1(a_ra, c_ra)
	ops = [op for op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))]
	@test length(ops) == 3
	end

	ir = @code_hlo optimize = false call3(y_ra)
	ops = [op for op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))]
	@test length(ops) == 5 # call3, .+, .*, _call3 (2X)
	end

	@testset "call: Caching struct arguments" begin
	a = rand(10)
	b = rand(10)
	foo = Foo(Reactant.to_rarray(a))
	foo2 = Foo(Reactant.to_rarray(b))
	bar = Foo(Bar(Reactant.to_rarray(b))) # typeof(foo) == typeof(bar), but these don't match!
	ir = @code_hlo optimize = false call4(foo, foo2, bar)
	ops = [op for op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))]
	@test length(ops) == 3 # call4, _call4 for {foo, foo2}, and _call4 for bar
	end

Insert call instructions (without any caching) #1276

Are you sure you want to change the base?

Insert call instructions (without any caching) #1276

Uh oh!

Conversation

jumerckx commented May 12, 2025

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

jumerckx commented May 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wsmoses May 21, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

jumerckx Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jumerckx commented May 14, 2025 •

edited

Loading