feat: High-Level IFRT dispatches (EnzymeAD#828)

* feat: env var to control runtime

* feat: skeleton impl of IFRT Array & Number types

* feat: working single device execution for IFRT

* fix: ambiguity

* fix: missing runtime

* fix: avoid using args...

* ci: run tests with IFRT runtime

* fix: use preferences + alias ConcreteRArray and concreteRNumber

* feat: more concreteifrt* dispatches

* fix: update tests and more coverage

* fix: python number type

* ci: temporarily disable precompilation

* fix: check for no null client

* fix: more sharding + precompile fix

* fix: copy data over for sharding

* test: fix new tests

* fix: add check

* ci: reduce localjll runners

* fix: IFRT un-sharded inputs with partially sharded inputs

* feat: non-divisible dims for IFRT

* test: remove specialized test

* feat: distributed workflow working

* ci: fix test call

* fix: try disabling precompilation in aarch64

* ci: downgrade testing for ifrt

* fix: remove convert dispatch

* fix: KA backend

* chore: test precompilation on aarch64

* ci: properly use preferences

* chore: run formatter

* ci: dont overwrite preferences

Author: avik-pal

.buildkite/pipeline.yml

@@ -1,7 +1,7 @@
 steps:
   - group: ":test_tube: Tests"
     steps:
-      - label: ":julia: :linux: CUDA Julia v{{matrix.version}} -- {{matrix.group}}"
+      - label: ":julia: :linux: CUDA Julia v{{matrix.version}} -- {{matrix.group}} -- {{matrix.runtime}}"
         matrix:
           setup:
             version:
@@ -10,6 +10,9 @@ steps:
               - core
               - neural_networks
               - integration
+            runtime:
+              - "PJRT"
+              - "IFRT"
         plugins:
           - JuliaCI/julia#v1:
               version: "{{matrix.version}}"
@@ -20,6 +23,13 @@ steps:
                 - ext
                 - lib/ReactantCore/src
         commands: |
+          touch LocalPreferences.toml
+
+          echo "[Reactant]" >> LocalPreferences.toml
+          echo "xla_runtime = \"{{matrix.runtime}}\"" >> LocalPreferences.toml
+
+          cat LocalPreferences.toml
+
           julia --project=. -e 'println("--- :julia: Instantiating project")
               using Pkg
               Pkg.develop([PackageSpec(path="lib/ReactantCore")])'
@@ -33,6 +43,7 @@ steps:
         env:
           REACTANT_TEST_GROUP: "{{matrix.group}}"
           CUDA_VISIBLE_DEVICES: 0
+          JULIA_DEBUG: "Reactant,Reactant_jll"
         if: build.message !~ /\[skip tests\]/
         timeout_in_minutes: 120
 

.github/workflows/CI-localjll.yml

@@ -62,6 +62,16 @@ jobs:
           julia --color=yes --project=deps -e 'using Pkg; Pkg.instantiate()'
           julia --color=yes --project=deps deps/build_local.jl
           cp LocalPreferences.toml test/
+      - name: "Setup Runtime Preferences"
+        run: |
+          import Pkg
+          Pkg.Registry.update()
+          Pkg.instantiate()
+          using Preferences
+          Preferences.set_preferences!("Reactant", "xla_runtime" => "PJRT"; force=true)
+        shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0}
+        env:
+          JULIA_PKG_PRECOMPILE_AUTO: 0
       - name: "Install Dependencies"
         run: |
           import Pkg
@@ -77,16 +87,36 @@ jobs:
         if: ${{ matrix.version == '1.10' }}
         env:
           JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
-      - name: "Run Tests"
+      - name: "Run Tests: PJRT"
+        run: |
+          import Pkg
+          Pkg.Registry.update()
+          Pkg.test(; coverage="user")
+        shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0}
+        env:
+          JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
+          XLA_FLAGS: "--xla_force_host_platform_device_count=8"
+          JULIA_DEBUG: "Reactant,Reactant_jll"
+      - name: "Setup Runtime Preferences"
+        run: |
+          import Pkg
+          Pkg.Registry.update()
+          Pkg.instantiate()
+          using Preferences
+          Preferences.set_preferences!("Reactant", "xla_runtime" => "IFRT"; force=true)
+        shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0}
+        env:
+          JULIA_PKG_PRECOMPILE_AUTO: 0
+      - name: "Run Tests: IFRT"
         run: |
           import Pkg
           Pkg.Registry.update()
           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
           XLA_FLAGS: "--xla_force_host_platform_device_count=8"
+          JULIA_DEBUG: "Reactant,Reactant_jll"
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v5
         with:

.github/workflows/CI.yml

@@ -31,7 +31,7 @@ concurrency:
 jobs:
   test:
     timeout-minutes: 90
-    name: Julia ${{ matrix.version }} - ${{ matrix.test_group }} - ${{ matrix.os }} - assertions=${{ matrix.assertions }} - ${{ github.event_name }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.test_group }} - ${{ matrix.os }} - ${{ matrix.runtime }} - assertions=${{ matrix.assertions }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -51,21 +51,27 @@ jobs:
           - core
           - neural_networks
           - integration
+        runtime:
+          - "PJRT"
+          - "IFRT"
         assertions:
           - false
         include:
           - os: ubuntu-24.04
             version: '1.10'
             assertions: true
             test_group: core
+            runtime: "PJRT"
           - os: ubuntu-24.04
             version: '1.10'
             assertions: true
             test_group: neural_networks
+            runtime: "PJRT"
           - os: ubuntu-24.04
             version: '1.10'
             assertions: true
             test_group: integration
+            runtime: "PJRT"
           # - os: ubuntu-24.04
           #   libReactant: packaged
           #   version: '1.10'
@@ -97,6 +103,14 @@ jobs:
           sed -i.bak 's/exit 2/exit 0/g' julia/deps/tools/jlchecksum
           make -C julia -j $(nproc) FORCE_ASSERTIONS=1 LLVM_ASSERTIONS=1 JULIA_PRECOMPILE=0
           echo $PWD/julia/usr/bin >> $GITHUB_PATH
+      - name: "Setup Runtime Preferences"
+        uses: "DamianReeves/write-file-action@master"
+        with:
+          path: "LocalPreferences.toml"
+          write-mode: "overwrite"
+          contents: |
+            [Reactant]
+            xla_runtime = "${{ matrix.runtime }}"
       - name: "Install Dependencies"
         run: |
           import Pkg
@@ -124,6 +138,7 @@ jobs:
           JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
           REACTANT_TEST_GROUP: ${{ matrix.test_group }}
           XLA_FLAGS: "--xla_force_host_platform_device_count=8"
+          JULIA_DEBUG: "Reactant,Reactant_jll"
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v5
         with:

.github/workflows/downgrade.yml

@@ -38,6 +38,9 @@ jobs:
           - core
           - neural_networks
           - integration
+        runtime:
+          - PJRT
+          - IFRT
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -47,6 +50,14 @@ jobs:
       - uses: julia-actions/julia-downgrade-compat@v1
         with:
           skip: "ReactantCore"
+      - name: "Setup Runtime Preferences"
+        uses: "DamianReeves/write-file-action@master"
+        with:
+          path: "LocalPreferences.toml"
+          write-mode: "overwrite"
+          contents: |
+            [Reactant]
+            xla_runtime = "${{ matrix.runtime }}"
       - name: "Install Dependencies and Run Tests"
         run: |
           import Pkg
@@ -64,6 +75,7 @@ jobs:
           JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
           REACTANT_TEST_GROUP: ${{ matrix.test_group }}
           XLA_FLAGS: "--xla_force_host_platform_device_count=8"
+          JULIA_DEBUG: "Reactant,Reactant_jll"
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v5
         with:

deps/ReactantExtra/API.cpp

@@ -31,9 +31,9 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Transform/Transforms/Passes.h"
 #include "mlir/InitAllPasses.h"
+#include "mlir/Parser/Parser.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Transforms/Passes.h"
-#include "mlir/Parser/Parser.h"
 #include "src/enzyme_ad/jax/Dialect/Dialect.h"
 #include "src/enzyme_ad/jax/Implementations/XLADerivatives.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
@@ -230,7 +230,7 @@ extern "C" MlirAttribute mlirComplexAttrDoubleGetChecked(MlirLocation loc,
 extern "C" MlirOperation mlirOperationParse(MlirContext ctx,
                                             MlirStringRef code) {
   ParserConfig config(unwrap(ctx));
-  OwningOpRef<Operation*> owning_op = parseSourceString(unwrap(code), config);
+  OwningOpRef<Operation *> owning_op = parseSourceString(unwrap(code), config);
   if (!owning_op)
     return MlirOperation{nullptr};
   return MlirOperation{owning_op.release()};
@@ -1964,6 +1964,33 @@ extern "C" void ifrt_sharding_to_device_list(
   }
 }
 
+extern "C" void ifrt_sharding_to_index_domains(
+    HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding,
+    int64_t *array_size_list, int32_t array_size_len,
+    int64_t *index_domain_origins, int64_t *index_domain_shapes) {
+  std::vector<int64_t> array_size(array_size_len);
+  for (int i = 0; i < array_size_len; i++) {
+    array_size[i] = array_size_list[i];
+  }
+  auto array_size_span = absl::MakeSpan(array_size);
+  auto array_shape = xla::ifrt::Shape(array_size_span);
+
+  std::vector<ifrt::IndexDomain> index_domains =
+      MyValueOrThrow(sharding->obj()->IndexDomains(array_shape));
+
+  for (int i = 0; i < index_domains.size(); i++) {
+    auto index_domain = index_domains[i];
+    absl::Span<const int64_t> origin = index_domain.origin().elements();
+    absl::Span<const int64_t> shape = index_domain.shape().dims();
+
+    for (int j = 0; j < origin.size(); j++) {
+      auto idx = i * origin.size() + j;
+      index_domain_origins[idx] = origin[j];
+      index_domain_shapes[idx] = shape[j];
+    }
+  }
+}
+
 #pragma endregion
 
 typedef ifrt::Future<> IfRtFutureType;

ext/ReactantArrayInterfaceExt.jl

@@ -4,8 +4,7 @@ using ArrayInterface: ArrayInterface
 using Reactant:
     Reactant,
     RArray,
-    ConcretePJRTArray,
-    ConcretePJRTNumber,
+    AbstractConcreteNumber,
     TracedRNumber,
     TracedRArray,
     AnyTracedRArray,
@@ -14,8 +13,9 @@ using Reactant:
 ArrayInterface.can_setindex(::Type{<:RArray}) = false
 ArrayInterface.fast_scalar_indexing(::Type{<:RArray}) = false
 
-for aType in
-    (AbstractArray{<:ConcretePJRTNumber}, AbstractArray{<:TracedRNumber}, AnyTracedRArray)
+for aType in (
+    AbstractArray{<:AbstractConcreteNumber}, AbstractArray{<:TracedRNumber}, AnyTracedRArray
+)
     @eval ArrayInterface.aos_to_soa(x::$aType) = Reactant.aos_to_soa(x)
 end
 

ext/ReactantCUDAExt.jl

@@ -1093,8 +1093,10 @@ end
 Base.@nospecializeinfer function Reactant.traced_type_inner(
     @nospecialize(A::Type{<:CuTracedArray}),
     seen,
-    mode::Reactant.TraceMode,
-    @nospecialize(track_numbers::Type)
+    @nospecialize(mode::Reactant.TraceMode),
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
 )
     return A
 end
@@ -1104,28 +1106,34 @@ Base.@nospecializeinfer function Reactant.traced_type_inner(
     seen,
     mode::Reactant.TraceMode,
     @nospecialize(track_numbers::Type),
-    @nospecialize(sharding)
+    @nospecialize(sharding),
+    @nospecialize(runtime)
 )
     T = eltype(A)
     N = ndims(A)
     if mode == Reactant.ArrayToConcrete && T <: Reactant.ReactantPrimitive
-        if !Reactant.Sharding.is_sharded(sharding)
-            return Reactant.ConcretePJRTArray{T,N,1,Reactant.Sharding.NoShardInfo}
-        else
-            error("TODO: implement sharding")
-        end
-    else
-        TT = Reactant.traced_type_inner(T, seen, mode, track_numbers, sharding)
-        if TT === T
-            return A
-        else
-            return Array{
-                Reactant.traced_type_inner(
-                    T, seen, mode, track_numbers, Base.getproperty(sharding, 1)
-                ),
+        if runtime isa Val{:PJRT}
+            return Reactant.ConcretePJRTArray{
+                T,
                 N,
+                Reactant.Sharding.ndevices(sharding),
+                Reactant.Sharding.shard_type(typeof(sharding), N),
+            }
+        elseif runtime isa Val{:IFRT}
+            return Reactant.ConcreteIFRTArray{
+                T,N,Reactant.Sharding.shard_type(typeof(sharding), N)
             }
         end
+        error("Unsupported runtime $runtime")
+    else
+        TT = Reactant.traced_type_inner(T, seen, mode, track_numbers, sharding, runtime)
+        TT === T && return A
+        return Array{
+            Reactant.traced_type_inner(
+                T, seen, mode, track_numbers, Base.getproperty(sharding, 1), runtime
+            ),
+            N,
+        }
     end
 end
 
@@ -1136,6 +1144,7 @@ function Reactant.make_tracer(
     mode;
     @nospecialize(track_numbers::Type = Union{}),
     @nospecialize(sharding = Reactant.Sharding.NoSharding()),
+    @nospecialize(runtime),
     kwargs...,
 )
     RT = Core.Typeof(prev)
@@ -1145,9 +1154,14 @@ function Reactant.make_tracer(
         return seen[prev]
     end
     if mode == Reactant.ArrayToConcrete && eltype(RT) <: Reactant.ReactantPrimitive
-        return seen[prev] = Reactant.ConcretePJRTArray(Array(prev); sharding)
+        if runtime isa Val{:PJRT}
+            return seen[prev] = Reactant.ConcretePJRTArray(Array(prev); sharding)
+        elseif runtime isa Val{:IFRT}
+            return seen[prev] = Reactant.ConcreteIFRTArray(Array(prev); sharding)
+        end
+        error("Unsupported runtime $runtime")
     end
-    TT = Reactant.traced_type(eltype(RT), Val(mode), track_numbers, sharding)
+    TT = Reactant.traced_type(eltype(RT), Val(mode), track_numbers, sharding, runtime)
     if TT === eltype(RT)
         return prev
     end
@@ -1164,6 +1178,7 @@ function Reactant.make_tracer(
                 mode;
                 track_numbers,
                 sharding=Base.getproperty(sharding, I),
+                runtime,
                 kwargs...,
             )
             if pv !== nv
@@ -1192,7 +1207,15 @@ end
 @static if !Sys.isapple()
     Reactant.PrecompileTools.@setup_workload begin
         Reactant.initialize_dialect()
-        client = Reactant.XLA.PJRT.CPUClient(; checkcount=false)
+
+        if Reactant.XLA.REACTANT_XLA_RUNTIME == "PJRT"
+            client = Reactant.XLA.PJRT.CPUClient(; checkcount=false)
+        elseif Reactant.XLA.REACTANT_XLA_RUNTIME == "IFRT"
+            client = Reactant.XLA.IFRT.CPUClient(; checkcount=false)
+        else
+            error("Unsupported runtime: $(Reactant.XLA.REACTANT_XLA_RUNTIME)")
+        end
+
         Reactant.PrecompileTools.@compile_workload begin
             @static if Reactant.precompilation_supported() && VERSION != v"1.11.3"
                 function square_kernel!(x)
@@ -1205,10 +1228,11 @@ end
                     CUDA.@cuda blocks = 1 threads = length(x) square_kernel!(x)
                     return nothing
                 end
-                y = Reactant.ConcretePJRTArray([2.0]; client)
+                y = Reactant.ConcreteRArray([2.0]; client)
                 Reactant.Compiler.compile_mlir(square!, (y,); optimize=false)
             end
         end
+
         Reactant.XLA.free_client(client)
         client.client = C_NULL
         Reactant.deinitialize_dialect()