mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 13:44:15 +08:00
59dd12042e
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/62336 This PR was generated by removing `const` for all types of nodes in NNC IR, and fixing compilation errors that were the result of this change. This is the first step in making all NNC mutations in-place. Test Plan: Imported from OSS Reviewed By: iramazanli Differential Revision: D30049829 Pulled By: navahgar fbshipit-source-id: ed14e2d2ca0559ffc0b92ac371f405579c85dd63
45 lines
1.2 KiB
C++
45 lines
1.2 KiB
C++
#include <torch/csrc/jit/tensorexpr/cpp_codegen.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
namespace tensorexpr {
|
|
|
|
void CppPrinter::visit(Allocate* alloc) {
|
|
constexpr size_t kAllocOnStackThresholdSize = 512;
|
|
|
|
size_t size = 1;
|
|
for (auto dim : alloc->dims()) {
|
|
IntImm* v = dynamic_cast<IntImm*>(dim);
|
|
if (v) {
|
|
size *= v->value();
|
|
} else {
|
|
throw std::runtime_error("Only IntImm dimensions are supported for now");
|
|
}
|
|
}
|
|
|
|
emitIndent();
|
|
if (size <= kAllocOnStackThresholdSize) {
|
|
os() << alloc->dtype().ToCppString() << " " << (*alloc->buffer_var()) << "["
|
|
<< size << "];" << std::endl;
|
|
} else {
|
|
size *= alloc->dtype().byte_size();
|
|
os() << alloc->dtype().ToCppString() << "* " << (*alloc->buffer_var())
|
|
<< " = static_cast<" << alloc->dtype().ToCppString() << "*>(malloc("
|
|
<< size << "));" << std::endl;
|
|
allocated_on_heap_.insert(alloc->buffer_var());
|
|
}
|
|
}
|
|
|
|
void CppPrinter::visit(Free* free) {
|
|
Var* var = free->buffer_var();
|
|
if (allocated_on_heap_.count(var)) {
|
|
emitIndent();
|
|
os() << "free(" << name_manager()->get_unique_name(var) << ");"
|
|
<< std::endl;
|
|
}
|
|
}
|
|
|
|
} // namespace tensorexpr
|
|
} // namespace jit
|
|
} // namespace torch
|