mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
ac99d63f83
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/63414 Misuse of raw pointer in here where stack is never nullable. ghstack-source-id: 136938318 Test Plan: compiles. Imported from OSS Reviewed By: ejguan Differential Revision: D30375410 fbshipit-source-id: 9d65b620bb76d90d886c800f54308520095d58ee
52 lines
1.4 KiB
C++
52 lines
1.4 KiB
C++
#include <torch/csrc/jit/codegen/fuser/fallback.h>
|
|
|
|
#include <ATen/core/functional.h> //fmap
|
|
#include <ATen/core/stack.h>
|
|
#include <torch/csrc/jit/codegen/fuser/kernel_cache.h>
|
|
#include <torch/csrc/jit/ir/ir.h>
|
|
#include <torch/csrc/jit/runtime/custom_operator.h>
|
|
#include <torch/csrc/jit/runtime/interpreter.h>
|
|
|
|
#include <stdexcept>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
namespace fuser {
|
|
|
|
namespace {
|
|
c10::AliasAnalysisKind aliasAnalysisIsSpecialCase() {
|
|
return AliasAnalysisKind::INTERNAL_SPECIAL_CASE;
|
|
}
|
|
} // 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));
|
|
};
|
|
},
|
|
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
|