drop no-longer-supported lifetime annotations (#1280)

LLVM IR dropped support for lifetime intrinsics on everything but allocas and
poisons while we still seem to be producing these. Remove them immediately
before lowering. This gives us the possibility to rewrite them into something
else if we figure out what.

Author: ftynse

src/enzyme_ad/jax/Passes/ConvertPolygeistToLLVM.cpp

@@ -3513,6 +3513,24 @@ populateCStyleFuncLoweringPatterns(RewritePatternSet &patterns,
   patterns.add<CallOpLowering, FuncOpLowering, ReturnOpLowering>(typeConverter);
 }
 
+static void removeUnsupportedLifeTimes(mlir::Operation *root) {
+  llvm::SmallVector<mlir::Operation *> toErase;
+  root->walk([&](mlir::Operation *op) {
+    if (auto lifetimeStart = llvm::dyn_cast<mlir::LLVM::LifetimeStartOp>(op)) {
+      if (!llvm::isa_and_nonnull<mlir::LLVM::AllocaOp, mlir::LLVM::PoisonOp>(
+              lifetimeStart.getPtr().getDefiningOp()))
+        toErase.push_back(op);
+    } else if (auto lifetimeEnd =
+                   llvm::dyn_cast<mlir::LLVM::LifetimeEndOp>(op)) {
+      if (!llvm::isa_and_nonnull<mlir::LLVM::AllocaOp, mlir::LLVM::PoisonOp>(
+              lifetimeEnd.getPtr().getDefiningOp()))
+        toErase.push_back(op);
+    }
+  });
+  for (mlir::Operation *op : toErase)
+    op->erase();
+}
+
 //===-----------------------------------------------------------------------===/
 
 namespace {
@@ -3727,6 +3745,8 @@ struct ConvertPolygeistToLLVMPass
         }
       });
     }
+
+    removeUnsupportedLifeTimes(m);
   }
 
   void runOnOperation() override {

src/enzyme_ad/jax/raise.cpp

@@ -156,6 +156,10 @@ extern "C" std::string runLLVMToMLIRRoundTrip(std::string input,
   llvm::LLVMContext llvmContext;
   llvmContext.setDiscardValueNames(false);
   auto outModule = translateModuleToLLVMIR(*mod, llvmContext);
+  if (!outModule) {
+    llvm::errs() << "failed to translate MLIR to LLVM IR\n";
+    return "";
+  }
 
   if (auto F = outModule->getFunction("mgpuModuleLoad")) {
     for (auto U : llvm::make_early_inc_range(F->users())) {