From 3da42edaf8a495c8e8cd91e4ebb1759fd4d84ed8 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Wed, 3 Sep 2025 21:18:17 +0000 Subject: [PATCH 1/4] Add partial evaluation tutorial --- docs/src/tutorials/index.md | 1 + docs/src/tutorials/partial-evaluation.md | 70 ++++++++++++++++++++++++ 2 files changed, 71 insertions(+) create mode 100644 docs/src/tutorials/partial-evaluation.md diff --git a/docs/src/tutorials/index.md b/docs/src/tutorials/index.md index 39217a3e43..177ea7e89f 100644 --- a/docs/src/tutorials/index.md +++ b/docs/src/tutorials/index.md @@ -6,5 +6,6 @@ - [Control flow](@ref control-flow). - [Sharding](@ref sharding). - [Persistent Compilation Cache](@ref persistent_compile_cache). +- [Partial evaluation](@ref partial-evaluation). We are currently working on adding more tutorials to Reactant!! Please check back soon! diff --git a/docs/src/tutorials/partial-evaluation.md b/docs/src/tutorials/partial-evaluation.md new file mode 100644 index 0000000000..ceafde0f93 --- /dev/null +++ b/docs/src/tutorials/partial-evaluation.md @@ -0,0 +1,70 @@ +# [Partial Evaluation](@id partial-evaluation) + +When compiling functions with Reactant, the function arguments (and possible +closure fields) may contain non-Reactant values, so numbers and arrays that +are not of type `Reactant.AbstractConcreteNumber` or +`Reactant.AbstractConcreteArray`. + +The Reactant compiler may (but is not guaranteed to) treat these non-Reactant +values as constant and partially evaluate the function to be compiled based +on this. + +For example, the function + + +```@example partial_evaluation_tutorial +function add(a, b) + a + b +end +``` + +when compiled with two `ConcreteRNumber` arguments + +```@example partial_evaluation_tutorial +using Reactant + +x = ConcreteRNumber(3) +y = ConcreteRNumber(4) + +addxy = @compile add(x, y) + +addxy(x, y) +``` + +returns a result that depends on both arguments: + + +```@example partial_evaluation_tutorial +addxy(ConcreteRNumber(7), ConcreteRNumber(8)) +``` + +The StableHLO IR code generated here is: + +```@example partial_evaluation_tutorial +@code_hlo add(x, y) +``` + +and shows two variable inputs. + +However, if one of the arguments is a non-Reactant value, then the result + +```@example partial_evaluation_tutorial +addx4 = @compile add(x, 4) + +addx4(x, 4) +``` + +will only change based on the first argument, not the second (non-Reactant) +argument: + +```@example partial_evaluation_tutorial +addx4(ConcreteRNumber(7), 8) +``` + +The StableHLO code shows that the second argument has been replaced by a +constant during partial evaluation and is ignored during execution of the +compiled function: + +```@example partial_evaluation_tutorial +@code_hlo add(x, 4) +``` From b9e164728db6de3e22a45b53d4a93731447b2ab6 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Wed, 3 Sep 2025 21:44:11 +0000 Subject: [PATCH 2/4] Improve partial evaluation tutorial --- docs/src/tutorials/index.md | 2 +- docs/src/tutorials/partial-evaluation.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/src/tutorials/index.md b/docs/src/tutorials/index.md index 177ea7e89f..d780f451eb 100644 --- a/docs/src/tutorials/index.md +++ b/docs/src/tutorials/index.md @@ -3,9 +3,9 @@ - [Profiling](@ref profiling). - [Multi-Host Environments](@ref distributed). - [Local build of ReactantExtra](@ref local-build). +- [Partial evaluation](@ref partial-evaluation). - [Control flow](@ref control-flow). - [Sharding](@ref sharding). - [Persistent Compilation Cache](@ref persistent_compile_cache). -- [Partial evaluation](@ref partial-evaluation). We are currently working on adding more tutorials to Reactant!! Please check back soon! diff --git a/docs/src/tutorials/partial-evaluation.md b/docs/src/tutorials/partial-evaluation.md index ceafde0f93..4cc5943481 100644 --- a/docs/src/tutorials/partial-evaluation.md +++ b/docs/src/tutorials/partial-evaluation.md @@ -1,7 +1,7 @@ # [Partial Evaluation](@id partial-evaluation) When compiling functions with Reactant, the function arguments (and possible -closure fields) may contain non-Reactant values, so numbers and arrays that +closure fields) may contain non-Reactant values, i.e. numbers and arrays that are not of type `Reactant.AbstractConcreteNumber` or `Reactant.AbstractConcreteArray`. From 695ab01811d5c2ae421bdccd41b0219ec17da47e Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Thu, 4 Sep 2025 09:00:27 +0000 Subject: [PATCH 3/4] Reorder tutorials --- docs/src/tutorials/index.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/src/tutorials/index.md b/docs/src/tutorials/index.md index d780f451eb..ece4f0b3dd 100644 --- a/docs/src/tutorials/index.md +++ b/docs/src/tutorials/index.md @@ -1,11 +1,11 @@ # Tutorials -- [Profiling](@ref profiling). -- [Multi-Host Environments](@ref distributed). -- [Local build of ReactantExtra](@ref local-build). - [Partial evaluation](@ref partial-evaluation). - [Control flow](@ref control-flow). - [Sharding](@ref sharding). +- [Profiling](@ref profiling). +- [Multi-Host Environments](@ref distributed). +- [Local build of ReactantExtra](@ref local-build). - [Persistent Compilation Cache](@ref persistent_compile_cache). We are currently working on adding more tutorials to Reactant!! Please check back soon! From 3b5216cb034b4b1bebebd87c3f8ccf6992625759 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Thu, 4 Sep 2025 09:18:03 +0000 Subject: [PATCH 4/4] Use doctests in partial evaluation tutorial --- docs/src/tutorials/partial-evaluation.md | 76 +++++++++++++++++++----- 1 file changed, 61 insertions(+), 15 deletions(-) diff --git a/docs/src/tutorials/partial-evaluation.md b/docs/src/tutorials/partial-evaluation.md index 4cc5943481..6b4220388b 100644 --- a/docs/src/tutorials/partial-evaluation.md +++ b/docs/src/tutorials/partial-evaluation.md @@ -12,15 +12,20 @@ on this. For example, the function -```@example partial_evaluation_tutorial +```jldoctest partial_evaluation_tutorial +using Reactant function add(a, b) a + b -end +end; + +# output + +add (generic function with 1 method) ``` when compiled with two `ConcreteRNumber` arguments -```@example partial_evaluation_tutorial +```jldoctest partial_evaluation_tutorial; filter = r"I000.*" => s"" using Reactant x = ConcreteRNumber(3) @@ -29,42 +34,83 @@ y = ConcreteRNumber(4) addxy = @compile add(x, y) addxy(x, y) + +# output + +I0000 +I0000 +ConcretePJRTNumber{Int64, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(7) ``` -returns a result that depends on both arguments: +returns a result that depends on both arguments `x` and `y`: -```@example partial_evaluation_tutorial +```jldoctest partial_evaluation_tutorial addxy(ConcreteRNumber(7), ConcreteRNumber(8)) + +# output + +ConcretePJRTNumber{Int64, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(15) ``` The StableHLO IR code generated here is: -```@example partial_evaluation_tutorial +```jldoctest partial_evaluation_tutorial @code_hlo add(x, y) + +# output + +module @reactant_add attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} { + func.func @main(%arg0: tensor, %arg1: tensor) -> tensor { + %0 = stablehlo.add %arg0, %arg1 : tensor + return %0 : tensor + } +} ``` -and shows two variable inputs. +So at HLO-level, there a are two variable inputs `%arg0` and `%arg1`. -However, if one of the arguments is a non-Reactant value, then the result +However, if argument `y` has a non-Reactant value during compilation, (`4` in +this example) then the result when executing the compiled function -```@example partial_evaluation_tutorial +```jldoctest partial_evaluation_tutorial; filter = r"I000.*" => s"" addx4 = @compile add(x, 4) addx4(x, 4) + +# output + +I0000 +I0000 +ConcretePJRTNumber{Int64, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(7) ``` -will only change based on the first argument, not the second (non-Reactant) -argument: +will only change based on `x`, not on the non-Reactant argument `y`, we get +`7 + 4 == 11`, not `7 + 8 == 15`: -```@example partial_evaluation_tutorial +```jldoctest partial_evaluation_tutorial addx4(ConcreteRNumber(7), 8) + +# output + +ConcretePJRTNumber{Int64, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(11) ``` The StableHLO code shows that the second argument has been replaced by a -constant during partial evaluation and is ignored during execution of the -compiled function: +constant `%c` during partial evaluation. When the compiled function is +executed, the value of `y` is ignored - at HLO-level, there is only one +variable input `%arg0`: -```@example partial_evaluation_tutorial +```jldoctest partial_evaluation_tutorial @code_hlo add(x, 4) + +# output + +module @reactant_add attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} { + func.func @main(%arg0: tensor) -> tensor { + %c = stablehlo.constant dense<4> : tensor + %0 = stablehlo.add %arg0, %c : tensor + return %0 : tensor + } +} ```