Fix c++ import bugs (#1310)

Author: wsmoses

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;