PreserveNVVM.h

Go to the documentation of this file.

//===- PreserveNVVM.h - Mark NVVM attributes for preservation.  -------===//
//
//                             Enzyme Project
//
// Part of the Enzyme 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
//
// If using this code in an academic setting, please cite the following:
// @incollection{enzymeNeurips,
// title = {Instead of Rewriting Foreign Code for Machine Learning,
//          Automatically Synthesize Fast Gradients},
// author = {Moses, William S. and Churavy, Valentin},
// booktitle = {Advances in Neural Information Processing Systems 33},
// year = {2020},
// note = {To appear in},
// }
//
//===----------------------------------------------------------------------===//
//
// This file contains createPreserveNVVM, a transformation pass that marks
// calls to __nv_* functions, marking them as noinline as implementing the llvm
// intrinsic.
//
//===----------------------------------------------------------------------===//
 
#ifndef ENZYME_PRESERVE_NVVM_H
#define ENZYME_PRESERVE_NVVM_H
 
#include "PassUtils.h"
#include "llvm/IR/PassManager.h"
 
namespace llvm {
class ModulePass;
class FunctionPass;
} // namespace llvm
 
llvm::ModulePass *createPreserveNVVMPass(bool Begin);
llvm::FunctionPass *createPreserveNVVMFnPass(bool Begin);
 
class PreserveNVVMNewPM final : public PassParent<PreserveNVVMNewPM> {
  friend PassParent<PreserveNVVMNewPM>;
 
private:
  bool Begin;
  // Whether a function named as one half of a custom derivative rule is kept
  // external and uninlined so the rule can still be resolved later. A pipeline
  // that does not implement custom rules -- the MLIR one -- has nothing to
  // resolve them with, and holding the rule's functions in place only keeps
  // dead code, and the kernels it launches, alive. The globals are consumed
  // either way.
  bool PreserveCustomRuleLinkage;
  // Only meaningful with !Begin: the end of an LTO pre-link pipeline that
  // deferred differentiation to the post-link run. The functions Begin made
  // external get their original linkage back, so the bitcode does not export a
  // strong copy of every static or inline function a rule names, and are kept
  // alive through llvm.compiler.used instead. Their inlining attributes stay as
  // Begin left them; the post-link end pass restores those and drops the
  // llvm.compiler.used entries.
  bool LTOPreLink;
  static llvm::AnalysisKey Key;
 
public:
  using Result = llvm::PreservedAnalyses;
  PreserveNVVMNewPM(bool Begin, bool PreserveCustomRuleLinkage = true,
                    bool LTOPreLink = false)
      : Begin(Begin), PreserveCustomRuleLinkage(PreserveCustomRuleLinkage),
        LTOPreLink(LTOPreLink) {}
 
  Result run(llvm::Module &M, llvm::ModuleAnalysisManager &MAM);
 
  static bool isRequired() { return true; }
};
 
#endif // ENZYME_PRESERVE_NVVM_H