Start forward mode AD

.github/workflows/CI.yml

@@ -23,31 +23,31 @@ jobs:
         test_group: [
           'quality',
           'basic',
-          'rrules/avoiding_non_differentiable_code',
-          'rrules/blas',
-          'rrules/builtins',
-          'rrules/fastmath',
-          'rrules/foreigncall',
-          'rrules/functionwrappers',
-          'rrules/iddict',
-          'rrules/lapack',
-          'rrules/linear_algebra',
-          'rrules/low_level_maths',
-          'rrules/memory',
-          'rrules/misc',
-          'rrules/new',
-          'rrules/tasks',
-          'rrules/twice_precision',
+          # 'rrules/avoiding_non_differentiable_code',
+          # 'rrules/blas',
+          # 'rrules/builtins',
+          # 'rrules/fastmath',
+          # 'rrules/foreigncall',
+          # 'rrules/functionwrappers',
+          # 'rrules/iddict',
+          # 'rrules/lapack',
+          # 'rrules/linear_algebra',
+          # 'rrules/low_level_maths',
+          # 'rrules/memory',
+          # 'rrules/misc',
+          # 'rrules/new',
+          # 'rrules/tasks',
+          # 'rrules/twice_precision',
         ]
         version:
-          - 'lts'
+          # - 'lts'
           - '1'
         arch:
           - x64
-        include:
-          - test_group: 'basic'
-            version: '1.10'
-            arch: x86
+        # include:
+        #   - test_group: 'basic'
+        #     version: '1.10'
+        #     arch: x86
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -66,132 +66,4 @@ jobs:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: false
-  extra:
-    name: ${{matrix.test_group.test_type}}-${{ matrix.test_group.label }}-${{ matrix.version }}-${{ matrix.arch }}
-    runs-on: ubuntu-latest
-    if: github.event_name != 'schedule'
-    strategy:
-      fail-fast: false
-      matrix:
-        test_group: [
-          {test_type: 'ext', label: 'differentiation_interface'},
-          {test_type: 'ext', label: 'dynamic_ppl'},
-          {test_type: 'ext', label: 'luxlib'},
-          {test_type: 'ext', label: 'nnlib'},
-          {test_type: 'ext', label: 'special_functions'},
-          {test_type: 'integration_testing', label: 'array'},
-          {test_type: 'integration_testing', label: 'bijectors'},
-          {test_type: 'integration_testing', label: 'diff_tests'},
-          {test_type: 'integration_testing', label: 'distributions'},
-          {test_type: 'integration_testing', label: 'gp'},
-          {test_type: 'integration_testing', label: 'logexpfunctions'},
-          {test_type: 'integration_testing', label: 'lux'},
-          {test_type: 'integration_testing', label: 'battery_tests'},
-          {test_type: 'integration_testing', label: 'misc_abstract_array'},
-          {test_type: 'integration_testing', label: 'temporalgps'},
-          {test_type: 'integration_testing', label: 'turing'},
-        ]
-        version:
-          - '1'
-          - 'lts'
-        arch:
-          - x64
-    steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
-          include-all-prereleases: false
-      - uses: julia-actions/cache@v2
-      - uses: julia-actions/julia-buildpkg@v1
-      - run: |
-          if [ ${{ matrix.test_group.test_type }} == 'ext' ]; then
-            julia --code-coverage=user --eval 'include("test/run_extra.jl")'
-          else
-            julia --eval 'include("test/run_extra.jl")'
-          fi
-        env:
-          LABEL: ${{ matrix.test_group.label }}
-          TEST_TYPE: ${{ matrix.test_group.test_type }}
-      - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v5
-        with:
-          files: lcov.info
-          token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: false
-  perf:
-    name: "Performance (${{ matrix.perf_group }})"
-    runs-on: ubuntu-latest
-    strategy:
-      fail-fast: false
-      matrix:
-        perf_group:
-          - 'hand_written'
-          - 'derived'
-    steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: '1'
-          arch: x64
-          include-all-prereleases: false
-      - uses: julia-actions/cache@v2
-      - uses: julia-actions/julia-buildpkg@v1
-      - run: julia --project=bench --eval 'include("bench/run_benchmarks.jl"); main()'
-        env:
-          PERF_GROUP: ${{ matrix.perf_group }}
-        shell: bash
-  compperf:
-    name: "Performance (inter-AD)"
-    runs-on: ubuntu-latest
-    if: github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name == github.repository
-    strategy:
-      fail-fast: false
-    steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: '1'
-          arch: x64
-          include-all-prereleases: false
-      - uses: julia-actions/cache@v2
-      - uses: julia-actions/julia-buildpkg@v1
-      - run: mkdir bench_results
-      - run: julia --project=bench --eval 'include("bench/run_benchmarks.jl"); main()'
-        env:
-          PERF_GROUP: 'comparison'
-          GKSwstype: '100'
-        shell: bash
-      - uses: actions/upload-artifact@v4
-        with:
-          name: benchmarking-results
-          path: bench_results/
-      # Useful code for testing action.
-      # - run: |
-      #     text="this is line one
-      #     this is line two
-      #     this is line three"
-      #     echo "$text" > benchmark_results.txt
-      - name: Read file content
-        id: read-file
-        run: |
-          {
-            echo "table<<EOF"
-            cat bench/benchmark_results.txt
-            echo "EOF"
-          } >> $GITHUB_OUTPUT
-      - name: Find Comment
-        uses: peter-evans/find-comment@v3
-        id: fc
-        with:
-          issue-number: ${{ github.event.pull_request.number }}
-          comment-author: github-actions[bot]
-      - id: post-report-as-pr-comment
-        name: Post Report as Pull Request Comment
-        uses: peter-evans/create-or-update-comment@v4
-        with:
-          issue-number: ${{ github.event.pull_request.number }}
-          body: "Performance Ratio:\nRatio of time to compute gradient and time to compute function.\nWarning: results are very approximate! See [here](https://github.com/compintell/Mooncake.jl/tree/main/bench#inter-framework-benchmarking) for more context.\n```\n${{ steps.read-file.outputs.table }}\n```"
-          comment-id: ${{ steps.fc.outputs.comment-id }}
-          edit-mode: replace
+

.gitignore

@@ -6,3 +6,4 @@ profile.pb.gz
 scratch.jl
 docs/build/
 docs/site/
+playground.jl

src/Mooncake.jl

@@ -27,6 +27,7 @@ using Base:
     twiceprecision
 using Base.Experimental: @opaque
 using Base.Iterators: product
+using Base.Meta: isexpr
 using Core:
     Intrinsics,
     bitcast,
@@ -50,6 +51,13 @@ using FunctionWrappers: FunctionWrapper
 # Needs to be defined before various other things.
 function _foreigncall_ end
 
+"""
+    frule!!(f::Dual, x::Dual...)
+
+Performs AD in forward mode, possibly modifying the inputs, and returns a `Dual`.
+"""
+function frule!! end
+
 """
     rrule!!(f::CoDual, x::CoDual...)
 
@@ -75,8 +83,11 @@ pb!!(1.0)
 """
 function rrule!! end
 
+include("interpreter/diffractor_compiler_utils.jl")
+
 include("utils.jl")
 include("tangents.jl")
+include("dual.jl")
 include("fwds_rvs_data.jl")
 include("codual.jl")
 include("debug_mode.jl")
@@ -88,6 +99,7 @@ include(joinpath("interpreter", "ir_utils.jl"))
 include(joinpath("interpreter", "bbcode.jl"))
 include(joinpath("interpreter", "ir_normalisation.jl"))
 include(joinpath("interpreter", "zero_like_rdata.jl"))
+include(joinpath("interpreter", "s2s_forward_mode_ad.jl"))
 include(joinpath("interpreter", "s2s_reverse_mode_ad.jl"))
 
 include("tools_for_rules.jl")
@@ -133,9 +145,13 @@ export primal,
     _add_to_primal,
     _diff,
     _dot,
+    Dual,
+    zero_dual,
     zero_codual,
     codual_type,
+    frule!!,
     rrule!!,
+    build_frule,
     build_rrule,
     value_and_gradient!!,
     value_and_pullback!!,

src/debug_mode.jl

@@ -1,3 +1,4 @@
+DebugFRule(rule) = rule  # TODO: make it non-trivial
 
 """
     DebugPullback(pb, y, x)

src/dual.jl

@@ -0,0 +1,38 @@
+struct Dual{P,T}
+    primal::P
+    tangent::T
+end
+
+primal(x::Dual) = x.primal
+tangent(x::Dual) = x.tangent
+Base.copy(x::Dual) = Dual(copy(primal(x)), copy(tangent(x)))
+_copy(x::P) where {P<:Dual} = x
+
+zero_dual(x) = Dual(x, zero_tangent(x))
+randn_dual(rng::AbstractRNG, x) = Dual(x, randn_tangent(rng, x))
+
+function dual_type(::Type{P}) where {P}
+    P == DataType && return Dual
+    P isa Union && return Union{dual_type(P.a),dual_type(P.b)}
+    P <: UnionAll && return Dual # P is abstract, so we don't know its tangent type.
+    return isconcretetype(P) ? Dual{P,tangent_type(P)} : Dual
+end
+
+function dual_type(p::Type{Type{P}}) where {P}
+    return @isdefined(P) ? Dual{Type{P},NoTangent} : Dual{_typeof(p),NoTangent}
+end
+
+_primal(x) = x
+_primal(x::Dual) = primal(x)
+
+_dual(x) = zero_dual(x)
+_dual(x::Dual) = x
+
+"""
+    verify_dual_type(x::Dual)
+
+Check that the type of `tangent(x)` is the tangent type of the type of `primal(x)`.
+"""
+verify_dual_type(x::Dual) = tangent_type(typeof(primal(x))) == typeof(tangent(x))
+
+@inline uninit_dual(x::P) where {P} = Dual(x, uninit_tangent(x))

src/interpreter/bbcode.jl

@@ -889,8 +889,19 @@
 """
 inc_args(x::Expr) = Expr(x.head, map(__inc, x.args)...)
 inc_args(x::ReturnNode) = isdefined(x, :val) ? ReturnNode(__inc(x.val)) : x
+inc_args(x::GotoIfNot) = GotoIfNot(__inc(x.cond), x.dest)
 inc_args(x::IDGotoIfNot) = IDGotoIfNot(__inc(x.cond), x.dest)
 inc_args(x::IDGotoNode) = x
+inc_args(x::PiNode) = PiNode(__inc(x.val), x.typ)
+function inc_args(x::PhiNode)
+    new_values = Vector{Any}(undef, length(x.values))
+    for n in eachindex(x.values)
+        if isassigned(x.values, n)
+            new_values[n] = __inc(x.values[n])
+        end
+    end
+    return PhiNode(x.edges, new_values)
+end
 function inc_args(x::IDPhiNode)
     new_values = Vector{Any}(undef, length(x.values))
     for n in eachindex(x.values)