/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/jit/passes
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/jit/passes/cuda_graph_fuser.h (860B)
#pragma once
#include
#include
#include
#include
namespace torch {
namespace jit {
// Register CudaFuseGraph in custom passes
struct C10_EXPORT RegisterCudaFuseGraph
: public PassManager {
static bool registerPass(bool enabled) {
TORCH_CHECK(
at::globalContext().hasCUDA() && !at::globalContext().hasHIP(),
"Running CUDA fuser is only supported on CUDA builds.");
bool old_flag = PassManager::isRegistered();
if (enabled) {
PassManager::registerPass(fuser::cuda::fuseGraph);
} else {
PassManager::clearPass();
}
return old_flag;
}
static bool isRegistered() {
return PassManager::isRegistered();
}
};
} // namespace jit
} // namespace torch