Skip to content

Refactor lower_alias_memory.cpp #2170

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Nov 11, 2022
Merged

Refactor lower_alias_memory.cpp #2170

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
0732580
Refactoring of lower_alias_memory
naoyam Nov 5, 2022
dd7d352
Error msg
naoyam Nov 8, 2022
6cd8a9d
WIP: further refactoring using ExprMutaor
naoyam Nov 8, 2022
b0ce9dc
clang-format
naoyam Nov 9, 2022
78eef72
cleanup
naoyam Nov 9, 2022
664c907
format
naoyam Nov 9, 2022
d1fa87a
cleanup
naoyam Nov 9, 2022
0abf303
cleanup
naoyam Nov 10, 2022
ee2b84c
bug fix
naoyam Nov 10, 2022
25b56be
format
naoyam Nov 10, 2022
c27f22f
comments, cleanup
naoyam Nov 10, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions torch/csrc/jit/codegen/cuda/kernel_ir.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -449,12 +449,14 @@ Allocate::Allocate(
Val* buffer,
MemoryType memory_type,
std::vector<Val*> shape,
bool zero_init)
bool zero_init,
const Allocate* alias)
: Expr(passkey, ExprType::Allocate),
buffer_(buffer),
memory_type_(memory_type),
shape_(std::move(shape)),
zero_init_(zero_init) {
zero_init_(zero_init),
alias_(alias) {
TORCH_INTERNAL_ASSERT(
passkey.ir_container_->isA<kir::Kernel>(),
"IR type only valid for Kernel container.");
Expand Down Expand Up @@ -484,6 +486,12 @@ Allocate::Allocate(
size_ = FusionGuard::getCurFusion()->oneVal();
}

if (alias_ != nullptr) {
TORCH_INTERNAL_ASSERT(alias_ != this, "Invalid alias");
TORCH_INTERNAL_ASSERT(
alias_->memoryType() == memory_type_, "Invalid alias");
}

addInput(size_);
}

Expand Down
9 changes: 2 additions & 7 deletions torch/csrc/jit/codegen/cuda/kernel_ir.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,8 @@ class TORCH_CUDA_CU_API Allocate final : public Expr {
Val* buffer,
MemoryType memory_type,
std::vector<Val*> shape = {},
bool zero_init = false);
bool zero_init = false,
const Allocate* alias = nullptr);

//! Allocation of a non-dimensional buffer
//!
Expand Down Expand Up @@ -225,12 +226,6 @@ class TORCH_CUDA_CU_API Allocate final : public Expr {
return alias_;
}

void setAlias(const Allocate* alias) {
TORCH_INTERNAL_ASSERT(alias != this);
TORCH_INTERNAL_ASSERT(alias->memoryType() == memory_type_);
alias_ = alias;
}

private:
Val* buffer_ = nullptr;
MemoryType memory_type_ = MemoryType::Local;
Expand Down
10 changes: 10 additions & 0 deletions torch/csrc/jit/codegen/cuda/lower2device.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,6 +396,16 @@ bool GpuLower::hasCurrent() {

void GpuLower::propagateExprInfo(const Expr* old_expr, const Expr* new_expr) {
pred_elimination_.propagateRemovalInfo(old_expr, new_expr);
if (old_expr->isA<kir::Allocate>()) {
auto alloc_info_it =
localAllocationInfoMap().find(old_expr->as<kir::Allocate>());
if (alloc_info_it != localAllocationInfoMap().end()) {
auto alloc_info =
std::make_unique<LocalAllocationInfo>(*(alloc_info_it->second));
localAllocationInfoMap().emplace(
new_expr->as<kir::Allocate>(), std::move(alloc_info));
}
}
}

} // namespace cuda
Expand Down
Loading