enzymexla.pointer2memref derivative

Author: pengmai

src/enzyme_ad/jax/Implementations/EnzymeXLAAutoDiffOpInterfaceImpl.cpp

@@ -1,4 +1,4 @@
-//===- CHLOAutoDiffOpInterfaceImpl.cpp - Interface external model --------===//
+//===- EnzymeXLAAutoDiffOpInterfaceImpl.cpp - Interface external model ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -7,7 +7,7 @@
 //===----------------------------------------------------------------------===//
 //
 // This file contains the external model implementation of the automatic
-// differentiation op interfaces for the upstream MLIR arithmetic dialect.
+// differentiation op interfaces for the EnzymeXLA dialect.
 //
 //===----------------------------------------------------------------------===//
 
@@ -193,6 +193,31 @@ struct GPUWrapperOpInterfaceReverse
                           MGradientUtilsReverse *gutils) const {}
 };
 
+class Pointer2MemrefRev : public ReverseAutoDiffOpInterface::ExternalModel<
+                              Pointer2MemrefRev, enzymexla::Pointer2MemrefOp> {
+public:
+  LogicalResult createReverseModeAdjoint(Operation *orig, OpBuilder &builder,
+                                         MGradientUtilsReverse *gutils,
+                                         SmallVector<Value> caches) const {
+    return success();
+  }
+
+  SmallVector<Value> cacheValues(Operation *orig,
+                                 MGradientUtilsReverse *gutils) const {
+    return SmallVector<Value>();
+  }
+
+  void createShadowValues(Operation *op, OpBuilder &builder,
+                          MGradientUtilsReverse *gutils) const {
+    auto p2m = cast<enzymexla::Pointer2MemrefOp>(op);
+    if (!gutils->isConstantValue(p2m)) {
+      Value dres = gutils->invertPointerM(p2m.getSource(), builder);
+      Value shadow = builder.create<enzymexla::Pointer2MemrefOp>(
+          p2m.getLoc(), p2m.getType(), dres);
+      gutils->setDiffe(p2m, shadow, builder);
+    }
+  }
+};
 } // namespace
 
 void mlir::enzyme::registerEnzymeXLADialectAutoDiffInterface(
@@ -201,6 +226,7 @@ void mlir::enzyme::registerEnzymeXLADialectAutoDiffInterface(
     registerInterfaces(context);
     GPUWrapperOp::attachInterface<GPUWrapperOpInterfaceReverse>(*context);
     GPUWrapperOp::attachInterface<GPUWrapperOpEnzymeOpsRemover>(*context);
+    enzymexla::Pointer2MemrefOp::attachInterface<Pointer2MemrefRev>(*context);
 
     // Register batching interfaces
     JITCallOp::attachInterface<SHLOGenericBatchOpInterface<JITCallOp>>(

src/enzyme_ad/jax/Implementations/XLADerivatives.h

@@ -10,13 +10,15 @@
 
 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);