Merge branch 'main' into branch-3

Author: wsmoses

.buildkite/pipeline.yml

@@ -1,7 +1,7 @@
 steps:
   - group: ":test_tube: Tests"
     steps:
-      - label: "CUDA Julia v{{matrix.version}} -- {{matrix.group}}"
+      - label: ":julia: :linux: CUDA Julia v{{matrix.version}} -- {{matrix.group}}"
         matrix:
           setup:
             version:
@@ -33,7 +33,7 @@ steps:
         env:
           REACTANT_TEST_GROUP: "{{matrix.group}}"
         if: build.message !~ /\[skip tests\]/
-        timeout_in_minutes: 60
+        timeout_in_minutes: 120
 
       - label: ":julia: :linux: aarch64 - Julia v{{matrix.version}} -- {{matrix.group}}"
         matrix:
@@ -70,78 +70,78 @@ steps:
         env:
           REACTANT_TEST_GROUP: "{{matrix.group}}"
         if: build.message !~ /\[skip tests\]/
-        timeout_in_minutes: 60
+        timeout_in_minutes: 120
 
-  - group: ":racehorse: Benchmarks"
-    steps:
-      - label: "CPU: Run Benchmarks"
-        plugins:
-          - JuliaCI/julia#v1:
-              version: "1"
-        command: |
-          julia --project=benchmark -e 'println("--- :julia: Instantiating project")
-              using Pkg
-              Pkg.develop([PackageSpec(path=pwd()), PackageSpec(path="lib/ReactantCore")])'
+  # - group: ":racehorse: Benchmarks"
+  #   steps:
+  #     - label: "CPU: Run Benchmarks"
+  #       plugins:
+  #         - JuliaCI/julia#v1:
+  #             version: "1"
+  #       command: |
+  #         julia --project=benchmark -e 'println("--- :julia: Instantiating project")
+  #             using Pkg
+  #             Pkg.develop([PackageSpec(path=pwd()), PackageSpec(path="lib/ReactantCore")])'
 
-          julia --project=benchmark -e 'println("--- :julia: Run Benchmarks")
-              include("benchmark/runbenchmarks.jl")'
-        artifact_paths:
-          - "benchmark/results/*"
-        agents:
-          # Models are quite large so we need a decent sized machine. Don't tell Chris we
-          # are stealing SciMLBenchmarks machine :P
-          queue: "juliaecosystem"
-          sandbox_capable: true
-          exclusive: true
-          arch: "x86_64"
-        env:
-          BENCHMARK_GROUP: CPU
-          JULIA_NUM_THREADS: "auto"
-        timeout_in_minutes: 120
+  #         julia --project=benchmark -e 'println("--- :julia: Run Benchmarks")
+  #             include("benchmark/runbenchmarks.jl")'
+  #       artifact_paths:
+  #         - "benchmark/results/*"
+  #       agents:
+  #         # Models are quite large so we need a decent sized machine. Don't tell Chris we
+  #         # are stealing SciMLBenchmarks machine :P
+  #         queue: "juliaecosystem"
+  #         sandbox_capable: true
+  #         exclusive: true
+  #         arch: "x86_64"
+  #       env:
+  #         BENCHMARK_GROUP: CPU
+  #         JULIA_NUM_THREADS: "auto"
+  #       timeout_in_minutes: 120
 
-      - label: "CUDA: Run Benchmarks"
-        plugins:
-          - JuliaCI/julia#v1:
-              version: "1"
-        command: |
-          julia --project=benchmark -e 'println("--- :julia: Instantiating project")
-              using Pkg
-              Pkg.develop([PackageSpec(path=pwd()), PackageSpec(path="lib/ReactantCore")])'
+  #     - label: "CUDA: Run Benchmarks"
+  #       plugins:
+  #         - JuliaCI/julia#v1:
+  #             version: "1"
+  #       command: |
+  #         julia --project=benchmark -e 'println("--- :julia: Instantiating project")
+  #             using Pkg
+  #             Pkg.develop([PackageSpec(path=pwd()), PackageSpec(path="lib/ReactantCore")])'
 
-          julia --project=benchmark -e 'println("--- :julia: Run Benchmarks")
-              include("benchmark/runbenchmarks.jl")'
-        artifact_paths:
-          - "benchmark/results/*"
-        agents:
-          queue: "benchmark"
-          gpu: "rtx4070"
-          cuda: "*"
-        env:
-          BENCHMARK_GROUP: CUDA
-          JULIA_NUM_THREADS: "auto"
-        timeout_in_minutes: 120
+  #         julia --project=benchmark -e 'println("--- :julia: Run Benchmarks")
+  #             include("benchmark/runbenchmarks.jl")'
+  #       artifact_paths:
+  #         - "benchmark/results/*"
+  #       agents:
+  #         queue: "benchmark"
+  #         gpu: "rtx4070"
+  #         cuda: "*"
+  #       env:
+  #         BENCHMARK_GROUP: CUDA
+  #         JULIA_NUM_THREADS: "auto"
+  #       timeout_in_minutes: 120
 
-      - wait: ~
-        continue_on_failure: true
+  #     - wait: ~
+  #       continue_on_failure: true
 
-      - label: "Combine benchmarks"
-        plugins:
-          - JuliaCI/julia#v1:
-              version: "1"
-        command: |
-          buildkite-agent artifact download "benchmark/results/*" .
+  #     - label: "Combine benchmarks"
+  #       plugins:
+  #         - JuliaCI/julia#v1:
+  #             version: "1"
+  #       command: |
+  #         buildkite-agent artifact download "benchmark/results/*" .
 
-          julia -e 'println("--- :julia: Instantiating project")
-              using Pkg
-              Pkg.add("BenchmarkTools")
+  #         julia -e 'println("--- :julia: Instantiating project")
+  #             using Pkg
+  #             Pkg.add("BenchmarkTools")
 
-              println("--- :julia: Combining Benchmarks")
-              include("benchmark/aggregate.jl")'
-        artifact_paths:
-          - "benchmark/results/combinedbenchmarks.json"
-        agents:
-          queue: "juliagpu"
-        timeout_in_minutes: 10
+  #             println("--- :julia: Combining Benchmarks")
+  #             include("benchmark/aggregate.jl")'
+  #       artifact_paths:
+  #         - "benchmark/results/combinedbenchmarks.json"
+  #       agents:
+  #         queue: "juliagpu"
+  #       timeout_in_minutes: 10
 
   # - label: "AMDGPU Julia v{{matrix.version}}"
   #   matrix:

.github/workflows/CI.yml

@@ -23,7 +23,7 @@ jobs:
         version:
           - '1.10'
           - '1.11'
-          - 'nightly'
+          # - 'nightly'
         os:
           - ubuntu-20.04
           - macOS-latest
@@ -51,20 +51,26 @@ jobs:
             assertions: true
             test_group: neural_networks
           - os: ubuntu-20.04
-            arch: x86
-            libReactant: packaged
-            version: '1.10'
-            test_group: core
-          - os: ubuntu-20.04
-            arch: x86
-            libReactant: packaged
-            version: '1.10'
-            test_group: neural_networks
-          - os: ubuntu-20.04
-            arch: x86
+            arch: x64
             libReactant: packaged
             version: '1.10'
+            assertions: true
             test_group: integration
+          # - os: ubuntu-20.04
+          #   arch: x86
+          #   libReactant: packaged
+          #   version: '1.10'
+          #   test_group: core
+          # - os: ubuntu-20.04
+          #   arch: x86
+          #   libReactant: packaged
+          #   version: '1.10'
+          #   test_group: neural_networks
+          # - os: ubuntu-20.04
+          #   arch: x86
+          #   libReactant: packaged
+          #   version: '1.10'
+          #   test_group: integration
         exclude:
           # these are run on Buildkite
           - os: ubuntu-20.04

Project.toml

@@ -1,7 +1,7 @@
 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>", "Avik Pal <avikpal@mit.edu>"]
-version = "0.2.10"
+version = "0.2.11"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -15,32 +15,38 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 
 [weakdeps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
 
-[sources.ReactantCore]
-path = "lib/ReactantCore"
+[sources]
+ReactantCore = {path = "lib/ReactantCore"}
 
 [extensions]
 ReactantAbstractFFTsExt = "AbstractFFTs"
 ReactantArrayInterfaceExt = "ArrayInterface"
+ReactantCUDAExt = "CUDA"
 ReactantNNlibExt = "NNlib"
+ReactantRandom123Ext = "Random123"
 ReactantStatisticsExt = "Statistics"
 ReactantYaoBlocksExt = "YaoBlocks"
 
 [compat]
 AbstractFFTs = "1.5"
-Adapt = "4"
-ArrayInterface = "7.10"
-CEnum = "0.4, 0.5"
+Adapt = "4.1"
+ArrayInterface = "7.17.1"
+CEnum = "0.5"
+CUDA = "5.5"
 Downloads = "1.6"
 Enzyme = "0.13.22"
 EnzymeCore = "0.8.8"
@@ -50,8 +56,10 @@ NNlib = "0.9.26"
 OrderedCollections = "1"
 PrecompileTools = "1"
 Preferences = "1.4"
+Random = "1.10"
+Random123 = "1.7"
 ReactantCore = "0.1.3"
-Reactant_jll = "0.0.26"
+Reactant_jll = "0.0.27"
 Scratch = "1.2"
 Statistics = "1.10"
 YaoBlocks = "0.13"
@@ -60,4 +68,5 @@ julia = "1.10"
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"

deps/ReactantExtra/API.cpp

@@ -470,6 +470,10 @@ extern "C" void RegisterDialects(MlirContext cctx) {
   context.loadDialect<mlir::stablehlo::StablehloDialect>();
   context.loadDialect<mlir::chlo::ChloDialect>();
 }
+
+#include "mlir/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.h"
+#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
 extern "C" void InitializeRegistryAndPasses(MlirDialectRegistry creg) {
   mlir::DialectRegistry &registry = *unwrap(creg);
 
@@ -513,6 +517,11 @@ extern "C" void InitializeRegistryAndPasses(MlirDialectRegistry creg) {
   mlir::affine::registerAffinePasses();
   mlir::registerReconcileUnrealizedCasts();
 
+  mlir::registerLLVMDialectImport(registry);
+  mlir::registerNVVMDialectImport(registry);
+
+  mlir::LLVM::registerInlinerInterface(registry);
+
 /*
   registry.addExtension(+[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
     LLVM::LLVMFunctionType::attachInterface<MemRefInsider>(*ctx);
@@ -540,6 +549,81 @@ extern "C" void InitializeRegistryAndPasses(MlirDialectRegistry creg) {
   mlir::enzyme::registerEnzymeJaxTransformExtension(registry);
 }
 
+
+/// Returns an unused symbol in `module` for `oldSymbolName` by trying numeric
+/// suffix in `lastUsedID`.
+static mlir::StringAttr renameSymbol(llvm::StringRef oldSymName,
+                                     unsigned &lastUsedID,
+                                     mlir::ModuleOp source,
+                                     mlir::ModuleOp target) {
+  using namespace llvm;
+  using namespace mlir;
+  SmallString<64> newSymName(oldSymName);
+  newSymName.push_back('_');
+  while (true) {
+    auto possible = newSymName + Twine(++lastUsedID);
+    if (!SymbolTable::lookupSymbolIn(source, possible.str()) && !SymbolTable::lookupSymbolIn(target, possible.str())) {
+      return StringAttr::get(target.getContext(), possible);
+    }
+  }
+}
+
+
+/// Checks if a symbol with the same name as `op` already exists in `source`.
+/// If so, renames `op` and updates all its references in `target`.
+static mlir::LogicalResult
+updateSymbolAndAllUses(mlir::SymbolOpInterface op, mlir::ModuleOp source, mlir::ModuleOp target,
+                       unsigned &lastUsedID) {
+  using namespace llvm;
+  using namespace mlir;
+
+  auto opName = op.getName().str();
+
+  if (!SymbolTable::lookupSymbolIn(target, opName)) {
+    return success();
+  }
+
+  StringAttr newSymName =
+      renameSymbol(opName, lastUsedID, source, target);
+
+  if (failed(SymbolTable::replaceAllSymbolUses(op, newSymName, source)))
+    return op.emitError("unable to update all symbol uses for ")
+           << opName << " to " << newSymName;
+
+  SymbolTable::setSymbolName(op, newSymName);
+  return success();
+}
+
+extern "C" MlirOperation LinkInModule(MlirModule prevModC, MlirModule newModC, const char* entryfn) {
+    auto prevMod = cast<ModuleOp>(*unwrap(prevModC));
+    auto newMod = cast<ModuleOp>(*unwrap(newModC));
+
+    Operation* entryFn = nullptr;
+
+      unsigned lastUsedID = 0;
+
+      for (auto &op : *newMod.getBody()) {
+        auto symbolOp = dyn_cast<SymbolOpInterface>(op);
+        if (!symbolOp)
+          continue;
+
+        StringRef oldSymName = symbolOp.getName();
+        
+        if (oldSymName == entryfn) {
+          entryFn = &op;
+        }
+
+        if (failed(updateSymbolAndAllUses(symbolOp, newMod, prevMod,
+                                          lastUsedID))) {
+          assert(0 && "failed to update all uses");
+        }
+        SymbolTable::setSymbolVisibility(&op, SymbolTable::Visibility::Private);
+      }
+      prevMod.getBody()->getOperations().splice(prevMod.getBody()->getOperations().end(),
+                                   newMod.getBody()->getOperations());
+  return wrap(entryFn);
+}
+
 #pragma region xla::ifrt
 
 #pragma region xla::ifrt::Value