Merge branch 'main' of https://github.com/EnzymeAD/Reactant.jl into probprog

Author: sbrantq

.github/workflows/CI-localjll.yml

@@ -47,7 +47,7 @@ jobs:
         with:
           version: ${{ matrix.version }}
       - uses: julia-actions/cache@v2
-      - uses: bazel-contrib/setup-bazel@0.14.0
+      - uses: bazel-contrib/setup-bazel@0.15.0
         name: Set up Bazel
         with:
           # Avoid downloading Bazel every time.

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.122"
+version = "0.2.134"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -52,7 +52,7 @@ ReactantArrayInterfaceExt = "ArrayInterface"
 ReactantCUDAExt = ["CUDA", "GPUCompiler", "KernelAbstractions", "LLVM"]
 ReactantKernelAbstractionsExt = "KernelAbstractions"
 ReactantMPIExt = "MPI"
-ReactantNNlibExt = "NNlib"
+ReactantNNlibExt = ["NNlib", "Statistics"]
 ReactantOffsetArraysExt = "OffsetArrays"
 ReactantOneHotArraysExt = "OneHotArrays"
 ReactantPythonCallExt = "PythonCall"
@@ -69,8 +69,8 @@ CEnum = "0.5"
 CUDA = "5.6"
 Downloads = "1.6"
 EnumX = "1"
-Enzyme = "0.13.46"
-EnzymeCore = "0.8.9"
+Enzyme = "0.13.49"
+EnzymeCore = "0.8.11"
 Functors = "0.5"
 GPUArraysCore = "0.2"
 GPUCompiler = "1.3"
@@ -89,8 +89,8 @@ Preferences = "1.4"
 PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
-ReactantCore = "0.1.11"
-Reactant_jll = "0.0.191"
+ReactantCore = "0.1.12"
+Reactant_jll = "0.0.201"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

deps/ReactantExtra/.bazelrc

@@ -23,7 +23,7 @@ build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.8.1"
 build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.8.0"
 # "sm" means we emit only cubin, which is forward compatible within a GPU generation.
 # "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations.
-build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,compute_90"
+build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_60,sm_70,sm_80,compute_90"
 build:cuda --crosstool_top="@local_config_cuda//crosstool:toolchain"
 build:cuda --@local_config_cuda//:enable_cuda
 # Default hermetic CUDA and CUDNN versions.

deps/ReactantExtra/API.cpp

@@ -39,11 +39,14 @@
 #include "src/enzyme_ad/jax/Dialect/Dialect.h"
 #include "src/enzyme_ad/jax/Implementations/XLADerivatives.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
+#include "src/enzyme_ad/jax/RegistryUtils.h"
 #include "llvm/Support/TargetSelect.h"
 
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
 #include "stablehlo/dialect/ChloOps.h"
 #include "stablehlo/dialect/StablehloOps.h"
+#include "stablehlo/transforms/Passes.h"
+#include "stablehlo/transforms/optimization/Passes.h"
 
 #include "absl/log/globals.h"
 #include "absl/log/initialize.h"
@@ -75,6 +78,7 @@
 #include "xla/pjrt/pjrt_api.h"
 #include "xla/pjrt/pjrt_c_api_client.h"
 #include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
 
 // CPU collectives
 #include "xla/backends/cpu/collectives/mpi_collectives.h"
@@ -147,10 +151,6 @@
 
 #include "jaxlib/mosaic/dialect/tpu/tpu_dialect.h"
 
-// Triton did a dumb thing and their import is incompatible
-// We don't use so disabling until upstream fix
-// #include "triton/Dialect/Triton/IR/Dialect.h"
-
 #include "llvm/Support/ExtensibleRTTI.h"
 #include <llvm/Support/FileSystem.h>
 #include <llvm/Support/raw_ostream.h>
@@ -164,10 +164,6 @@ void registerRemoveTransformPass();
 void registerGenerateApplyPatternsPass();
 } // namespace enzyme
 
-namespace triton {
-class TritonDialect;
-}
-
 } // namespace mlir
 
 namespace reactant {
@@ -213,6 +209,8 @@ using HeldPjRtBuffer = HeldValue<std::shared_ptr<xla::PjRtBuffer>>;
 using HeldIfrtArray = HeldValue<tsl::RCReference<xla::ifrt::Array>>;
 using HeldHloModule = HeldValue<std::shared_ptr<xla::HloModule>>;
 using HeldIfrtSharding = HeldValue<std::shared_ptr<xla::ifrt::Sharding>>;
+using HeldIfrtLoadedExecutable =
+    HeldValue<std::shared_ptr<xla::ifrt::LoadedExecutable>>;
 
 extern "C" void (*ReactantThrowError)(const char *) = nullptr;
 
@@ -409,8 +407,7 @@ PjRtClient *MakeCPUClientInternal(
   if (collectives.has_value())
     options.collectives = collectives.value();
 
-  auto client = MyValueOrThrow(GetTfrtCpuClient(options));
-  return client.release();
+  return MyValueOrThrow(GetPjRtCpuClient(options)).release();
 }
 
 extern "C" PjRtClient *MakeCPUClient(uint8_t asynchronous, int node_id) {
@@ -1118,25 +1115,15 @@ extern "C" int PjRtLoadedExecutableNumPartitions(PjRtLoadedExecutable *exec) {
   return exec->num_partitions();
 }
 
-void prepareRegistry(mlir::DialectRegistry &registry);
-
 extern "C" void RegisterDialects(MlirContext cctx) {
   mlir::MLIRContext &context = *unwrap(cctx);
   DialectRegistry registry;
-  prepareRegistry(registry);
+  mlir::enzyme::prepareRegistry(registry);
+  mlir::enzyme::registerDialects(registry);
+  mlir::enzyme::registerInterfaces(registry);
+
   context.appendDialectRegistry(registry);
-  context.loadDialect<mlir::arith::ArithDialect>();
-  context.loadDialect<mlir::enzyme::EnzymeDialect>();
-  context.loadDialect<mlir::enzymexla::EnzymeXLADialect>();
-  // context.loadDialect<mlir::triton::TritonDialect>();
-  context.loadDialect<mlir::tpu::TPUDialect>();
-  context.loadDialect<mlir::tensor::TensorDialect>();
-  context.loadDialect<mlir::func::FuncDialect>();
-  context.loadDialect<mlir::mhlo::MhloDialect>();
-  context.loadDialect<mlir::stablehlo::StablehloDialect>();
-  context.loadDialect<mlir::chlo::ChloDialect>();
-  context.loadDialect<mlir::sdy::SdyDialect>();
-  context.loadDialect<mlir::LLVM::LLVMDialect>();
+  mlir::enzyme::loadAllRegisteredDialects(context);
 }
 
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
@@ -1145,53 +1132,14 @@ extern "C" void RegisterDialects(MlirContext cctx) {
 #include "xla/service/spmd/shardy/sdy_round_trip/pipelines.h"
 
 extern "C" void InitializePasses(MlirDialectRegistry creg) {
-  mlir::registerenzymePasses();
-  enzyme::registerenzymexlaPasses();
-
-  // Register the standard passes we want.
-  mlir::registerTransformsPasses();
-  mlir::registerLowerAffinePass();
-  mlir::registerSCCPPass();
-  mlir::registerInlinerPass();
-  mlir::registerSymbolDCEPass();
-  mlir::registerLoopInvariantCodeMotionPass();
-  mlir::registerConvertSCFToOpenMPPass();
-  mlir::affine::registerAffinePasses();
-  mlir::registerReconcileUnrealizedCastsPass();
-
-  /*
-    registry.addExtension(+[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
-      LLVM::LLVMFunctionType::attachInterface<MemRefInsider>(*ctx);
-      LLVM::LLVMArrayType::attachInterface<MemRefInsider>(*ctx);
-      LLVM::LLVMPointerType::attachInterface<MemRefInsider>(*ctx);
-      LLVM::LLVMStructType::attachInterface<MemRefInsider>(*ctx);
-      MemRefType::attachInterface<PtrElementModel<MemRefType>>(*ctx);
-      LLVM::LLVMStructType::attachInterface<
-          PtrElementModel<LLVM::LLVMStructType>>(*ctx);
-      LLVM::LLVMPointerType::attachInterface<
-          PtrElementModel<LLVM::LLVMPointerType>>(*ctx);
-      LLVM::LLVMArrayType::attachInterface<PtrElementModel<LLVM::LLVMArrayType>>(
-          *ctx);
-    });
-    */
-
-  // Transform dialect and extensions.
-  mlir::transform::registerInterpreterPass();
-  mlir::enzyme::registerGenerateApplyPatternsPass();
-  mlir::enzyme::registerRemoveTransformPass();
-
-  // xla + shardy specific passes
-  xla::sdy::registerSdyRoundTripExportPipeline();
-  xla::sdy::registerSdyRoundTripImportPipeline();
-  mlir::sdy::registerAllSdyPassesAndPipelines();
-  xla::sdy::registerStablehloExportPipeline();
-  xla::sdy::registerStablehloImportPipeline();
-  xla::sdy::registerStablehloImportShardingsPass();
+  mlir::enzyme::initializePasses();
 }
 
 extern "C" void InitializeRegistry(MlirDialectRegistry creg) {
   mlir::DialectRegistry &registry = *unwrap(creg);
-  prepareRegistry(registry);
+  mlir::enzyme::prepareRegistry(registry);
+  mlir::enzyme::registerDialects(registry);
+  mlir::enzyme::registerInterfaces(registry);
 
   mlir::registerLLVMDialectImport(registry);
   mlir::registerNVVMDialectImport(registry);
@@ -1424,7 +1372,7 @@ ifrt_pjrt_array_create(ifrt::PjRtClient *client,
 
 // we might me interested in the `Compiler::Compile` method variant that accepts
 // `Topology`
-extern "C" xla::ifrt::LoadedExecutable *
+extern "C" HeldIfrtLoadedExecutable *
 ifrt_compile(ifrt::Client *client, MlirModule cmod, int64_t device_id,
              const int64_t *mesh_ids, int64_t num_mesh_ids,
              const char *xla_gpu_cuda_data_dir, bool use_shardy_partitioner,
@@ -1452,9 +1400,8 @@ ifrt_compile(ifrt::Client *client, MlirModule cmod, int64_t device_id,
       std::make_unique<xla::ifrt::HloProgram>(xla::ifrt::HloProgram(cmod_op));
   auto compiler = client->GetDefaultCompiler();
 
-  return MyValueOrThrow(
-             compiler->Compile(std::move(program), std::move(options)))
-      .release();
+  return reactant::capture(MyValueOrThrow(
+      compiler->CompileAndLoad(std::move(program), std::move(options))));
 }
 
 extern "C" void
@@ -2325,19 +2272,20 @@ extern "C" mlir::sdy::TensorShardingAttr hloShardingToTensorShardingAttr(
 
   return mlir::sdy::TensorShardingAttr::get(
       context, meshName, tensorShardingAttr.getDimShardings(),
-      tensorShardingAttr.getReplicatedAxes());
+      tensorShardingAttr.getReplicatedAxes(),
+      tensorShardingAttr.getUnreducedAxes());
 }
 
 #pragma endregion
 
 #pragma region ifrt::LoadedExecutable
 
-extern "C" void ifrt_loaded_executable_dtor(ifrt::LoadedExecutable *exec) {
+extern "C" void ifrt_loaded_executable_dtor(HeldIfrtLoadedExecutable *exec) {
   delete exec;
 }
 
 extern "C" void ifrt_loaded_executable_execute(
-    ifrt::LoadedExecutable *exec, int num_args,
+    HeldIfrtLoadedExecutable *exec, int num_args,
     HeldValue<tsl::RCReference<ifrt::Array>> **op_args,
     uint8_t *is_arg_donatable, int num_results,
     HeldValue<tsl::RCReference<ifrt::Array>> **op_results, uint8_t *futures,
@@ -2355,7 +2303,7 @@ extern "C" void ifrt_loaded_executable_execute(
   }
   options.fill_status = true;
 
-  auto result = MyValueOrThrow(exec->Execute(
+  auto result = MyValueOrThrow(exec->obj()->Execute(
       static_cast<absl::Span<tsl::RCReference<xla::ifrt::Array>>>(args),
       options, /* devices */ std::nullopt));
 
@@ -2376,16 +2324,16 @@ extern "C" void ifrt_loaded_executable_execute(
 }
 
 extern "C" ifrt::Client *
-ifrt_loaded_executable_client(ifrt::LoadedExecutable *exec) {
-  return exec->client();
+ifrt_loaded_executable_client(HeldIfrtLoadedExecutable *exec) {
+  return exec->obj()->client();
 }
 
 extern "C" void
-ifrt_loaded_executable_get_parameter_shardings(ifrt::LoadedExecutable *exec,
+ifrt_loaded_executable_get_parameter_shardings(HeldIfrtLoadedExecutable *exec,
                                                xla::OpSharding **op_shardings,
                                                int32_t num_op_shardings) {
   std::optional<std::vector<xla::OpSharding>> shardings =
-      exec->GetParameterShardings();
+      exec->obj()->GetParameterShardings();
   if (!shardings.has_value()) {
     ReactantThrowError(
         "No sharding found for the output of the loaded executable");
@@ -2405,11 +2353,11 @@ ifrt_loaded_executable_get_parameter_shardings(ifrt::LoadedExecutable *exec,
 }
 
 extern "C" void
-ifrt_loaded_executable_get_output_shardings(ifrt::LoadedExecutable *exec,
+ifrt_loaded_executable_get_output_shardings(HeldIfrtLoadedExecutable *exec,
                                             xla::OpSharding **op_shardings,
                                             int32_t num_op_shardings) {
   std::optional<std::vector<xla::OpSharding>> shardings =
-      exec->GetOutputShardings();
+      exec->obj()->GetOutputShardings();
   if (!shardings.has_value()) {
     ReactantThrowError(
         "No sharding found for the output of the loaded executable");
@@ -2429,18 +2377,18 @@ ifrt_loaded_executable_get_output_shardings(ifrt::LoadedExecutable *exec,
 }
 
 extern "C" void
-ifrt_loaded_executable_get_hlo_modules(ifrt::LoadedExecutable *exec,
+ifrt_loaded_executable_get_hlo_modules(HeldIfrtLoadedExecutable *exec,
                                        void **hlo_modules, int32_t *nmodules) {
-  auto hlo_modules_vec = MyValueOrThrow(exec->GetHloModules());
+  auto hlo_modules_vec = MyValueOrThrow(exec->obj()->GetHloModules());
   *nmodules = hlo_modules_vec.size();
   for (int32_t i = 0; i < *nmodules; i++) {
     hlo_modules[i] = reactant::capture(hlo_modules_vec[i]);
   }
 }
 
 extern "C" int32_t
-ifrt_loaded_executable_num_devices(ifrt::LoadedExecutable *exec) {
-  return static_cast<int32_t>(exec->num_devices());
+ifrt_loaded_executable_num_devices(HeldIfrtLoadedExecutable *exec) {
+  return static_cast<int32_t>(exec->obj()->num_devices());
 }
 
 #pragma endregion

deps/ReactantExtra/BUILD

@@ -761,25 +761,18 @@ cc_library(
         ],
 
     ) + [
-        "@enzyme_ad//src/enzyme_ad/jax:RegistryUtils.cpp",
         "@enzyme_ad//src/enzyme_ad/jax:gpu.cc",
         "@enzyme_ad//src/enzyme_ad/jax:cpu.cc",
-        # "@com_google_protobuf//:src/google/protobuf/io/coded_stream.cc",
-        # "@xla//xla:xla.pb.cc",
-        "@xla//xla:xla_data.pb.cc",
-        # "@xla//xla/stream_executor:device_description.pb.cc",
-        "@xla//xla/service:hlo.pb.cc",
-        #  # "@tsl//tsl/protobuf:dnn.pb.cc",
-        #"@tsl//tsl/protobuf:histogram.pb.cc",
-        #"@tsl//tsl/protobuf:bfc_memory_map_proto",bfc_memory_map.pb.cc",
         "@xla//xla/service/gpu:backend_configs.pb.cc",
         "@xla//xla:autotuning.pb.cc",
         "@xla//xla:autotune_results.pb.cc",
         "@xla//xla/service:buffer_assignment.pb.cc",
         ],
     hdrs = glob([
         "*.h",
-    ]),
+    ]) + [
+        "@enzyme_ad//src/enzyme_ad/jax:RegistryUtils.h",
+    ],
     copts = [
         "-Werror=unused-variable",
         "-Werror=unused-but-set-variable",
@@ -901,6 +894,7 @@ cc_library(
 "-Wl,-exported_symbol,_addSdyPropagationPipeline",
     ]}),
     deps = [
+        "@enzyme_ad//src/enzyme_ad/jax:RegistryUtils",
         "@enzyme//:EnzymeMLIR",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:AllPassesAndDialects",
@@ -1028,6 +1022,24 @@ cc_library(
         "@xla//xla/stream_executor/cuda:cuda_compute_capability_proto_cc_impl",
         "@xla//xla/service:gpu_plugin",
         "@xla//xla/pjrt/c:pjrt_c_api_gpu",
+        "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
+        "@xla//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client",
+        "@xla//xla/pjrt/plugin/xla_tpu:xla_tpu_pjrt_client",
+        "@stablehlo//:linalg_passes",
+        "@stablehlo//:tosa_passes",
+        "@stablehlo//:stablehlo_passes",
+        "@stablehlo//:stablehlo_passes_optimization",
+        "@stablehlo//stablehlo/tests:check_ops",
+
+        "@tsl//tsl/platform:env",
+        "@xla//xla/tsl/protobuf:dnn_proto_cc_impl",
+        "@xla//xla/tsl/protobuf:histogram_proto_cc",
+        "@xla//xla/tsl/protobuf:histogram_proto_cc_impl",
+        "@xla//xla:xla_data_proto_cc_impl",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/service:hlo_proto_cc_impl",
+        "@com_google_absl//absl/status:statusor",
     ] + select({
         "@xla//xla/tsl:is_cuda_enabled_and_oss":[
             "@xla//xla/stream_executor:cuda_platform",

deps/ReactantExtra/WORKSPACE

@@ -9,7 +9,7 @@ http_archive(
     urls = ["https://github.com/wsmoses/nsync/archive/{commit}.tar.gz".format(commit = NSYNC_COMMIT)],
 )
 
-ENZYMEXLA_COMMIT = "36230482acbbf9b32cca2e3906b7f08f4acc1e45"
+ENZYMEXLA_COMMIT = "7872c7761de86316c45a4df72210b1af38e09f07"
 ENZYMEXLA_SHA256 = ""
 
 http_archive(