feat: function argument memory effects for kernels (#1439)

* feat: function argument memory effects for kernels

* chore: run fmt

* feat: use llvm dialect attributes

* feat: merge the arg effects pass into func effects

* fix: correct use of CallOpInterface + only mark llvm attrs for ptrs

* fix: avoid readnone for now (drop me)

Author: avik-pal

src/enzyme_ad/jax/BUILD

@@ -786,6 +786,7 @@ cc_library(
         "@stablehlo//:stablehlo_ops",
         "@stablehlo//:stablehlo_passes",
         "@stablehlo//:stablehlo_type_inference",
+        "@triton//:TritonDialects",
         "@xla//xla/mlir/utils:type_util",
         "@xla//xla/mlir_hlo",
     ],

src/enzyme_ad/jax/Passes/MarkFunctionMemoryEffectsPass.cpp

@@ -5,9 +5,15 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Pass/Pass.h"
+#include "llvm/ADT/BitVector.h"
 
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "src/enzyme_ad/jax/Dialect/Ops.h"
 #include "stablehlo/dialect/StablehloOps.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+#include <queue>
 
 namespace mlir {
 namespace enzyme {
@@ -50,33 +56,40 @@ struct MarkFunctionMemoryEffectsPass
   }
 
   void
-  insertMemoryEffects(SmallVector<uint8_t, 4> &effects,
+  insertMemoryEffects(BitVector &effects,
                       SmallVector<MemoryEffects::EffectInstance> memEffects) {
     for (auto &effect : memEffects)
       insertMemoryEffects(effects, effect);
   }
 
-  void insertMemoryEffects(SmallVector<uint8_t, 4> &effects) {
-    for (int i = 0; i < effects.size(); i++)
-      effects[i] = 1;
+  void insertMemoryEffects(BitVector &effects) {
+    effects.set(0, effects.size());
+  }
+
+  void insertMemoryEffects(BitVector &effects, BitVector &argEffects) {
+    for (int i = 0; i < effects.size(); i++) {
+      if (argEffects[i]) {
+        effects.set(i);
+      }
+    }
   }
 
-  void insertMemoryEffects(SmallVector<uint8_t, 4> &effects,
+  void insertMemoryEffects(BitVector &effects,
                            MemoryEffects::EffectInstance effect) {
     if (effect.getEffect() == MemoryEffects::Read::get()) {
-      effects[0] = 1;
+      effects.set(0);
     } else if (effect.getEffect() == MemoryEffects::Write::get()) {
-      effects[1] = 1;
+      effects.set(1);
     } else if (effect.getEffect() == MemoryEffects::Allocate::get()) {
-      effects[2] = 1;
+      effects.set(2);
     } else if (effect.getEffect() == MemoryEffects::Free::get()) {
-      effects[3] = 1;
+      effects.set(3);
     } else {
       assert(false && "unknown memory effect");
     }
   }
 
-  int64_t getNumEffects(SmallVector<uint8_t, 4> &effects) {
+  int64_t getNumEffects(BitVector &effects) {
     int64_t numEffects = 0;
     for (int i = 0; i < effects.size(); i++) {
       if (effects[i])
@@ -85,12 +98,171 @@ struct MarkFunctionMemoryEffectsPass
     return numEffects;
   }
 
+  struct EffectInfo {
+    ArrayAttr enzymexlaEffects;
+    bool readOnly;
+    bool writeOnly;
+    bool readNone;
+  };
+
+  EffectInfo getEffectInfo(OpBuilder &builder, BitVector &effects) {
+    EffectInfo info;
+    info.readOnly = effects[0];
+    info.writeOnly = effects[1];
+    info.readNone = !effects[0] && !effects[1];
+    SmallVector<Attribute> effectsAttrs;
+
+    if (effects[0]) {
+      info.writeOnly = false;
+      effectsAttrs.push_back(builder.getStringAttr("read"));
+    }
+
+    if (effects[1]) {
+      info.readOnly = false;
+      effectsAttrs.push_back(builder.getStringAttr("write"));
+    }
+
+    if (effects[2]) {
+      info.writeOnly = false;
+      info.readOnly = false;
+      info.readNone = false;
+      effectsAttrs.push_back(builder.getStringAttr("allocate"));
+    }
+
+    if (effects[3]) {
+      info.writeOnly = false;
+      info.readOnly = false;
+      info.readNone = false;
+      effectsAttrs.push_back(builder.getStringAttr("free"));
+    }
+
+    info.enzymexlaEffects = builder.getArrayAttr(effectsAttrs);
+    return info;
+  }
+
+  int32_t getArgIndex(CallOpInterface callOp, OpOperand *operand) {
+    auto callOperands = callOp.getArgOperands();
+    for (unsigned i = 0; i < callOperands.size(); i++) {
+      if (callOperands[i] == operand->get())
+        return i;
+    }
+    assert(false && "operand not found");
+    return -1;
+  }
+
+  // TODO: at some point, we should reuse pre-existing attributes (see
+  // jitcallsideeffect2.mlir)
+  void handleCallOpInterface(
+      CallOpInterface callOp, OpOperand *operand, BitVector &effects,
+      DenseMap<SymbolRefAttr, SmallVector<BitVector>> &funcArgEffects) {
+    if (auto calleeAttr = callOp.getCallableForCallee()) {
+      if (auto symRef = dyn_cast<SymbolRefAttr>(calleeAttr)) {
+        if (funcArgEffects.contains(symRef)) {
+          auto &argEffects = funcArgEffects[symRef];
+          insertMemoryEffects(effects,
+                              argEffects[getArgIndex(callOp, operand)]);
+          return;
+        } else {
+          insertMemoryEffects(effects);
+          return;
+        }
+      }
+    } else {
+      insertMemoryEffects(effects);
+    }
+  }
+
+  bool isPointerType(Value v) { return isPointerType(v.getType()); }
+
+  bool isPointerType(Type t) {
+    return isa<LLVM::LLVMPointerType, MemRefType, triton::PointerType>(t);
+  }
+
+  void analyzeMemoryEffects(
+      Operation *op, OpOperand *operand, BitVector &effects,
+      DenseMap<SymbolRefAttr, SmallVector<BitVector>> &funcArgEffects) {
+    auto memEffectsOrNothing = getEffectsRecursively(op);
+    if (!memEffectsOrNothing.has_value()) {
+      insertMemoryEffects(effects);
+      return;
+    }
+    auto &memEffects = memEffectsOrNothing.value();
+
+    for (const auto &effect : memEffects) {
+      if (effect.getValue() && effect.getValue() == operand->get()) {
+        if (isa<MemoryEffects::Read>(effect.getEffect())) {
+          effects.set(0);
+        } else if (isa<MemoryEffects::Write>(effect.getEffect())) {
+          effects.set(1);
+        } else if (isa<MemoryEffects::Allocate>(effect.getEffect())) {
+          effects.set(2);
+        } else if (isa<MemoryEffects::Free>(effect.getEffect())) {
+          effects.set(3);
+        } else {
+          assert(false && "unknown memory effect");
+        }
+      }
+    }
+  }
+
+  void analyzeFunctionArgumentMemoryEffects(
+      FunctionOpInterface funcOp, SmallVector<BitVector> &argEffects,
+      DenseMap<SymbolRefAttr, SmallVector<BitVector>> &funcArgEffects) {
+    auto *ctx = funcOp->getContext();
+    OpBuilder builder(ctx);
+
+    DenseMap<Value, unsigned> valueToArgIndex;
+    for (unsigned i = 0; i < funcOp.getNumArguments(); i++) {
+      valueToArgIndex[funcOp.getArgument(i)] = i;
+    }
+
+    // BFS traversal starting from arguments
+    std::queue<Value> worklist;
+    DenseSet<Value> visited;
+    for (unsigned i = 0; i < funcOp.getNumArguments(); i++) {
+      Value arg = funcOp.getArgument(i);
+      worklist.push(arg);
+      visited.insert(arg);
+    }
+
+    // BFS through the graph
+    while (!worklist.empty()) {
+      Value cur = worklist.front();
+      worklist.pop();
+
+      auto argIt = valueToArgIndex.find(cur);
+      if (argIt == valueToArgIndex.end())
+        continue;
+      unsigned argIndex = argIt->second;
+
+      for (OpOperand &use : cur.getUses()) {
+        Operation *user = use.getOwner();
+
+        if (auto callOp = dyn_cast<CallOpInterface>(user)) {
+          handleCallOpInterface(callOp, &use, argEffects[argIndex],
+                                funcArgEffects);
+        } else {
+          analyzeMemoryEffects(user, &use, argEffects[argIndex],
+                               funcArgEffects);
+        }
+
+        for (auto result : user->getResults()) {
+          if (visited.insert(result).second) {
+            valueToArgIndex[result] = argIndex;
+            worklist.push(result);
+          }
+        }
+      }
+    }
+  }
+
   void runOnOperation() override {
     ModuleOp module = getOperation();
     auto *ctx = module->getContext();
     OpBuilder builder(ctx);
 
-    DenseMap<SymbolRefAttr, SmallVector<uint8_t, 4>> funcEffects;
+    DenseMap<SymbolRefAttr, BitVector> funcEffects;
+    DenseMap<SymbolRefAttr, SmallVector<BitVector>> funcArgEffects;
 
     CallGraph callGraph(module);
 
@@ -114,7 +286,12 @@ struct MarkFunctionMemoryEffectsPass
       if (!funcOp)
         return signalPassFailure();
 
-      SmallVector<uint8_t, 4> effects(4, 0);
+      BitVector effects(4, 0);
+      SmallVector<BitVector> argEffects;
+      argEffects.reserve(funcOp.getNumArguments());
+      for (unsigned i = 0; i < funcOp.getNumArguments(); i++) {
+        argEffects.push_back(BitVector(4, 0));
+      }
 
       funcOp.walk([&](Operation *op) {
         if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>()) {
@@ -154,12 +331,12 @@ struct MarkFunctionMemoryEffectsPass
         return WalkResult::advance();
       });
 
-      funcEffects[SymbolRefAttr::get(funcOp.getOperation())] =
-          std::move(effects);
+      auto symRef = SymbolRefAttr::get(funcOp.getOperation());
+      funcEffects[symRef] = std::move(effects);
+      funcArgEffects[symRef] = std::move(argEffects);
     }
 
-    auto propagate = [&](FunctionOpInterface funcOp,
-                         SmallVector<uint8_t, 4> &effects) {
+    auto propagate = [&](FunctionOpInterface funcOp, BitVector &effects) {
       funcOp.walk([&](Operation *op) {
         if (auto callOp = dyn_cast<CallOpInterface>(op)) {
           if (auto calleeAttr = callOp.getCallableForCallee()) {
@@ -168,7 +345,7 @@ struct MarkFunctionMemoryEffectsPass
               auto funcEffectsSymRef = funcEffects.lookup(symRef);
               for (int i = 0; i < funcEffectsSymRef.size(); i++) {
                 if (funcEffectsSymRef[i])
-                  effects[i] = 1;
+                  effects.set(i);
               }
             }
           }
@@ -197,8 +374,10 @@ struct MarkFunctionMemoryEffectsPass
           if (!funcOp)
             continue;
 
-          auto &effects =
-              funcEffects[SymbolRefAttr::get(ctx, funcOp.getName())];
+          auto symRef = SymbolRefAttr::get(ctx, funcOp.getName());
+          analyzeFunctionArgumentMemoryEffects(funcOp, funcArgEffects[symRef],
+                                               funcArgEffects);
+          auto &effects = funcEffects[symRef];
           size_t before = getNumEffects(effects);
           propagate(funcOp, effects);
           changed = getNumEffects(effects) != before;
@@ -211,8 +390,7 @@ struct MarkFunctionMemoryEffectsPass
           insertMemoryEffects(effects);
       }
     } else {
-      // No cycles: reverse topological order and propagate
-      for (CallGraphNode *node : llvm::reverse(topoOrder)) {
+      for (CallGraphNode *node : topoOrder) {
         if (node->isExternal())
           continue;
 
@@ -225,7 +403,10 @@ struct MarkFunctionMemoryEffectsPass
         if (!funcOp)
           continue;
 
-        auto &effects = funcEffects[SymbolRefAttr::get(ctx, funcOp.getName())];
+        auto symRef = SymbolRefAttr::get(ctx, funcOp.getName());
+        analyzeFunctionArgumentMemoryEffects(funcOp, funcArgEffects[symRef],
+                                             funcArgEffects);
+        auto &effects = funcEffects[symRef];
         propagate(funcOp, effects);
       }
     }
@@ -237,26 +418,37 @@ struct MarkFunctionMemoryEffectsPass
       if (!funcOp)
         continue;
 
-      SmallVector<Attribute> effectsAttrs;
-      for (int i = 0; i < effectsSet.size(); i++) {
-        if (effectsSet[i]) {
-          if (i == 0) {
-            effectsAttrs.push_back(builder.getStringAttr("read"));
-          } else if (i == 1) {
-            effectsAttrs.push_back(builder.getStringAttr("write"));
-          } else if (i == 2) {
-            effectsAttrs.push_back(builder.getStringAttr("allocate"));
-          } else if (i == 3) {
-            effectsAttrs.push_back(builder.getStringAttr("free"));
-          } else {
-            assert(false && "unknown memory effect");
+      auto funcEffectInfo = getEffectInfo(builder, effectsSet);
+      funcOp->setAttr("enzymexla.memory_effects",
+                      funcEffectInfo.enzymexlaEffects);
+
+      auto &argEffects = funcArgEffects[symbol];
+      for (unsigned i = 0; i < funcOp.getNumArguments(); i++) {
+        auto argEffectInfo = getEffectInfo(builder, argEffects[i]);
+        funcOp.setArgAttr(i, "enzymexla.memory_effects",
+                          argEffectInfo.enzymexlaEffects);
+
+        if (isPointerType(funcOp.getArgument(i))) {
+          if (argEffectInfo.readOnly) {
+            funcOp.setArgAttr(i, LLVM::LLVMDialect::getReadonlyAttrName(),
+                              builder.getUnitAttr());
+          }
+          if (argEffectInfo.writeOnly) {
+            funcOp.setArgAttr(i, LLVM::LLVMDialect::getWriteOnlyAttrName(),
+                              builder.getUnitAttr());
+          }
+          // if (argEffectInfo.readNone) {
+          //   funcOp.setArgAttr(i, LLVM::LLVMDialect::getReadnoneAttrName(),
+          //                     builder.getUnitAttr());
+          // }
+          if (!argEffects[i][3]) {
+            funcOp.setArgAttr(i, LLVM::LLVMDialect::getNoFreeAttrName(),
+                              builder.getUnitAttr());
           }
         }
       }
-
-      funcOp->setAttr("enzymexla.memory_effects",
-                      builder.getArrayAttr(effectsAttrs));
     }
   }
 };
+
 } // namespace

src/enzyme_ad/jax/Passes/Passes.td

@@ -63,6 +63,11 @@ def PropagateConstantBoundsPass
 
 def MarkFunctionMemoryEffectsPass : Pass<"mark-func-memory-effects", "ModuleOp"> {
   let summary = "Attach enzymexla.memory_effects attribute summarizing memory access";
+  let dependentDialects = [
+    "mlir::LLVM::LLVMDialect",
+    "triton::TritonDialect",
+    "memref::MemRefDialect",
+  ];
   let options = [
     Option<
       /*C++ variable name=*/"max_iterations",