PreserveNVVM.cpp File Reference

#include <llvm/Config/llvm-config.h>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/TimeProfiler.h"
#include "llvm/Pass.h"
#include "llvm/Transforms/Utils.h"
#include "llvm/Transforms/Utils/ModuleUtils.h"
#include <map>
#include "PreserveNVVM.h"
#include "Utils.h"
#include <llvm-c/Core.h>
#include <llvm-c/Types.h>
#include "llvm/IR/LegacyPassManager.h"

Include dependency graph for PreserveNVVM.cpp:

Go to the source code of this file.

Macros
#define	DEBUG_TYPE   "preserve-nvvm"
#define	ENZYME_ENABLE_NVVM_ATTRIBUTION   1

Functions
void	markFunctionInactive (Function &F)
	Mark F itself as inactive, and additionally mark everything in its body, so that differentiating through the body is a no-op rather than merely calls to F being inactive.
bool	copyForwardedCalleeAttrs (Function &F)
	If F is a trivial forwarding wrapper – one basic block holding a single call whose operands are F's own arguments in order, followed by a return of that call's result – copy the callee's function attributes onto F, except those that only steer inlining or code placement.
bool	preserveLinkage (bool Begin, Function &F, bool Inlining=true)
	Returns whether changed.
static void	handleFunctionLike (bool Begin, Value *Target, StringRef FunctionName)
bool	isTargetNVPTX (llvm::Module &M)
template<const char * handlername, DerivativeMode Mode, int numargs>
static void	handleCustomDerivative (llvm::Module &M, llvm::GlobalVariable &g, SmallVectorImpl< GlobalVariable * > &globalsToErase, bool PreserveCustomRuleLinkage)
bool	preserveNVVM (bool Begin, Module &M, bool PreserveCustomRuleLinkage=true)
static bool	preserveNVVMLTOPreLinkEnd (Module &M)
static RegisterPass< PreserveNVVM >	X ("preserve-nvvm", "Preserve NVVM Pass")
static RegisterPass< PreserveNVVMFn >	XFn ("preserve-nvvm-fn", "Preserve NVVM Pass")
ModulePass *	createPreserveNVVMPass (bool Begin)
FunctionPass *	createPreserveNVVMFnPass (bool Begin)
void	AddPreserveNVVMPass (LLVMPassManagerRef PM, uint8_t Begin)

Macro Definition Documentation

DEBUG_TYPE

#define DEBUG_TYPE   "preserve-nvvm"

Definition at line 58 of file PreserveNVVM.cpp.

ENZYME_ENABLE_NVVM_ATTRIBUTION

#define ENZYME_ENABLE_NVVM_ATTRIBUTION   1

Definition at line 65 of file PreserveNVVM.cpp.

Function Documentation

AddPreserveNVVMPass()

void AddPreserveNVVMPass	(	LLVMPassManagerRef	PM,
		uint8_t	Begin )

Definition at line 1303 of file PreserveNVVM.cpp.

References createPreserveNVVMPass().

copyForwardedCalleeAttrs()

bool copyForwardedCalleeAttrs

(

Function &

F

)

If F is a trivial forwarding wrapper – one basic block holding a single call whose operands are F's own arguments in order, followed by a return of that call's result – copy the callee's function attributes onto F, except those that only steer inlining or code placement.

preserveLinkage marks libdevice wrappers such as __nv_sqrt (whose body is just call @llvm.nvvm.sqrt.rn.d; ret) noinline so Enzyme can recognize them by name. Inlining is how the wrapper would normally pick up the intrinsic's attributes, and FunctionAttrs never infers speculatable, so without this the wrapper stays an opaque call that LICM cannot hoist and SimplifyCFG cannot if-convert. In a loop, that forces Enzyme to cache per-iteration values that would otherwise be loop invariant.

Returns whether F changed.

Definition at line 100 of file PreserveNVVM.cpp.

References startsWith().

Referenced by preserveLinkage().

createPreserveNVVMFnPass()

FunctionPass * createPreserveNVVMFnPass

(

bool

Begin

)

Definition at line 1294 of file PreserveNVVM.cpp.

Referenced by loadNVVMPass().

createPreserveNVVMPass()

ModulePass * createPreserveNVVMPass

(

bool

Begin

)

Definition at line 1290 of file PreserveNVVM.cpp.

Referenced by AddPreserveNVVMPass(), and loadPass().

handleCustomDerivative()

template<const char * handlername, DerivativeMode Mode, int numargs>

static void handleCustomDerivative ( llvm::Module & M, llvm::GlobalVariable & g, SmallVectorImpl< GlobalVariable * > & globalsToErase, bool PreserveCustomRuleLinkage )

static

Definition at line 244 of file PreserveNVVM.cpp.

References ForwardMode, ForwardModeSplit, isTargetNVPTX(), preserveLinkage(), ReverseModeGradient, startsWith(), and str().

Referenced by preserveNVVM().

handleFunctionLike()

static void handleFunctionLike ( bool Begin, Value * Target, StringRef FunctionName )

static

Definition at line 197 of file PreserveNVVM.cpp.

References preserveLinkage().

Referenced by preserveNVVM().

isTargetNVPTX()

bool isTargetNVPTX

(

llvm::Module &

M

)

Definition at line 233 of file PreserveNVVM.cpp.

Referenced by handleCustomDerivative().

markFunctionInactive()

void markFunctionInactive

(

Function &

F

)

Mark F itself as inactive, and additionally mark everything in its body, so that differentiating through the body is a no-op rather than merely calls to F being inactive.

Definition at line 71 of file PreserveNVVM.cpp.

Referenced by preserveNVVM().

preserveLinkage()

bool preserveLinkage	(	bool	Begin,
		Function &	F,
		bool	Inlining = true )

Returns whether changed.

Definition at line 178 of file PreserveNVVM.cpp.

References copyForwardedCalleeAttrs().

Referenced by handleCustomDerivative(), handleFunctionLike(), and preserveNVVM().

preserveNVVM()

bool preserveNVVM	(	bool	Begin,
		Module &	M,
		bool	PreserveCustomRuleLinkage = true )

Definition at line 480 of file PreserveNVVM.cpp.

References attributeKnownFunctions(), Call, ForwardMode, ForwardModeSplit, handleCustomDerivative(), handleFunctionLike(), markFunctionInactive(), preserveLinkage(), ReverseModeGradient, and startsWith().

Referenced by PreserveNVVMNewPM::run().

preserveNVVMLTOPreLinkEnd()

static bool preserveNVVMLTOPreLinkEnd ( Module & M )

static

Definition at line 1230 of file PreserveNVVM.cpp.

Referenced by PreserveNVVMNewPM::run().

X()

static RegisterPass< PreserveNVVM > X ( "preserve-nvvm" , "Preserve NVVM Pass"  )

static

XFn()

static RegisterPass< PreserveNVVMFn > XFn ( "preserve-nvvm-fn" , "Preserve NVVM Pass"  )

static