From bf874750182a28e56b186636e4994da75e496ac7 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko@google.com>
Date: Thu, 4 Jan 2024 10:21:38 +0000
Subject: [PATCH] rebase MLIR-Enzyme back on upstream LLVM

The relevant changes from the @pengmai fork have been integrated.
---
 .github/workflows/enzyme-mlir.yml             | 22 ++++-----
 .../enzymemlir-translate.cpp                  | 47 ++++++++-----------
 2 files changed, 30 insertions(+), 39 deletions(-)

diff --git a/.github/workflows/enzyme-mlir.yml b/.github/workflows/enzyme-mlir.yml
index 834efa465619..ecf909bb5b32 100644
--- a/.github/workflows/enzyme-mlir.yml
+++ b/.github/workflows/enzyme-mlir.yml
@@ -13,30 +13,30 @@ jobs:
   build-linux:
     name: MLIR ${{ matrix.build }} ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
-    
+
     strategy:
       fail-fast: false
       matrix:
         build: ["Release", "Debug"] # "RelWithDebInfo"
         llbuild: ["Release"]
         os: [openstack22]
-  
+
     timeout-minutes: 500
 
     steps:
-    - name: Install dependencies 
+    - name: Install dependencies
       run: |
           sudo apt-get update
-          sudo apt-get install -y binutils ninja-build cmake gcc g++ python3 python3-dev 
-    
+          sudo apt-get install -y binutils ninja-build cmake gcc g++ python3 python3-dev
+
     - uses: actions/checkout@v3
       with:
         path: 'Enzyme'
 
     - uses: actions/checkout@v3
       with:
-        repository: 'pengmai/llvm-project'
-        ref: '205fb43c57fbf2307bf7ec17f07da2324250d26a'
+        repository: 'llvm/llvm-project'
+        ref: '5ed11e767c0c39a3bc8e035588e7a383849d46a8'
         path: 'llvm-project'
 
     - name: Get MLIR commit hash
@@ -50,7 +50,7 @@ jobs:
       with:
         path: llvm-project/mlir-build
         key: ${{ matrix.llbuild }}-${{ matrix.os }}-mlir-${{ steps.mlir-commit.outputs.sha_short }}
-    
+
     - name: MLIR build
       if: steps.cache-mlir.outputs.cache-hit != 'true'
       working-directory: 'llvm-project'
@@ -66,11 +66,11 @@ jobs:
     - name: Enzyme build
       working-directory: 'Enzyme'
       run: |
-          mkdir enzyme-build && cd enzyme-build 
+          mkdir enzyme-build && cd enzyme-build
           cmake ../enzyme -GNinja -DCMAKE_BUILD_TYPE=${{ matrix.build }} -DLLVM_DIR=${{ github.workspace }}/llvm-project/mlir-build -DENZYME_MLIR=ON
           ninja
-    
+
     - name: Check enzyme-mlir
       working-directory: 'Enzyme/enzyme-build'
       run: ninja check-enzymemlir
-  
+
diff --git a/enzyme/Enzyme/MLIR/enzymemlir-translate/enzymemlir-translate.cpp b/enzyme/Enzyme/MLIR/enzymemlir-translate/enzymemlir-translate.cpp
index 9d778333724c..0404febe83a4 100644
--- a/enzyme/Enzyme/MLIR/enzymemlir-translate/enzymemlir-translate.cpp
+++ b/enzyme/Enzyme/MLIR/enzymemlir-translate/enzymemlir-translate.cpp
@@ -32,53 +32,44 @@ using namespace llvm;
 class ActivityToMetadataTranslation : public LLVMTranslationDialectInterface {
   using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;
 
-  void annotateActivity(Operation *op, StringRef key,
-                        LLVM::ModuleTranslation &moduleTranslation) const {
-    LLVMContext &llvmCtx = moduleTranslation.getLLVMContext();
-    MDNode *md = MDNode::get(llvmCtx, {});
-
-    if (op->getNumResults() == 1) {
-      llvm::Value *val = moduleTranslation.lookupValue(op->getResult(0));
-      if (auto *inst = dyn_cast<llvm::Instruction>(val)) {
-        inst->setMetadata(key, md);
-      }
-    } else if (op->getNumResults() == 0) {
-      llvm::Instruction *inst = moduleTranslation.lookupOperation(op);
-      if (!inst)
-        return;
+  void annotateActivity(StringRef key,
+                        ArrayRef<llvm::Instruction *> instructions, ) const {
+    if (instructions.empty())
+      return;
 
+    LLVMContext &llvmCtx = instructions.front()->getContext();
+    MDNode *md = MDNode::get(llvmCtx, {});
+    for (llvm::Instruction *inst : instructions) {
       inst->setMetadata(key, md);
     }
   }
 
   LogicalResult
-  amendOperation(Operation *op, NamedAttribute attribute,
+  amendOperation(Operation *op, ArrayRef<llvm::Instruction *> instructions,
+                 NamedAttribute attribute,
                  LLVM::ModuleTranslation &moduleTranslation) const override {
-    if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
-      // auto *fn = moduleTranslation.lookupFunction(funcOp.getName());
-      // errs() << "[translation] fn: " << *fn << "\n";
+    if (auto funcOp = dyn_cast<FunctionOpInterface>(op))
       return success();
-    }
 
     auto iciAttr = op->getAttrOfType<BoolAttr>("enzyme.ici");
     auto icvAttr = op->getAttrOfType<BoolAttr>("enzyme.icv");
 
-    // Op was already processed
+    // Op was already processed.
     if (!(iciAttr && icvAttr))
       return success();
 
-    // Convert the attributes to the appropriate metadata
-    if (iciAttr.getValue() && icvAttr.getValue())
-      annotateActivity(op, "enzyme_inactive", moduleTranslation);
-    else if (!iciAttr.getValue() && !icvAttr.getValue())
-      annotateActivity(op, "enzyme_active", moduleTranslation);
-    else {
+    // Convert the attributes to the appropriate metadata.
+    if (iciAttr.getValue() && icvAttr.getValue()) {
+      annotateActivity("enzyme_inactive", instructions);
+    } else if (!iciAttr.getValue() && !icvAttr.getValue()) {
+      annotateActivity("enzyme_active", instructions);
+    } else {
       StringRef instActivity =
           iciAttr.getValue() ? "enzyme_inactive_inst" : "enzyme_active_inst";
       StringRef valActivity =
           icvAttr.getValue() ? "enzyme_inactive_val" : "enzyme_active_val";
-      annotateActivity(op, instActivity, moduleTranslation);
-      annotateActivity(op, valActivity, moduleTranslation);
+      annotateActivity(instActivity, instructions);
+      annotateActivity(valActivity, instructions);
     }
 
     op->removeAttr("enzyme.ici");