mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
bdf10380d6
Summary: This takes a lot of pressure off of the C++ typechecker as well as generating much more efficient and smaller code. In my not-super-rigorous testing, compile time for register_prim_ops.cpp went from 68s to 35s, and the size of libtorch went from 72MB to 70MB. Pull Request resolved: https://github.com/pytorch/pytorch/pull/26560 Differential Revision: D17507305 fbshipit-source-id: 8bbd2c08304739432efda96da71f0fa80eb7668b
54 lines
1.5 KiB
C++
54 lines
1.5 KiB
C++
#include <torch/csrc/jit/fuser/fallback.h>
|
|
|
|
#include <ATen/core/functional.h> //fmap
|
|
#include <ATen/core/stack.h>
|
|
#include <torch/csrc/jit/custom_operator.h>
|
|
#include <torch/csrc/jit/fuser/kernel_cache.h>
|
|
#include <torch/csrc/jit/interpreter.h>
|
|
#include <torch/csrc/jit/ir.h>
|
|
|
|
#include <stdexcept>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
namespace fuser {
|
|
|
|
namespace {
|
|
c10::OperatorOptions aliasAnalysisIsSpecialCase() {
|
|
c10::OperatorOptions options;
|
|
options.setAliasAnalysis(AliasAnalysisKind::INTERNAL_SPECIAL_CASE);
|
|
return options;
|
|
}
|
|
} // namespace
|
|
|
|
// Registers fused operators so that fused graphs can properly generate fallback
|
|
// code.
|
|
RegisterOperators reg_fused_operators({Operator(
|
|
prim::FusedConcat,
|
|
[](const Node* node) -> Operation {
|
|
int64_t dim = node->i(attr::dim);
|
|
int64_t num_inputs = node->inputs().size();
|
|
return [dim, num_inputs](Stack& stack) {
|
|
auto result = at::cat(
|
|
fmap(
|
|
last(stack, num_inputs),
|
|
[](const IValue& i) { return i.toTensor(); }),
|
|
dim);
|
|
drop(stack, num_inputs);
|
|
pack(stack, std::move(result));
|
|
return 0;
|
|
};
|
|
},
|
|
aliasAnalysisIsSpecialCase())});
|
|
|
|
void runFallback(int64_t key, Stack& stack) {
|
|
auto maybe_spec = retrieve(key);
|
|
if (!maybe_spec)
|
|
throw std::runtime_error("Failed to find fusion spec to run fallback.");
|
|
|
|
InterpreterState{(*maybe_spec)->code()}.run(stack);
|
|
}
|
|
|
|
} // namespace fuser
|
|
} // namespace jit
|
|
} // namespace torch
|