Merge branch 'main' into wfelix_xla_dev

Author: felixwqp

.github/workflows/test-gb-25.yml

@@ -52,7 +52,7 @@ jobs:
           - 'mg/sharded-factors'
           # - '0123456789abcdef0123456789abcdef01234567'
         reactant_commit:
-          - 'main'
+          - 'ap/persistent_compile_cache'
           # - "ap/updated_no_nan"
 
     steps:
@@ -189,12 +189,12 @@ jobs:
         timeout-minutes: 60
         run: |
           export XLA_FLAGS='--xla_dump_to=${{ env.GB25_DIR }}/xla_dump'
-          timeout --signal=TERM --verbose 59m mpiexecjl -np 1 julia  --color=yes --project -O0 sharding/sharded_baroclinic_instability_simulation_run.jl
+          timeout --signal=TERM --verbose 59m mpiexecjl -np 1 julia --color=yes --project -O0 --startup-file=no --threads=16 --compiled-modules=strict sharding/sharded_baroclinic_instability_simulation_run.jl
         working-directory: ${{ env.GB25_DIR }}
       - name: Test correctness in GB-25 code
         timeout-minutes: 20
         run: |
-          timeout --signal=TERM --verbose 19m mpiexecjl -np 1 julia --color=yes --project -O0 correctness/correctness_sharded_baroclinic_instability_simulation_run.jl
+          timeout --signal=TERM --verbose 19m mpiexecjl -np 1 julia --color=yes --project -O0 --startup-file=no --threads=16 --compiled-modules=strict correctness/correctness_sharded_baroclinic_instability_simulation_run.jl
         working-directory: ${{ env.GB25_DIR }}
       - name: Upload MLIR and XLA modules
         uses: actions/upload-artifact@v4

src/enzyme_ad/jax/Passes/AffineCFG.cpp

@@ -2344,7 +2344,7 @@ struct MoveSelectToAffine : public OpRewritePattern<arith::SelectOp> {
 
     bool changed = false;
     auto condOp = ifOp.getCondition().getDefiningOp();
-    if (isa<AndIOp, OrIOp>(condOp)) {
+    if (condOp && isa<AndIOp, OrIOp>(condOp)) {
       // condition, Negated
 
       for (auto &opv : condOp->getOpOperands()) {

src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp

@@ -30,6 +30,7 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
+#include "Interfaces/AutoDiffTypeInterface.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 
 #include "src/enzyme_ad/jax/Dialect/Ops.h"
@@ -2529,6 +2530,7 @@ struct AffineToStableHLORaisingPass
         for (auto arg : operands0) {
 
           Attribute attr;
+
           if (matchPattern(arg, m_Constant(&attr))) {
             affine::AffineValueMap accessMap(AffineMap::get(arg.getContext()),
                                              {});
@@ -2539,15 +2541,21 @@ struct AffineToStableHLORaisingPass
             auto unrankedTensorType = RankedTensorType::get({}, ET);
             OpBuilder builder(arg.getContext());
             builder.setInsertionPointToEnd(newBlock);
-            auto newConst = builder.create<stablehlo::ConstantOp>(
-                arg.getLoc(), unrankedTensorType,
-                SplatElementsAttr::get(
-                    unrankedTensorType,
-                    ArrayRef<Attribute>(
-                        isIndex ? IntegerAttr::get(
-                                      ET, cast<IntegerAttr>(attr).getValue())
-                                : attr)));
-            auto newVal = newConst.getResult();
+            Value newVal;
+            if (arg.getDefiningOp<ub::PoisonOp>()) {
+              newVal = cast<mlir::enzyme::AutoDiffTypeInterface>(arg.getType())
+                           .createNullValue(builder, arg.getLoc());
+            } else {
+              auto newConst = builder.create<stablehlo::ConstantOp>(
+                  arg.getLoc(), unrankedTensorType,
+                  SplatElementsAttr::get(
+                      unrankedTensorType,
+                      ArrayRef<Attribute>(
+                          isIndex ? IntegerAttr::get(
+                                        ET, cast<IntegerAttr>(attr).getValue())
+                                  : attr)));
+              newVal = newConst.getResult();
+            }
             mapping.map(arg, newVal);
             maps[newVal] = accessMap;
             continue;

src/enzyme_ad/jax/Passes/LowerJIT.cpp

@@ -80,6 +80,22 @@
 
 #include "mlir/Target/LLVMIR/Export.h"
 
+#if (defined(_WIN32) || defined(__CYGWIN__)) &&                                \
+    !defined(MLIR_CAPI_ENABLE_WINDOWS_DLL_DECLSPEC)
+// Visibility annotations disabled.
+#define MLIR_CAPI_EXPORTED
+#elif defined(_WIN32) || defined(__CYGWIN__)
+// Windows visibility declarations.
+#if MLIR_CAPI_BUILDING_LIBRARY
+#define MLIR_CAPI_EXPORTED __declspec(dllexport)
+#else
+#define MLIR_CAPI_EXPORTED __declspec(dllimport)
+#endif
+#else
+// Non-windows: use visibility attributes.
+#define MLIR_CAPI_EXPORTED __attribute__((visibility("default")))
+#endif
+
 #define DEBUG_TYPE "lower-jit"
 
 namespace mlir {
@@ -274,7 +290,8 @@ bool initJIT() {
   return true;
 }
 
-extern "C" void EnzymeJaXMapSymbol(const char *name, void *symbol) {
+extern "C" MLIR_CAPI_EXPORTED void EnzymeJaXMapSymbol(const char *name,
+                                                      void *symbol) {
   initJIT();
   MappedSymbols[JIT->mangleAndIntern(name)] = llvm::orc::ExecutorSymbolDef(
       llvm::orc::ExecutorAddr::fromPtr(symbol), llvm::JITSymbolFlags());

src/enzyme_ad/jax/cpu.cc

@@ -2,6 +2,22 @@
 #include "xla/service/custom_call_target_registry.h"
 #include <cstring>
 
+#if (defined(_WIN32) || defined(__CYGWIN__)) &&                                \
+    !defined(MLIR_CAPI_ENABLE_WINDOWS_DLL_DECLSPEC)
+// Visibility annotations disabled.
+#define MLIR_CAPI_EXPORTED
+#elif defined(_WIN32) || defined(__CYGWIN__)
+// Windows visibility declarations.
+#if MLIR_CAPI_BUILDING_LIBRARY
+#define MLIR_CAPI_EXPORTED __declspec(dllexport)
+#else
+#define MLIR_CAPI_EXPORTED __declspec(dllimport)
+#endif
+#else
+// Non-windows: use visibility attributes.
+#define MLIR_CAPI_EXPORTED __attribute__((visibility("default")))
+#endif
+
 template <bool withError> struct CallInfo;
 
 template <> struct CallInfo<false> {
@@ -28,7 +44,7 @@ void forwarding_custom_call(void *out, const void **in, const void *opaque_ptr,
   }
 }
 
-extern "C" void RegisterEnzymeXLACPUHandler() {
+extern "C" MLIR_CAPI_EXPORTED void RegisterEnzymeXLACPUHandler() {
   xla::CustomCallTargetRegistry::Global()->Register(
       "enzymexla_compile_cpu", (void *)&forwarding_custom_call<false>, "Host");
   xla::CustomCallTargetRegistry::Global()->Register(

src/enzyme_ad/jax/gpu.cc

@@ -1,6 +1,22 @@
 #include "xla/ffi/api/ffi.h"
 #include "xla/ffi/ffi_api.h"
 
+#if (defined(_WIN32) || defined(__CYGWIN__)) &&                                \
+    !defined(MLIR_CAPI_ENABLE_WINDOWS_DLL_DECLSPEC)
+// Visibility annotations disabled.
+#define MLIR_CAPI_EXPORTED
+#elif defined(_WIN32) || defined(__CYGWIN__)
+// Windows visibility declarations.
+#if MLIR_CAPI_BUILDING_LIBRARY
+#define MLIR_CAPI_EXPORTED __declspec(dllexport)
+#else
+#define MLIR_CAPI_EXPORTED __declspec(dllimport)
+#endif
+#else
+// Non-windows: use visibility attributes.
+#define MLIR_CAPI_EXPORTED __attribute__((visibility("default")))
+#endif
+
 template <bool withError> struct CallInfo;
 
 template <> struct CallInfo<false> {
@@ -138,7 +154,7 @@ XLA_FFI_Error *execute(XLA_FFI_CallFrame *call_frame) {
   return nullptr;
 }
 
-extern "C" void RegisterEnzymeXLAGPUHandler() {
+extern "C" MLIR_CAPI_EXPORTED void RegisterEnzymeXLAGPUHandler() {
   XLA_FFI_Handler_Bundle bundle = {instantiate, prepare, initialize<false>,
                                    execute<false>};