Hoist wrapper-invariant scalar computations to the host

src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp

@@ -3576,6 +3576,39 @@ struct AffineToStableHLORaisingPass
   // raising identifies buffers by SSA root: a memory_space_cast view would
   // split one buffer into two roots and lose store propagation. Retarget the
   // accesses to the source and drop the cast.
+  // A pure scalar computed entirely from values defined outside the wrapper
+  // (a null check of an optional buffer, a host-side flag chain) is the
+  // host's to compute: hoist it out, so the kernel captures the resulting
+  // scalar instead of pointers no tensor can stand for.
+  static void hoistWrapperInvariantScalars(Operation *g) {
+    auto definedOutside = [&](Value v) {
+      if (auto ba = dyn_cast<BlockArgument>(v))
+        return !g->isProperAncestor(ba.getOwner()->getParentOp()) &&
+               ba.getOwner()->getParentOp() != g;
+      return !g->isProperAncestor(v.getDefiningOp());
+    };
+    bool changed = true;
+    while (changed) {
+      changed = false;
+      SmallVector<Operation *> toHoist;
+      g->walk([&](Operation *op) {
+        if (op->getNumOperands() == 0 || op->getNumRegions() ||
+            op->hasTrait<OpTrait::IsTerminator>() || !isMemoryEffectFree(op))
+          return;
+        if (!llvm::all_of(op->getResultTypes(),
+                          [](Type t) { return t.isIntOrIndexOrFloat(); }))
+          return;
+        if (!llvm::all_of(op->getOperands(), definedOutside))
+          return;
+        toHoist.push_back(op);
+      });
+      for (Operation *op : toHoist) {
+        op->moveBefore(g);
+        changed = true;
+      }
+    }
+  }
+
   static void stripAccessMemorySpaceCasts(Operation *root) {
     SmallVector<memref::MemorySpaceCastOp> casts;
     root->walk([&](memref::MemorySpaceCastOp c) { casts.push_back(c); });
@@ -3737,6 +3770,7 @@ struct AffineToStableHLORaisingPass
     op->walk([&](enzymexla::GPUWrapperOp g) { gwrap.push_back(g); });
     for (auto g : gwrap) {
       stripAccessMemorySpaceCasts(g);
+      hoistWrapperInvariantScalars(g);
       peelDynamicParallelDims(g);
     }
     size_t raised_count = 0;

test/lit_tests/raising/raise_hoist_nullcheck.mlir

@@ -0,0 +1,29 @@
+// RUN: enzymexlamlir-opt %s --raise-affine-to-stablehlo --split-input-file | FileCheck %s
+
+// The guard flag of an optional buffer is computed inside the kernel as a
+// null check of a captured pointer. The comparison only uses values defined
+// outside the wrapper, so it hoists to the host, and the kernel captures the
+// resulting scalar instead of pointers no tensor can stand for.
+func.func @nullguard(%out: memref<32xf64, 1>, %opt: !llvm.ptr<1>) {
+  %c1 = arith.constant 1 : index
+  %c32 = arith.constant 32 : index
+  %null = llvm.mlir.zero : !llvm.ptr<1>
+  %0 = "enzymexla.gpu_wrapper"(%c1, %c1, %c1, %c32, %c1, %c1) ({
+    %has = llvm.icmp "ne" %opt, %null : !llvm.ptr<1>
+    affine.parallel (%t) = (0) to (32) {
+      %v = arith.uitofp %has : i1 to f64
+      affine.store %v, %out[%t] : memref<32xf64, 1>
+    }
+    "enzymexla.polygeist_yield"() : () -> ()
+  }) : (index, index, index, index, index, index) -> index
+  return
+}
+
+// CHECK-LABEL: func.func @nullguard(
+// CHECK-SAME: %[[OUT:.+]]: memref<32xf64, 1>, %[[OPT:.+]]: !llvm.ptr<1>
+// CHECK: %[[NULL:.+]] = llvm.mlir.zero : !llvm.ptr<1>
+// CHECK-NEXT: %[[CMP:.+]] = llvm.icmp "ne" %[[OPT]], %[[NULL]] : !llvm.ptr<1>
+// CHECK-NEXT: %[[F:.+]] = arith.uitofp %[[CMP]] : i1 to f64
+// CHECK: affine.store %[[F]], %{{.+}}[] : memref<f64>
+// CHECK: enzymexla.xla_wrapper @rxla$raised_0 (%{{.+}}, %[[OUT]]) : (memref<f64, 1>, memref<32xf64, 1>) -> ()
+// CHECK: func.func private @rxla$raised_0(%{{.+}}: tensor<f64>, %{{.+}}: tensor<32xf64>) -> (tensor<f64>, tensor<32xf64>)