pytorch/torch/csrc/jit/codegen/cuda/interface.h
jiej 127c9402d0 Revert "Revert D30752939: [pytorch][PR] nvfuser update" (#65137)
Summary:
This reverts commit 03389dc851.

Attempt again for PR: https://github.com/pytorch/pytorch/issues/63745
Fixes the windows build failure.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/65137

Reviewed By: seemethere, dzhulgakov, heitorschueroff

Differential Revision: D30994556

Pulled By: malfet

fbshipit-source-id: f1925b6c5cc1a1a441a96499667c91e8dfc1b53d
2021-09-22 04:54:51 -07:00

48 lines
1.5 KiB
C++

#pragma once
#include <torch/csrc/WindowsTorchApiMacro.h>
#include <torch/csrc/jit/ir/ir.h>
#include <torch/csrc/jit/runtime/profiling_record.h>
/*
* This file contains APIs for cuda fuser;
*
* We use an empty static struct to hold the function pointers, which are
* registered separately. This is to support cpu-only compilation.
* Registration is done in torch/csrc/jit/codegen/cuda/register_interface.cpp
*/
namespace torch {
namespace jit {
namespace fuser {
namespace cuda {
TORCH_API std::atomic<bool>& getCudaFusionGuardMode();
// dummy struct to allow API registration
struct CudaFuserInterface {
void (*fn_compile_n)(Node*) = nullptr;
void (*fn_run_n_s)(const Node*, Stack&) = nullptr;
void (*fn_fuse_graph)(std::shared_ptr<Graph>&) = nullptr;
bool (*fn_can_fuse_n)(const Node*) = nullptr;
void (*fn_insert_profile_inodes)(ProfilingRecord* pr) = nullptr;
};
// Get interface, this is used by registration and user facing API internally
C10_EXPORT CudaFuserInterface* getFuserInterface();
C10_EXPORT void compileFusionGroup(Node* fusion_node);
C10_EXPORT void runFusionGroup(const Node* fusion_node, Stack& stack);
C10_EXPORT void fuseGraph(std::shared_ptr<Graph>&);
C10_EXPORT bool canFuseNode(const Node* node);
C10_EXPORT void InsertProfileNodesForCUDAFuser(ProfilingRecord* pr);
C10_EXPORT bool complyWith(
const at::Tensor& tensor,
const c10::TensorTypePtr& guard_tensor_type);
} // namespace cuda
} // namespace fuser
} // namespace jit
} // namespace torch