mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
04939a4745
Summary:
More clang tidy cleanups in `torch/csrc`. This time:
1. `hicpp-use-equals-default` recommends `= default` instead of `{}` for constructors/destructors. This is better practice because it expresses the intent better (https://stackoverflow.com/questions/6502828/what-does-default-mean-after-a-class-function-declaration)
2. `readability-inconsistent-declaration-parameter-name` enforces that parameter names in the declaration match parameter names in the definition. This is just generally useful and can prevent confusion and bugs.
Also updated my script a little bit.
apaszke ezyang
Pull Request resolved: https://github.com/pytorch/pytorch/pull/9737
Differential Revision: D9069069
Pulled By: goldsborough
fbshipit-source-id: f7b3f3a4eb4c9fadc30425a153566d3b613a41ae
112 lines
3.4 KiB
C++
112 lines
3.4 KiB
C++
#pragma once
|
|
#include "ATen/ATen.h"
|
|
|
|
#include "torch/csrc/jit/ivalue.h"
|
|
|
|
namespace torch { namespace jit {
|
|
|
|
using Stack = std::vector<IValue>;
|
|
using Operation = std::function<int(Stack&)>;
|
|
|
|
// An operation with N inputs and M outputs pops the last N inputs off
|
|
// the stack and pushes its M inputs onto the stack
|
|
// before: <other stack items> I0, I1, ... IN <- stack.back()
|
|
// after: <other stack items> O0, O1, ... OM
|
|
// operations are defined this way so that ownership of inputs can be transferred
|
|
// to the operation and it can incrementally drop ownership of tensors
|
|
// when they become unneeded. For large operations, like 'run an entire subgraph',
|
|
// this functionality is very important for minimizing gpu memory usage
|
|
// return value is the relative 'offset' to jump to for the next operation:
|
|
// pc += 1 + offset
|
|
// so a return value of 0 goes to the next instruction
|
|
|
|
// treat the last N elements of the stack as a list, looking up
|
|
// element i
|
|
static inline IValue & peek(Stack & stack, size_t i, size_t N) {
|
|
return *(stack.end() - N + i);
|
|
}
|
|
// treat the last N elements of the stack as a list, looking up the
|
|
// slice starting at index i and having length len
|
|
static inline at::ArrayRef<IValue> peekSlice(const Stack & stack, size_t i, size_t len, size_t N) {
|
|
return at::ArrayRef<IValue>(stack).slice(stack.size() - N + i, len);
|
|
}
|
|
static inline at::ArrayRef<IValue> last(const Stack & stack, size_t N) {
|
|
return peekSlice(stack, 0, N, N);
|
|
}
|
|
static inline void drop(Stack & stack, size_t n) {
|
|
stack.erase(stack.end() - n, stack.end());
|
|
}
|
|
static inline IValue pop(Stack & stack) {
|
|
auto r = std::move(stack.back());
|
|
stack.pop_back();
|
|
return r;
|
|
}
|
|
|
|
// variadic pop:
|
|
// int64_t a; at::Tensor b;
|
|
// pop(stack, a, b);
|
|
// equivalent to:
|
|
// b = pop(stack).toTensor();
|
|
// a = pop(stack).toInt();
|
|
template<typename... Types>
|
|
static inline void pop(Stack& stack, Types&... args) {
|
|
size_t i = 0;
|
|
constexpr size_t N = sizeof...(args);
|
|
int result[N] = {
|
|
(args = std::move(peek(stack,i++, N)).template to<Types>(),0)...
|
|
};
|
|
(void) result;
|
|
drop(stack, N);
|
|
}
|
|
template<typename... Types>
|
|
static inline void push(Stack& stack, Types... args) {
|
|
constexpr size_t N = sizeof...(args);
|
|
int result[N] = {
|
|
(stack.push_back(std::forward<Types>(args)), 0)...
|
|
};
|
|
(void) result;
|
|
}
|
|
|
|
// The packer here is carefully written not to make any unnecessary
|
|
// copies.
|
|
|
|
// pack takes the return values of aten functions pushes them onto the stack
|
|
template<typename T>
|
|
inline void pack(Stack & stack, T&& v) {
|
|
stack.push_back(IValue(std::move(v)));
|
|
}
|
|
|
|
template<>
|
|
inline void pack(Stack & stack, std::vector<at::Tensor>&& v) {
|
|
for(auto& t : v) {
|
|
stack.push_back(IValue(std::move(t)));
|
|
}
|
|
}
|
|
|
|
template<std::size_t remaining, typename... Args>
|
|
struct TuplePacker
|
|
{
|
|
// NB: *Not* a universal reference.
|
|
static void execute(Stack & stack, std::tuple<Args...> && t)
|
|
{
|
|
// NB: The move here does not "destroy" the entire tuple, that is
|
|
// not what std::move does; only the particular tuple index
|
|
// processed here gets stolen.
|
|
pack(stack, std::get<sizeof...(Args) - remaining>(std::move(t)));
|
|
TuplePacker<remaining - 1, Args...>::execute(stack, std::move(t));
|
|
}
|
|
};
|
|
|
|
template<typename... Args>
|
|
struct TuplePacker<0, Args...>
|
|
{
|
|
static void execute(Stack & stack, std::tuple<Args...> && t) {};
|
|
};
|
|
|
|
template<typename... Args>
|
|
inline void pack(Stack & stack, std::tuple<Args...> && t) {
|
|
TuplePacker<sizeof...(Args), Args...>::execute(stack, std::move(t));
|
|
}
|
|
|
|
}}
|