Skip to content

Fix alias analysis #185

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
naoyam merged 3 commits into from
Apr 20, 2023
Merged

Fix alias analysis #185

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e671341
Fix alias analysis
Apr 19, 2023
5b188cc
Do not consider transpose as pointwise
Apr 19, 2023
36b638d
Merge branch 'main' into fix_alias_analysis
Apr 19, 2023
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
9 changes: 9 additions & 0 deletions csrc/ir_utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -572,6 +572,15 @@ bool isReductionTvOp(const Expr* expr) {
return ir_utils::isTvOp(expr) && isReductionOp(expr);
}

bool isPointwiseTvOp(const Expr* expr) {
// LoadStoreOp with rfactor domain means transpose, which is not
// considered pointwise
return isTvOp(expr) &&
(expr->isOneOf<UnaryOp, BinaryOp, TernaryOp>() ||
(expr->isA<LoadStoreOp>() &&
!ir_utils::getTvOutput(expr)->hasRFactor()));
}

std::vector<ViewOp*> getViewOps(Fusion* fusion) {
auto all_exprs = fusion->exprs();

Expand Down
3 changes: 3 additions & 0 deletions csrc/ir_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -338,6 +338,9 @@ TORCH_CUDA_CU_API bool isReductionOp(const Expr*);
// Returns if Expr is a reduction op with TensorView or TensorIndex
TORCH_CUDA_CU_API bool isReductionTvOp(const Expr*);

// Returns if Expr is a pointwise op op with TensorView or TensorIndex
TORCH_CUDA_CU_API bool isPointwiseTvOp(const Expr* expr);

// Returns all non-trivial view operations. We shouldn't have trivial view
// operations but this function is to simply make sure if we ever do we don't
// pull them in.
Expand Down
13 changes: 2 additions & 11 deletions csrc/lower_alias_memory.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1225,7 +1225,8 @@ class ReusableAllocationFinder : private kir::IrVisitor {
if (!tv_def) {
continue;
}
if (!isPointwiseTvOp(tv_def) && !ir_utils::isReductionTvOp(tv_def)) {
if (!ir_utils::isPointwiseTvOp(tv_def) &&
!ir_utils::isReductionTvOp(tv_def)) {
if (isBroadcastTvOp(tv_def)) {
info.has_broadcast_between = true;
} else {
Expand Down Expand Up @@ -1266,16 +1267,6 @@ class ReusableAllocationFinder : private kir::IrVisitor {
}
}

// Do we have a true pointwise op?
// (ie. a TV op, excluding direct assignments and reductions)
bool isPointwiseTvOp(const Expr* expr) {
if (ir_utils::isTvOp(expr)) {
return expr->isA<UnaryOp>() || expr->isA<BinaryOp>() ||
expr->isA<TernaryOp>();
}
return false;
}

// Utility to capture broadcast ops
bool isBroadcastTvOp(const Expr* expr) {
if (!ir_utils::isTvOp(expr)) {
Expand Down