mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
3d9c419648
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/25475 I got sucked into this rabbit hole when I was trying to understand what I should do with TensorTypeId occurrences in torch/csrc/utils/tensor_new.cpp. I eventually concluded that all of my problems were because Tensor.new_empty was hand implemented and not actually a native function. So I made it a native function. There are a bunch of other new_* functions which should get this treatment, but I'm sending out this PR just to show how it can be done. The general recipe: 1. Implement a concept of TensorOptions merging (TensorOptions::merge_in). This represents the notion of taking a tensor, but "overriding" some of its values with specific overrides. One subtlety here is how devices get merged; see the comments for what our existing behavior is, and how I preserve it. 2. Implement new_empty as a native function, using options merging. 3. Add another special case to Python binding generation to treat new_* similar to *_like (i.e., handle TensorOptions correctly). The logic here is probably wrong, actually; we should codegen TensorOptions correctly no matter what happens, but new_empty follows the same pattern as empty_like so I opted not to touch this code too much. 4. Delete the now defunct manual binding code. 5. Delete manual type annotations that are no longer necessary since we're going through native. I didn't handle memory format correctly here. I don't know if this function should accept memory format; prior memory format patches didn't add support for memory format to new_like. If we had put memory format in TensorOptions this wouldn't have been a question. ghstack-source-id: 89294185 Test Plan: sandcastle & ossci Differential Revision: D17133000 fbshipit-source-id: 00f4e98bd5174f6fd54e8aba2910ea91824771d9
34 lines
974 B
C++
34 lines
974 B
C++
#pragma once
|
|
|
|
#include <c10/core/TensorOptions.h>
|
|
|
|
// cuda_lazy_init() is always compiled, even for CPU-only builds.
|
|
// Thus, it does not live in the cuda/ folder.
|
|
|
|
namespace torch {
|
|
namespace utils {
|
|
|
|
// The INVARIANT is that this function MUST be called before you attempt
|
|
// to get a CUDA Type object from ATen, in any way. Here are some common
|
|
// ways that a Type object may be retrieved:
|
|
//
|
|
// - You call getNonVariableType or getNonVariableTypeOpt
|
|
// - You call toBackend() on a Type
|
|
//
|
|
// It's important to do this correctly, because if you forget to add it
|
|
// you'll get an oblique error message about "Cannot initialize CUDA without
|
|
// ATen_cuda library" if you try to use CUDA functionality from a CPU-only
|
|
// build, which is not good UX.
|
|
//
|
|
void cuda_lazy_init();
|
|
void set_run_yet_variable_to_false();
|
|
|
|
static void maybe_initialize_cuda(const at::TensorOptions& options) {
|
|
if (options.device().is_cuda()) {
|
|
torch::utils::cuda_lazy_init();
|
|
}
|
|
}
|
|
|
|
}
|
|
}
|