fix: mark memory effects for nested modules correctly

Author: avik-pal

src/enzyme_ad/jax/Passes/MarkFunctionMemoryEffectsPass.cpp

@@ -256,13 +256,34 @@ struct MarkFunctionMemoryEffectsPass
     }
   }
 
+  void collectAllFunctions(
+      Operation *op,
+      DenseMap<SymbolRefAttr, FunctionOpInterface> &symbolToFunc) {
+    if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
+      // Create the symbol reference for this function
+      auto symbolRef = SymbolRefAttr::get(funcOp.getOperation());
+      symbolToFunc[symbolRef] = funcOp;
+    }
+    for (Region &region : op->getRegions()) {
+      for (Block &block : region) {
+        for (Operation &childOp : block) {
+          collectAllFunctions(&childOp, symbolToFunc);
+        }
+      }
+    }
+  }
+
   void runOnOperation() override {
     ModuleOp module = getOperation();
     auto *ctx = module->getContext();
     OpBuilder builder(ctx);
 
     DenseMap<SymbolRefAttr, BitVector> funcEffects;
     DenseMap<SymbolRefAttr, SmallVector<BitVector>> funcArgEffects;
+    DenseMap<SymbolRefAttr, FunctionOpInterface> symbolToFunc;
+
+    // Collect all functions from the module and nested modules
+    collectAllFunctions(module, symbolToFunc);
 
     CallGraph callGraph(module);
 
@@ -413,10 +434,10 @@ struct MarkFunctionMemoryEffectsPass
 
     // Finally, attach attributes
     for (auto &[symbol, effectsSet] : funcEffects) {
-      auto funcOp = dyn_cast_or_null<FunctionOpInterface>(
-          module.lookupSymbol(symbol.getLeafReference()));
-      if (!funcOp)
+      auto it = symbolToFunc.find(symbol);
+      if (it == symbolToFunc.end())
         continue;
+      auto &funcOp = it->second;
 
       auto funcEffectInfo = getEffectInfo(builder, effectsSet);
       funcOp->setAttr("enzymexla.memory_effects",

test/lit_tests/memoryeffects/nestedmodule.mlir

@@ -0,0 +1,22 @@
+// RUN: enzymexlamlir-opt --mark-func-memory-effects %s | FileCheck %s
+
+module {
+    module @nested {
+        func.func public @single_dim(%output: memref<3xi64>, %values: memref<3xi64>) {
+        // CHECK: func.func public @single_dim(%arg0: memref<3xi64> {enzymexla.memory_effects = ["write"], llvm.nofree, llvm.writeonly}, %arg1: memref<3xi64> {enzymexla.memory_effects = ["read"], llvm.nofree, llvm.readonly}) attributes {enzymexla.memory_effects = ["read", "write"]}
+            affine.parallel (%i) = (0) to (3) {
+                %val = memref.load %values[%i] : memref<3xi64>
+                affine.store %val, %output[%i] : memref<3xi64>
+            }
+            return
+        }
+    }
+    func.func @main(%output: memref<3xi64>, %values: memref<3xi64>) {
+    // CHECK: func.func @main(%arg0: memref<3xi64> {enzymexla.memory_effects = ["write"], llvm.nofree, llvm.writeonly}, %arg1: memref<3xi64> {enzymexla.memory_effects = ["read"], llvm.nofree, llvm.readonly}) attributes {enzymexla.memory_effects = ["read", "write"]}
+        affine.parallel (%i) = (0) to (3) {
+            %val = memref.load %values[%i] : memref<3xi64>
+            affine.store %val, %output[%i] : memref<3xi64>
+        }
+        return
+    }
+}