feat: support conditionals (#185)

* feat: support conditionals

* chore: run formatter

* feat: run post-processing on the expression to avoid boxing

* feat: overload traced_if for Reactant types

* fix: incorrect region

* fix: use `stablehlo.return_`

* feat: use ExpressionsExplorer for parsing

* fix: replace args

* feat: compiling `if-else` blocks now work 🎉

* feat: compile scalars in branches

* fix: bug in unrolling the generated code

* feat: support elseif

* feat: introduce a `MissingTracedValue`

* feat: cleanup missing values from the IR

* chore: apply formatting suggestion

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

* fix: codegen for else statements

* refactor: move code to ReactantCore

* chore: add to authors list

* test: test if conditions

* fix: check for elseif instead of not if

* chore: apply suggestions from code review

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

* ci(buildkite): fix the workflow

* test: more test cases

* fix: nested if handling

* fix: partial fix to mutation issue

* fix: code generation

* docs: rename to type instability

* test: mark some tests as broken

---------

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

Author: avik-pal

.buildkite/pipeline.yml

@@ -10,13 +10,20 @@ steps:
     plugins:
       - JuliaCI/julia#v1:
           version: "{{matrix.version}}"
-      - JuliaCI/julia-test#v1:
-          test_args: "--gpu"
       - JuliaCI/julia-coverage#v1:
           codecov: true
           dirs:
             - src
             - ext
+            - lib/ReactantCore/src
+    commands: |
+      julia --project=. -e 'println("--- :julia: Instantiating project")
+          using Pkg
+          Pkg.develop([PackageSpec(path="lib/ReactantCore")])'
+
+      julia --project=. -e 'println("--- :julia: Run Tests")
+          using Pkg
+          Pkg.test(; coverage="user")'
     agents:
       queue: "juliagpu"
       cuda: "*"
@@ -34,7 +41,7 @@ steps:
         command: |
           julia --project=benchmark -e 'println("--- :julia: Instantiating project")
               using Pkg
-              Pkg.develop([PackageSpec(path=pwd())])'
+              Pkg.develop([PackageSpec(path=pwd()), PackageSpec(path="lib/ReactantCore")])'
 
           julia --project=benchmark -e 'println("--- :julia: Run Benchmarks")
               include("benchmark/runbenchmarks.jl")'
@@ -59,7 +66,7 @@ steps:
         command: |
           julia --project=benchmark -e 'println("--- :julia: Instantiating project")
               using Pkg
-              Pkg.develop([PackageSpec(path=pwd())])'
+              Pkg.develop([PackageSpec(path=pwd()), PackageSpec(path="lib/ReactantCore")])'
 
           julia --project=benchmark -e 'println("--- :julia: Run Benchmarks")
               include("benchmark/runbenchmarks.jl")'

.github/workflows/CI.yml

@@ -84,8 +84,19 @@ jobs:
           julia --color=yes --project=deps -e 'using Pkg; Pkg.instantiate()'
           SDKROOT=`xcrun --show-sdk-path` julia --color=yes --project=deps deps/build_local.jl
           cp LocalPreferences.toml test/
-      - uses: julia-actions/julia-runtest@v1
-        # if: steps.buildpkg.outcome == 'success'
+      - name: "Install Dependencies and Run Tests"
+        run: |
+          import Pkg
+          Pkg.Registry.update()
+          # Install packages present in subdirectories
+          dev_pks = Pkg.PackageSpec[]
+          for path in ("lib/ReactantCore",)
+              push!(dev_pks, Pkg.PackageSpec(; path))
+          end
+          Pkg.develop(dev_pks)
+          Pkg.instantiate()
+          Pkg.test(; coverage="user")
+        shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0}
         id: run_tests
         env:
           JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
@@ -108,7 +119,11 @@ jobs:
       - run: |
           julia --color=yes --project=docs -e '
             using Pkg
-            Pkg.develop([PackageSpec(path=pwd()), PackageSpec("Reactant_jll")])
+            Pkg.develop([
+              PackageSpec(path=pwd()),
+              PackageSpec("Reactant_jll"),
+              PackageSpec(path="lib/ReactantCore")
+            ])
             Pkg.instantiate()'
         env:
           JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager

.gitignore

@@ -81,7 +81,6 @@ dist/
 downloads/
 eggs/
 .eggs/
-lib/
 lib64/
 parts/
 sdist/

Project.toml

@@ -1,6 +1,6 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
-authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>", "Sergio Sánchez Ramírez <sergio.sanchez.ramirez@bsc.es>", "Paul Berg <paul@plutojl.org>"]
+authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>", "Sergio Sánchez Ramírez <sergio.sanchez.ramirez@bsc.es>", "Paul Berg <paul@plutojl.org>", "Avik Pal <avikpal@mit.edu>"]
 version = "0.2.3"
 
 [deps]
@@ -11,6 +11,7 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 
@@ -33,6 +34,7 @@ Enzyme = "0.13"
 NNlib = "0.9"
 OrderedCollections = "1"
 Preferences = "1.4"
+ReactantCore = "0.1"
 Reactant_jll = "0.0.22"
 Scratch = "1.2"
 Statistics = "1.10"

lib/ReactantCore/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Enzyme Automatic Differentiation Compiler
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lib/ReactantCore/Project.toml

@@ -0,0 +1,13 @@
+name = "ReactantCore"
+uuid = "a3311ec8-5e00-46d5-b541-4f83e724a433"
+authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>", "Sergio Sánchez Ramírez <sergio.sanchez.ramirez@bsc.es>", "Paul Berg <paul@plutojl.org>", "Avik Pal <avikpal@mit.edu>"]
+version = "0.1.0"
+
+[deps]
+ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+
+[compat]
+ExpressionExplorer = "1"
+MacroTools = "0.5.13"
+julia = "1.10"