continuing gpu backend (#1035)

* continuing

* fmt

* raise fix

* fmt

* fix

* fmt

* fix debug1

* debug info fix

* fmt

Author: wsmoses

src/enzyme_ad/jax/Passes/ConvertParallelToGPU.cpp

@@ -1541,6 +1541,11 @@ struct ParallelToGPULaunch : public OpRewritePattern<enzymexla::GPUWrapperOp> {
 
     rewriter.setInsertionPoint(wrapper);
     auto errOp = rewriter.create<enzymexla::GPUErrorOp>(loc);
+
+    for (auto atname : {"passthrough", "target_features"})
+      if (auto attr = wrapper->getAttr(atname)) {
+        errOp->setAttr(atname, attr);
+      }
     rewriter.setInsertionPointToStart(errOp.getBody());
     rewriter.eraseOp(wrapper.getBody()->getTerminator());
     rewriter.inlineBlockBefore(wrapper.getBody(),
@@ -2238,6 +2243,51 @@ gdgo->erase();
       signalPassFailure();
       return;
     }
+    SymbolTableCollection symbolTable;
+    symbolTable.getSymbolTable(getOperation());
+    getOperation()->walk([&](GPUErrorOp err) {
+      std::string sm;
+      if (auto attr =
+              dyn_cast_or_null<ArrayAttr>(err->getAttr("passthrough"))) {
+        for (auto a : attr) {
+          if (auto ar = dyn_cast<ArrayAttr>(a)) {
+            if (ar.size() != 2)
+              continue;
+            auto s0 = dyn_cast<StringAttr>(ar[0]);
+            auto s1 = dyn_cast<StringAttr>(ar[1]);
+            if (!s0 || !s1)
+              continue;
+            if (s0.getValue() == "target-cpu")
+              sm = s1.getValue();
+          }
+        }
+      }
+      std::string feat;
+      if (auto attr = dyn_cast_or_null<LLVM::TargetFeaturesAttr>(
+              err->getAttr("target_features"))) {
+        feat = attr.getFeaturesString();
+      }
+
+      err->walk([&](gpu::LaunchFuncOp launch) {
+        auto gfunc = dyn_cast_or_null<gpu::GPUFuncOp>(
+            symbolTable.lookupNearestSymbolFrom(launch, launch.getKernel()));
+        if (!gfunc)
+          return;
+        auto gmod = cast<gpu::GPUModuleOp>(gfunc->getParentOp());
+        if (!gmod.getTargetsAttr()) {
+          auto chip = sm;
+          if (chip.size() == 0)
+            chip = "sm_50";
+          auto features = feat;
+          if (features.size() == 0)
+            features = "+ptx60";
+          auto target = NVVM::NVVMTargetAttr::get(
+              gmod.getContext(), /*optLevel*/ 2,
+              /*triple*/ "nvptx64-nvidia-cuda", chip, features);
+          gmod.setTargetsAttr(ArrayAttr::get(gmod.getContext(), target));
+        }
+      });
+    });
   }
 };
 

src/enzyme_ad/jax/Passes/ConvertPolygeistToLLVM.cpp

@@ -38,6 +38,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -1718,8 +1719,18 @@ LogicalResult ConvertLaunchFuncOpToGpuRuntimeCallPattern::matchAndRewrite(
             ctorBuilder.create<LLVM::AddressOfOp>(loc, fatBinWrapper);
         auto bitcastOfWrapper = ctorBuilder.create<LLVM::BitcastOp>(
             loc, llvmPointerType, addressOfWrapper);
-        auto module = rtRegisterFatBinaryCallBuilder.create(loc, ctorBuilder,
-                                                            {bitcastOfWrapper});
+
+        auto cudaRegisterFatbinFn = LLVM::lookupOrCreateFn(
+            rewriter, moduleOp, "__cudaRegisterFatBinary", llvmPointerType,
+            llvmPointerType);
+        if (failed(cudaRegisterFatbinFn)) {
+          llvm::errs() << " cudamalloc already exists with different types\n";
+          return failure();
+        }
+
+        auto module = rewriter.create<LLVM::CallOp>(
+            loc, cudaRegisterFatbinFn.value(), ValueRange(bitcastOfWrapper));
+
         auto moduleGlobalName =
             std::string(llvm::formatv("polygeist_{0}_module_ptr", moduleName));
         {
@@ -1771,12 +1782,32 @@ LogicalResult ConvertLaunchFuncOpToGpuRuntimeCallPattern::matchAndRewrite(
             auto aoo = ctorBuilder.create<LLVM::AddressOfOp>(loc, stub);
             auto bitcast =
                 ctorBuilder.create<LLVM::BitcastOp>(loc, llvmPointerType, aoo);
-            auto ret = rtRegisterFunctionCallBuilder.create(
-                loc, ctorBuilder,
-                {module.getResult(), bitcast, kernelName, kernelName,
-                 /* TODO I have no idea what the following params are */
-                 ctorBuilder.create<LLVM::ConstantOp>(loc, llvmInt32Type, -1),
-                 nullPtr, nullPtr, nullPtr, nullPtr, nullPtr});
+
+            Type tys[] = {llvmPointerType, llvmPointerType, llvmPointerType,
+                          llvmPointerType, llvmInt32Type,   llvmPointerType,
+                          llvmPointerType, llvmPointerType, llvmPointerType,
+                          llvmPointerType};
+            auto cudaRegisterFn = LLVM::lookupOrCreateFn(
+                rewriter, moduleOp, "__cudaRegisterFunction", tys,
+                llvmInt32Type);
+            if (failed(cudaRegisterFn)) {
+              llvm::errs()
+                  << " cudamalloc already exists with different types\n";
+              return failure();
+            }
+            Value args[] = {
+                module.getResult(),
+                bitcast,
+                kernelName,
+                kernelName,
+                ctorBuilder.create<LLVM::ConstantOp>(loc, llvmInt32Type, -1),
+                nullPtr,
+                nullPtr,
+                nullPtr,
+                nullPtr,
+                nullPtr};
+
+            rewriter.create<LLVM::CallOp>(loc, cudaRegisterFn.value(), args);
           } else if (LLVM::GlobalOp g = dyn_cast<LLVM::GlobalOp>(op)) {
             int addrSpace = g.getAddrSpace();
             if (addrSpace != 1 /* device */ && addrSpace != 4 /* constant */)
@@ -1825,9 +1856,18 @@ LogicalResult ConvertLaunchFuncOpToGpuRuntimeCallPattern::matchAndRewrite(
           }
         }
         // TODO this has to happen only for some CUDA versions
-        if (gpuTarget == "cuda")
-          rtRegisterFatBinaryEndCallBuilder.create(loc, ctorBuilder,
-                                                   {module.getResult()});
+        if (gpuTarget == "cuda") {
+          auto cudaRegisterFatbinFn = LLVM::lookupOrCreateFn(
+              rewriter, moduleOp, "__cudaRegisterFatBinaryEnd", llvmPointerType,
+              llvmVoidType);
+          if (failed(cudaRegisterFatbinFn)) {
+            llvm::errs() << " cudamalloc already exists with different types\n";
+            return failure();
+          }
+
+          rewriter.create<LLVM::CallOp>(loc, cudaRegisterFatbinFn.value(),
+                                        ValueRange(module->getResult(0)));
+        }
         ctorBuilder.create<LLVM::ReturnOp>(loc, ValueRange());
       }
       auto ctorSymbol = FlatSymbolRefAttr::get(ctor);
@@ -1847,8 +1887,17 @@ LogicalResult ConvertLaunchFuncOpToGpuRuntimeCallPattern::matchAndRewrite(
         auto aoo = dtorBuilder.create<LLVM::AddressOfOp>(loc, moduleGlobal);
         auto module = dtorBuilder.create<LLVM::LoadOp>(
             loc, llvmPointerPointerType, aoo->getResult(0));
-        rtUnregisterFatBinaryCallBuilder.create(loc, dtorBuilder,
-                                                module.getResult());
+
+        auto cudaUnRegisterFatbinFn = LLVM::lookupOrCreateFn(
+            rewriter, moduleOp, "__cudaUnregisterFatBinary", llvmPointerType,
+            llvmVoidType);
+        if (failed(cudaUnRegisterFatbinFn)) {
+          llvm::errs() << " cudamalloc already exists with different types\n";
+          return failure();
+        }
+
+        rewriter.create<LLVM::CallOp>(loc, cudaUnRegisterFatbinFn.value(),
+                                      ValueRange(module));
         dtorBuilder.create<LLVM::ReturnOp>(loc, ValueRange());
         auto dtorSymbol = FlatSymbolRefAttr::get(dtor);
         {
@@ -2469,6 +2518,34 @@ struct ReturnOpLowering : public ConvertOpToLLVMPattern<func::ReturnOp> {
   }
 };
 
+/// Pattern for returning from a function, packs the results into a struct.
+struct GPUReturnOpLowering : public ConvertOpToLLVMPattern<gpu::ReturnOp> {
+public:
+  using ConvertOpToLLVMPattern<gpu::ReturnOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(gpu::ReturnOp returnOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (returnOp->getNumOperands() <= 1) {
+      rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(returnOp,
+                                                  adaptor.getOperands());
+      return success();
+    }
+
+    auto returnedType = LLVM::LLVMStructType::getLiteral(
+        returnOp->getContext(),
+        llvm::to_vector(adaptor.getOperands().getTypes()));
+    Value packed =
+        rewriter.create<LLVM::UndefOp>(returnOp->getLoc(), returnedType);
+    for (const auto &[index, value] : llvm::enumerate(adaptor.getOperands())) {
+      packed = rewriter.create<LLVM::InsertValueOp>(returnOp->getLoc(), packed,
+                                                    value, index);
+    }
+    rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(returnOp, packed);
+    return success();
+  }
+};
+
 /// TODO: Temporary until we migrate everything to opaque pointers
 struct ReconcileUnrealizedPointerCasts
     : public OpRewritePattern<UnrealizedConversionCastOp> {
@@ -2558,6 +2635,23 @@ populateCStyleMemRefLoweringPatterns(RewritePatternSet &patterns,
   patterns.add<CMemcpyOpLowering>(typeConverter);
 }
 
+/// Appends the patterns lowering operations from the Func dialect to the LLVM
+/// dialect using the C-style type conversion, i.e. converting memrefs to
+/// pointer to arrays of arrays.
+static void
+populateCStyleGPUFuncLoweringPatterns(RewritePatternSet &patterns,
+                                      LLVMTypeConverter &typeConverter,
+                                      std::string gpuTarget) {
+  patterns.add<GPUReturnOpLowering>(typeConverter);
+  patterns.add<GPUFuncOpLowering>(
+      typeConverter,
+      /*allocaAddrSpace=*/0,
+      StringAttr::get(&typeConverter.getContext(),
+                      gpuTarget == "cuda"
+                          ? NVVM::NVVMDialect::getKernelFuncAttrName()
+                          : ROCDL::ROCDLDialect::getKernelFuncAttrName()));
+}
+
 /// Appends the patterns lowering operations from the Func dialect to the LLVM
 /// dialect using the C-style type conversion, i.e. converting memrefs to
 /// pointer to arrays of arrays.
@@ -2618,6 +2712,13 @@ struct ConvertPolygeistToLLVMPass
 
     RewritePatternSet patterns(&getContext());
 
+    auto gpuTarget = "cuda";
+
+    // Insert our custom version of GPUFuncLowering
+    if (useCStyleMemRef) {
+      populateCStyleGPUFuncLoweringPatterns(patterns, converter, gpuTarget);
+    }
+
     populatePolygeistToLLVMConversionPatterns(converter, patterns);
     populateSCFToControlFlowConversionPatterns(patterns);
     // populateForBreakToWhilePatterns(patterns);
@@ -2642,7 +2743,6 @@ struct ConvertPolygeistToLLVMPass
 
     // Our custom versions of the gpu patterns
     if (useCStyleMemRef) {
-      auto gpuTarget = "cuda";
       patterns.add<ConvertLaunchFuncOpToGpuRuntimeCallPattern>(
           converter, "gpu.binary", gpuTarget);
       // patterns.add<LegalizeLaunchFuncOpPattern>(

src/enzyme_ad/jax/Passes/GPULaunchRecognition.cpp

@@ -91,7 +91,6 @@ struct GPULaunchRecognitionPass
               cop.getArgOperands()[0].getDefiningOp<LLVM::AddressOfOp>();
           if (!argop)
             continue;
-          llvm::errs() << "argop: " << argop << "\n";
           auto cur = argop.getFunction(symbolTable);
           if (!cur)
             continue;
@@ -156,14 +155,12 @@ struct GPULaunchRecognitionPass
             });
           }
 
-          auto loc = launchFunc->getLoc();
+          auto loc = cop->getLoc();
           builder.setInsertionPointAfter(cop);
 
           auto shMemSize = builder.create<LLVM::TruncOp>(
               loc, builder.getI32Type(), cop.getArgOperands()[7]);
           auto stream = cop.getArgOperands()[8];
-          llvm::errs() << " stream: " << stream << "\n";
-          // TODO stream is arg 8
           llvm::SmallVector<mlir::Value> args;
           for (unsigned i = 9; i < cop.getArgOperands().size(); i++)
             args.push_back(cop.getArgOperands()[i]);
@@ -194,8 +191,8 @@ struct GPULaunchRecognitionPass
                   ValueRange(args));
             } else {
               auto op = builder.create<mlir::gpu::LaunchOp>(
-                  loc, grid[0], grid[1], grid[2], block[0], block[1], block[2],
-                  shMemSize, nullptr, ValueRange());
+                  launchFunc->getLoc(), grid[0], grid[1], grid[2], block[0],
+                  block[1], block[2], shMemSize, nullptr, ValueRange());
               builder.setInsertionPointToStart(&op.getRegion().front());
               builder.create<LLVM::CallOp>(loc, cur, args);
               builder.create<gpu::TerminatorOp>(loc);
@@ -208,8 +205,9 @@ struct GPULaunchRecognitionPass
                   ValueRange(args), stream.getType(), ValueRange(stream));
             } else {
               auto op = builder.create<mlir::gpu::LaunchOp>(
-                  loc, grid[0], grid[1], grid[2], block[0], block[1], block[2],
-                  shMemSize, stream.getType(), ValueRange(stream));
+                  launchFunc->getLoc(), grid[0], grid[1], grid[2], block[0],
+                  block[1], block[2], shMemSize, stream.getType(),
+                  ValueRange(stream));
               builder.setInsertionPointToStart(&op.getRegion().front());
               builder.create<LLVM::CallOp>(loc, cur, args);
               builder.create<gpu::TerminatorOp>(loc);

src/enzyme_ad/jax/Passes/ParallelLower.cpp

@@ -40,6 +40,7 @@ namespace mlir {
 namespace enzyme {
 #define GEN_PASS_DEF_PARALLELLOWER
 #define GEN_PASS_DEF_FIXGPUFUNC
+#define GEN_PASS_DEF_STRIPGPUINFO
 #include "src/enzyme_ad/jax/Passes/Passes.h.inc"
 } // namespace enzyme
 } // namespace mlir
@@ -113,6 +114,10 @@ struct FixGPUFunc : public enzyme::impl::FixGPUFuncBase<FixGPUFunc> {
   using FixGPUFuncBase::FixGPUFuncBase;
   void runOnOperation() override;
 };
+struct StripGPUInfo : public enzyme::impl::StripGPUInfoBase<StripGPUInfo> {
+  using StripGPUInfoBase::StripGPUInfoBase;
+  void runOnOperation() override;
+};
 } // end anonymous namespace
 
 /// Creates a pass to perform optimizations relying on memref dataflow such as
@@ -412,11 +417,16 @@ void ParallelLower::runOnOperation() {
       for (auto op : ops)
         callInliner(op);
     }
+    LLVM::LLVMFuncOp lfn = nullptr;
     {
       SmallVector<LLVM::CallOp> lops;
       launchOp.walk([&](LLVM::CallOp caller) { lops.push_back(caller); });
-      for (auto op : lops)
+      for (auto op : lops) {
+        if (!lfn)
+          lfn = dyn_cast_or_null<LLVM::LLVMFuncOp>(
+              op.resolveCallableInTable(&symbolTable));
         LLVMcallInliner(op);
+      }
     }
 
     mlir::IRRewriter builder(launchOp.getContext());
@@ -449,6 +459,14 @@ void ParallelLower::runOnOperation() {
           ValueRange({launchOp.getGridSizeX(), launchOp.getGridSizeY(),
                       launchOp.getGridSizeZ(), launchOp.getBlockSizeX(),
                       launchOp.getBlockSizeY(), launchOp.getBlockSizeZ()}));
+      if (lfn) {
+        if (auto passthrough = lfn.getPassthrough()) {
+          pw->setAttr("passthrough", *passthrough);
+        }
+        if (auto passthrough = lfn.getTargetFeatures()) {
+          pw->setAttr("target_features", *passthrough);
+        }
+      }
       builder.setInsertionPointToStart(pw.getBody());
     }
 
@@ -893,6 +911,22 @@ void ConvertCudaRTtoCPU::runOnOperation() {
 }
 #endif
 
+void StripGPUInfo::runOnOperation() {
+  getOperation()->walk([](gpu::GPUModuleOp v) {
+    auto unknown = OpBuilder(v).getUnknownLoc();
+    v->walk([&](Operation *op) {
+      op->setLoc(unknown);
+      for (auto &region : op->getRegions()) {
+        for (auto &blk : region) {
+          for (auto &arg : blk.getArguments()) {
+            arg.setLoc(unknown);
+          }
+        }
+      }
+    });
+  });
+}
+
 // Returns a list of all symbols provided by cudart (obtained from
 // libcudart_static.a)
 static std::vector<llvm::StringRef> getCudartSymbols();

src/enzyme_ad/jax/Passes/Passes.td

@@ -787,4 +787,9 @@ def FixGPUFunc : Pass<"fix-gpu-func", "mlir::gpu::GPUModuleOp"> {
   let dependentDialects = ["func::FuncDialect", "LLVM::LLVMDialect", "gpu::GPUDialect"];
 }
 
+def StripGPUInfo : Pass<"strip-gpu-info"> {
+  let summary = "Stirng GPU Debug info";
+  let dependentDialects = ["gpu::GPUDialect"];
+}
+
 #endif