mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +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
27 lines
520 B
C++
27 lines
520 B
C++
#pragma once
|
|
|
|
#include <torch/csrc/jit/tensorexpr/ir_printer.h>
|
|
|
|
#include <unordered_set>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
namespace tensorexpr {
|
|
|
|
// Generates C++ code from the IR.
|
|
class TORCH_API CppPrinter : public IRPrinter {
|
|
public:
|
|
explicit CppPrinter(std::ostream* os) : IRPrinter(*os) {}
|
|
|
|
using IRPrinter::visit;
|
|
void visit(Allocate*) override;
|
|
void visit(Free*) override;
|
|
|
|
private:
|
|
std::unordered_set<Var*> allocated_on_heap_;
|
|
};
|
|
|
|
} // namespace tensorexpr
|
|
} // namespace jit
|
|
} // namespace torch
|