removal interface for enzymexla.gpu_wrapper

Author: pengmai

src/enzyme_ad/jax/Implementations/EnzymeXLAAutoDiffOpInterfaceImpl.cpp

@@ -15,8 +15,13 @@
 #include "Enzyme/MLIR/Interfaces/AutoDiffOpInterface.h"
 #include "Enzyme/MLIR/Interfaces/GradientUtils.h"
 #include "Enzyme/MLIR/Interfaces/GradientUtilsReverse.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/RegionUtils.h"
 #include "src/enzyme_ad/jax/Implementations/SHLOGenericBatchOpInterface.h"
 
 #include "Dialect/Ops.h"
@@ -69,12 +74,92 @@ struct GPUWrapperOpEnzymeOpsRemover
     if (gradients.empty() && pushedCaches.empty())
       return success();
 
-    if (gradients.size())
-      return failure();
+    llvm::MapVector<Value, CacheInfo> cachesMap;
+    for (auto &it : *wrapOp.getBody()) {
+      Operation *op = &it;
+      if (auto pushOp = dyn_cast<enzyme::PushOp>(op)) {
+        CacheInfo info(pushOp.getCache());
+        if (cachesMap.contains(pushOp.getValue()))
+          info = info.merge(cachesMap.lookup(pushOp.getValue()), rewriter);
+        cachesMap[pushOp.getValue()] = info;
+      }
+    }
+    SmallVector<CacheInfo> caches =
+        llvm::map_to_vector(cachesMap, [](auto p) { return std::get<1>(p); });
+
+    if (caches.empty())
+      return success();
+
+    SetVector<Value> visited;
+    getUsedValuesDefinedAbove(wrapOp.getBodyRegion(), visited);
+    SmallVector<Value> frontier = llvm::map_to_vector(
+        caches, [](CacheInfo info) { return info.pushedValue(); });
+    SetVector<Operation *> opsToMove;
+    // Traverse backward from pushed values to find operations that the pushed
+    // value depends on
+    while (!frontier.empty()) {
+      Value v = frontier.back();
+      Operation *definingOp = v.getDefiningOp();
+      frontier.pop_back();
+
+      if (!definingOp)
+        continue;
+
+      // Assume allocations and frees are legal to move
+      if (hasEffect<MemoryEffects::Read>(definingOp) ||
+          hasEffect<MemoryEffects::Write>(definingOp)) {
+        definingOp->emitError() << "cannot move op with side effects";
+        return failure();
+      }
+      opsToMove.insert(definingOp);
+
+      for (Value operand : definingOp->getOperands()) {
+        if (visited.contains(operand))
+          continue;
+
+        frontier.push_back(operand);
+        visited.insert(operand);
+      }
+    }
 
-    if (pushedCaches.size())
-      return failure();
+    // Move the push and dependent values outside of the wrapper
+    OpBuilder::InsertionGuard guard(rewriter);
+    IRMapping map;
+    rewriter.setInsertionPoint(wrapOp);
+    for (Operation *toMove : llvm::reverse(opsToMove)) {
+      Operation *cloned = rewriter.clone(*toMove, map);
+      toMove->replaceAllUsesWith(cloned->getResults());
+
+      if (auto allocOp = dyn_cast<memref::AllocOp>(cloned)) {
+        // Assume GPU allocations need to be in address space 1
+        auto gpuAlloc = gpu::AllocOp::create(
+            rewriter, allocOp.getLoc(),
+            *allocOp.getType().clonePtrWith(rewriter.getI64IntegerAttr(1),
+                                            std::nullopt),
+            /*asyncDependencies=*/ValueRange(), allocOp.getDynamicSizes(),
+            /*symbolOperands=*/ValueRange());
+        allocOp.replaceAllUsesWith(gpuAlloc.getResult(0));
+        rewriter.eraseOp(allocOp);
+      }
+    }
 
+    for (auto &info : caches) {
+      rewriter.moveOpBefore(info.pushOp, wrapOp);
+      auto revWrapper = info.popOp->getParentOfType<enzymexla::GPUWrapperOp>();
+      assert(revWrapper && "failed to find reverse gpu_wrapper");
+      rewriter.moveOpBefore(info.popOp, revWrapper);
+
+      for (auto user : info.popOp.getResult().getUsers()) {
+        if (isa<memref::DeallocOp>(user)) {
+          rewriter.eraseOp(user);
+        }
+      }
+      rewriter.setInsertionPointAfter(revWrapper);
+      gpu::DeallocOp::create(rewriter, wrapOp.getLoc(), TypeRange(),
+                             info.popOp.getResult());
+    }
+
+    return success();
     // TODO need to convert to gpu allocations and conversion/copy
 
     /*
@@ -214,7 +299,7 @@ class Pointer2MemrefRev : public ReverseAutoDiffOpInterface::ExternalModel<
       Value dres = gutils->invertPointerM(p2m.getSource(), builder);
       Value shadow = builder.create<enzymexla::Pointer2MemrefOp>(
           p2m.getLoc(), p2m.getType(), dres);
-      gutils->setDiffe(p2m, shadow, builder);
+      gutils->setInvertedPointer(p2m, shadow);
     }
   }
 };

src/enzyme_ad/jax/Implementations/XLADerivatives.h

@@ -10,15 +10,13 @@
 
 namespace mlir {
 namespace enzyme {
-void registerEnzymeXLADialectAutoDiffInterface(mlir::DialectRegistry &registry);
 void registerMHLODialectAutoDiffInterface(mlir::DialectRegistry &registry);
 void registerStableHLODialectAutoDiffInterface(mlir::DialectRegistry &registry);
 void registerCHLODialectAutoDiffInterface(mlir::DialectRegistry &registry);
 void registerEnzymeXLADialectAutoDiffInterface(mlir::DialectRegistry &registry);
 
 static inline void
 registerXLAAutoDiffInterfaces(mlir::DialectRegistry &registry) {
-  registerEnzymeXLADialectAutoDiffInterface(registry);
   registerMHLODialectAutoDiffInterface(registry);
   registerStableHLODialectAutoDiffInterface(registry);
   registerCHLODialectAutoDiffInterface(registry);