Polygeist: handle occupancy and related utilities (#1395)

* WIP: occupance

* wip

* b

* occ op

* fmt

* fix

* fix

* fix

* fix

* no comdat

* More kernel support

* fix

* attr

* Fix

* fix

* fix

* fn attr

* fix

* fix

* fix

* fix

Author: wsmoses

src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td

@@ -137,6 +137,23 @@ def GetStreamOp : EnzymeXLA_Op<"get_stream", [Pure]> {
 }
 
 
+def GPUOccupancyOp : EnzymeXLA_Op<"gpu_occupancy", [Pure, DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let arguments = (ins
+	SymbolRefAttr:$fn,
+	AnyType:$blockSize,
+	AnyType:$dynamicSMemSize,
+	AnyType:$flags
+);
+  let results = (outs AnyType : $result);
+}
+
+def GPUKernelAddressOp : EnzymeXLA_Op<"gpu_kernel_address", [Pure, DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let arguments = (ins
+	SymbolRefAttr:$fn
+  );
+  let results = (outs AnyType : $result);
+}
+
 def GPUWrapperOp : EnzymeXLA_Op<"gpu_wrapper", [
   RecursiveMemoryEffects,
   AffineScope,

src/enzyme_ad/jax/Dialect/Ops.cpp

@@ -68,6 +68,32 @@ static std::optional<int64_t> getConstant(Value v) {
   return {};
 }
 
+LogicalResult
+GPUOccupancyOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  // TODO: Verify that the result type is same as the type of the referenced
+  // func.func op.
+  auto global = symbolTable.lookupNearestSymbolFrom<FunctionOpInterface>(
+      *this, getFnAttr());
+  if (!global)
+    return emitOpError("'")
+           << getFn() << "' does not reference a valid global funcOp";
+
+  return success();
+}
+
+LogicalResult
+GPUKernelAddressOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  // TODO: Verify that the result type is same as the type of the referenced
+  // func.func op.
+  auto global = symbolTable.lookupNearestSymbolFrom<FunctionOpInterface>(
+      *this, getFnAttr());
+  if (!global)
+    return emitOpError("'")
+           << getFn() << "' does not reference a valid global funcOp";
+
+  return success();
+}
+
 LogicalResult
 KernelCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   // TODO: Verify that the result type is same as the type of the referenced

src/enzyme_ad/jax/Passes/ConvertParallelToGPU.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -2529,6 +2530,30 @@ gdgo->erase();
               gmod.getContext(), /*optLevel*/ 2,
               /*triple*/ "nvptx64-nvidia-cuda", chip, features);
           gmod.setTargetsAttr(ArrayAttr::get(gmod.getContext(), target));
+
+          DataLayoutSpecInterface dataLayout = {};
+          // Set index type size to 32 bits
+          {
+            auto ctx = gmod.getContext();
+            llvm::DenseMap<mlir::TypeAttr, mlir::DataLayoutEntryInterface>
+                typeEntries;
+            auto type = IndexType::get(ctx);
+            auto key = mlir::TypeAttr::get(type);
+            uint64_t size = 32;
+            auto params =
+                IntegerAttr::get(mlir::IntegerType::get(ctx, 64), size);
+            typeEntries.try_emplace(key,
+                                    DataLayoutEntryAttr::get(type, params));
+            SmallVector<DataLayoutEntryInterface> entries;
+            entries.reserve(typeEntries.size());
+            for (const auto &it : typeEntries)
+              entries.push_back(it.second);
+            dataLayout = DataLayoutSpecAttr::get(ctx, entries);
+          }
+          // gpuModule->setAttr(
+          //     LLVM::LLVMDialect::getDataLayoutAttrName(),
+          //     deviceModule->getAttr(LLVM::LLVMDialect::getDataLayoutAttrName()));
+          gmod->setAttr(DLTIDialect::kDataLayoutAttrName, dataLayout);
         }
       });
     });

src/enzyme_ad/jax/Passes/ConvertPolygeistToLLVM.cpp

@@ -2556,6 +2556,101 @@ class ConvertAllocOpToGpuRuntimeCallPattern
   }
 };
 
+class ConvertOccupancyOp
+    : public ConvertOpToGpuRuntimeCallPattern<enzymexla::GPUOccupancyOp> {
+public:
+  /// The attribute name to use instead of `gpu.kernel`.
+  StringRef backend;
+
+  ConvertOccupancyOp(LLVMTypeConverter &typeConverter, StringRef backend)
+      : ConvertOpToGpuRuntimeCallPattern<enzymexla::GPUOccupancyOp>(
+            typeConverter),
+        backend(backend) {}
+
+private:
+  LogicalResult
+  matchAndRewrite(enzymexla::GPUOccupancyOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)))
+      return failure();
+
+    if (backend != "cuda")
+      return rewriter.notifyMatchFailure(
+          op, "Occupancy op lowering only supported for CUDA");
+
+    auto moduleOp = op->getParentOfType<ModuleOp>();
+    auto i64 = rewriter.getIntegerType(64);
+    auto i32 = rewriter.getIntegerType(32);
+
+    auto intty = adaptor.getBlockSize().getType();
+    auto loc = op.getLoc();
+
+    auto ptrty = LLVM::LLVMPointerType::get(rewriter.getContext());
+    Type tys[] = {ptrty, ptrty, intty, adaptor.getDynamicSMemSize().getType(),
+                  adaptor.getFlags().getType()};
+
+    auto cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlagsFn =
+        LLVM::lookupOrCreateFn(
+            rewriter, moduleOp,
+            "cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags", tys, i32);
+    if (failed(cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlagsFn)) {
+      llvm::errs() << " cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags "
+                      "already exists with different types\n";
+      return failure();
+    }
+
+    auto one = rewriter.create<LLVM::ConstantOp>(loc, i64,
+                                                 rewriter.getI64IntegerAttr(1));
+
+    auto ptr = rewriter.create<LLVM::AllocaOp>(loc, ptrty, intty, one);
+
+    std::string funcStubName =
+        getFuncStubName(op.getFn().getRootReference().getValue(),
+                        op.getFn().getLeafReference().getValue());
+    auto addr = rewriter.create<LLVM::AddressOfOp>(loc, ptrty, funcStubName);
+    Value args[] = {ptr, addr, adaptor.getBlockSize(),
+                    adaptor.getDynamicSMemSize(), adaptor.getFlags()};
+    rewriter.create<LLVM::CallOp>(
+        loc, cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlagsFn.value(),
+        args);
+    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(op, intty, ptr);
+
+    return success();
+  }
+};
+
+class ConvertGPUKernelAddressOp
+    : public ConvertOpToGpuRuntimeCallPattern<enzymexla::GPUKernelAddressOp> {
+public:
+  /// The attribute name to use instead of `gpu.kernel`.
+  StringRef backend;
+
+  ConvertGPUKernelAddressOp(LLVMTypeConverter &typeConverter, StringRef backend)
+      : ConvertOpToGpuRuntimeCallPattern<enzymexla::GPUKernelAddressOp>(
+            typeConverter),
+        backend(backend) {}
+
+private:
+  LogicalResult
+  matchAndRewrite(enzymexla::GPUKernelAddressOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    if (backend != "cuda")
+      return rewriter.notifyMatchFailure(
+          op, "KernelAddress lowering only supported for CUDA");
+
+    std::string funcStubName =
+        getFuncStubName(op.getFn().getRootReference().getValue(),
+                        op.getFn().getLeafReference().getValue());
+
+    rewriter.replaceOpWithNewOp<LLVM::AddressOfOp>(op, op.getType(),
+                                                   funcStubName);
+
+    return success();
+  }
+};
+
 /// A rewrite pattern to convert gpu.alloc operations into a GPU runtime
 /// call. Currently it supports CUDA, CPU, and XLA.
 template <bool cStyle>
@@ -3938,6 +4033,10 @@ struct ConvertPolygeistToLLVMPass
       //     /*kernelIntersperseSizeCallConv*/ false);
       patterns.add<ConvertAllocOpToGpuRuntimeCallPattern<true>>(converter,
                                                                 gpuTarget);
+      patterns.add<ConvertOccupancyOp>(converter, gpuTarget);
+
+      patterns.add<ConvertGPUKernelAddressOp>(converter, gpuTarget);
+
       patterns.add<ConvertDeallocOpToGpuRuntimeCallPattern<true>>(converter,
                                                                   gpuTarget);
       patterns.add<ConvertXLAWrapperPattern<true>>(converter, gpuTarget);