1#pragma once
2
3#include <c10/core/TensorOptions.h>
4
5// cuda_lazy_init() is always compiled, even for CPU-only builds.
6// Thus, it does not live in the cuda/ folder.
7
8namespace torch {
9namespace utils {
10
11// The INVARIANT is that this function MUST be called before you attempt
12// to get a CUDA Type object from ATen, in any way. Here are some common
13// ways that a Type object may be retrieved:
14//
15// - You call getNonVariableType or getNonVariableTypeOpt
16// - You call toBackend() on a Type
17//
18// It's important to do this correctly, because if you forget to add it
19// you'll get an oblique error message about "Cannot initialize CUDA without
20// ATen_cuda library" if you try to use CUDA functionality from a CPU-only
21// build, which is not good UX.
22//
23void cuda_lazy_init();
24void set_requires_cuda_init(bool value);
25
26static void maybe_initialize_cuda(const at::TensorOptions& options) {
27 if (options.device().is_cuda()) {
28 torch::utils::cuda_lazy_init();
29 }
30}
31
32} // namespace utils
33} // namespace torch
34