Skip to content

Commit

Permalink
rebase MLIR-Enzyme back on upstream LLVM
Browse files Browse the repository at this point in the history
The relevant changes from the @pengmai fork have been integrated.
  • Loading branch information
ftynse committed Jan 4, 2024
1 parent a6e8dc0 commit bf87475
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 39 deletions.
22 changes: 11 additions & 11 deletions .github/workflows/enzyme-mlir.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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'
Expand All @@ -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

47 changes: 19 additions & 28 deletions enzyme/Enzyme/MLIR/enzymemlir-translate/enzymemlir-translate.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down

0 comments on commit bf87475

Please sign in to comment.