Skip to content
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

Jump to bottom

Minor cleanup #1992

Merged
merged 1 commit into from
Sep 26, 2022
Merged

Minor cleanup #1992

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
14 changes: 13 additions & 1 deletion torch/csrc/jit/codegen/cuda/lower_shift.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -440,8 +440,20 @@ void HaloInfo::build(TensorDomain* td) {
} else {
setHaloWidth(merge->out(), 0);
}
} else if (expr->getExprType().value() == ExprType::Swizzle2D) {
} else if (auto swizzle = dynamic_cast<Swizzle2D*>(expr)) {
// Assume no halo on swizzled domain for now.
TORCH_INTERNAL_ASSERT(
getExtent(swizzle->inX()) == nullptr,
"Halo is not supported with swizzle. Halo-extended ID: ",
swizzle->inX()->toString(),
" used in ",
swizzle->toString());
TORCH_INTERNAL_ASSERT(
getExtent(swizzle->inY()) == nullptr,
"Halo is not supported with swizzle. Halo-extended ID: ",
swizzle->inY()->toString(),
" used in ",
swizzle->toString());
for (auto id : ir_utils::filterByType<IterDomain>(expr->outputs())) {
setHaloWidth(id, 0);
}
Expand Down
4 changes: 2 additions & 2 deletions torch/csrc/jit/codegen/cuda/transform_iter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ void ReplayTransformations::handle(Swizzle2D* swizzle_2d) {
auto id_in_y = swizzle_2d->inY();

// Make sure we have a corresponding entry in our map pointing to the ID we're
// going to replay the split on
// going to replay the swizzle on
auto it_x = id_map_.find(id_in_x);
auto it_y = id_map_.find(id_in_y);

Expand All @@ -162,7 +162,7 @@ void ReplayTransformations::handle(Swizzle2D* swizzle_2d) {
auto outs = std::make_pair(mapped_x, mapped_y);

if (replay_swizzle_) {
// Replay the split onto mapped
// Replay the swizzle onto mapped
outs = IterDomain::swizzle(swizzle_2d->swizzleType(), mapped_x, mapped_y);

// Remove mapped from the leaf IDs
Expand Down