LowerLLVMExtPass.cpp

Go to the documentation of this file.

//===- LowerLLVMExtPass.cpp - Lower LLVM Ext operations  ------------------ //
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass to lower custom ops generated by the LLVM Ext
// dialect.
//
//===----------------------------------------------------------------------===//
 
#include "Dialect/LLVMExt/LLVMExt.h"
#include "Passes/Passes.h"
 
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
 
namespace mlir {
namespace enzyme {
using namespace mlir::enzyme;
#define GEN_PASS_DEF_LOWERLLVMEXTPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir
 
namespace {
using namespace mlir;
using namespace enzyme;
 
LogicalResult tryLoweringToAlloca(llvm_ext::AllocOp alloc,
                                  uint64_t staticThreshold) {
  llvm::APInt size;
  if (!matchPattern(alloc.getSize(), m_ConstantInt(&size)))
    return failure();
 
  if (size.getZExtValue() > staticThreshold)
    return failure();
 
  Operation *free = nullptr;
  for (auto user : alloc.getResult().getUsers()) {
    auto oFree = dyn_cast<llvm_ext::FreeOp>(user);
    if (!oFree)
      continue;
 
    if (free)
      return failure(); // multiple frees
 
    if (oFree->getBlock() != alloc->getBlock()) // free not in same block
      return failure();
 
    free = user;
 
    if (!alloc->isBeforeInBlock(free))
      return failure();
  }
 
  OpBuilder builder(alloc);
  auto alloca = LLVM::AllocaOp::create(builder, alloc.getLoc(),
                                       alloc.getResult().getType(),
                                       builder.getI8Type(), alloc.getSize());
  alloc.getResult().replaceAllUsesWith(alloca.getResult());
 
  LLVM::LifetimeStartOp::create(builder, alloc.getLoc(), alloca.getResult());
 
  builder.setInsertionPoint(free);
  LLVM::LifetimeEndOp::create(builder, alloc.getLoc(), alloca.getResult());
 
  free->erase();
  alloc->erase();
 
  return success();
}
 
void lowerAlloc(llvm_ext::AllocOp alloc, uint64_t staticThreshold) {
  if (tryLoweringToAlloca(alloc, staticThreshold).succeeded())
    return;
 
  SymbolTable symtable(SymbolTable::getNearestSymbolTable(alloc));
 
  auto mallocFn = symtable.lookup<LLVM::LLVMFuncOp>("malloc");
  if (!mallocFn) {
    Block *b = &symtable.getOp()->getRegion(0).front();
    OpBuilder builder(b, b->begin());
 
    auto fnType = mlir::LLVM::LLVMFunctionType::get(
        LLVM::LLVMPointerType::get(alloc.getContext()), builder.getI64Type(),
        /*isVarArg=*/false);
 
    mallocFn =
        LLVM::LLVMFuncOp::create(builder, alloc.getLoc(), "malloc", fnType);
  }
 
  OpBuilder builder(alloc);
  auto mallocCall = LLVM::CallOp::create(builder, alloc.getLoc(), mallocFn,
                                         alloc->getOperands());
  alloc.getResult().replaceAllUsesWith(mallocCall.getResult());
  alloc.erase();
}
 
void lowerFree(llvm_ext::FreeOp free) {
  SymbolTable symtable(SymbolTable::getNearestSymbolTable(free));
 
  auto freeFn = symtable.lookup<LLVM::LLVMFuncOp>("free");
  if (!freeFn) {
    Block *b = &symtable.getOp()->getRegion(0).front();
    OpBuilder builder(b, b->begin());
 
    auto fnType = mlir::LLVM::LLVMFunctionType::get(
        LLVM::LLVMVoidType::get(free.getContext()),
        LLVM::LLVMPointerType::get(free.getContext()),
        /*isVarArg=*/false);
 
    freeFn = LLVM::LLVMFuncOp::create(builder, free.getLoc(), "free", fnType);
  }
 
  OpBuilder builder(free);
  LLVM::CallOp::create(builder, free.getLoc(), freeFn, free->getOperands());
 
  free.erase();
}
 
struct LowerLLVMExtPass
    : public enzyme::impl::LowerLLVMExtPassBase<LowerLLVMExtPass> {
  using LowerLLVMExtPassBase::LowerLLVMExtPassBase;
 
  void runOnOperation() override {
    Operation *op = getOperation();
 
    op->walk([](llvm_ext::PtrSizeHintOp psh) { psh.erase(); });
 
    SmallVector<llvm_ext::AllocOp> allocs;
    op->walk([&](llvm_ext::AllocOp alloc) { allocs.push_back(alloc); });
 
    for (auto alloc : allocs)
      lowerAlloc(alloc, lowerToAllocaThreshold);
 
    SmallVector<llvm_ext::FreeOp> frees;
    op->walk([&](llvm_ext::FreeOp free) { frees.push_back(free); });
 
    for (auto free : frees)
      lowerFree(free);
  }
};
 
} // end anonymous namespace