Fix mutator and sameAs for expanded IterDomain (#1902)

csarofeen · Aug 11, 2022 · 3381793 · 3381793
1 parent b7435af
commit 3381793
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/torch/csrc/jit/codegen/cuda/ir_nodes.cpp b/torch/csrc/jit/codegen/cuda/ir_nodes.cpp
@@ -1137,6 +1137,10 @@ bool IterDomain::sameAs(const Statement* other) const {
  is_same = is_same && ScalarCheck::sameAs(start(), other_id->start());
  is_same =
  is_same && ScalarCheck::sameAs(stopOffset(), other_id->stopOffset());
+ is_same = is_same && (hasExpandedExtent() == other_id->hasExpandedExtent());
+ if (is_same && hasExpandedExtent()) {
+ is_same = ScalarCheck::sameAs(expandedExtent(), other_id->expandedExtent());
+ }
 
  return is_same;
 }

diff --git a/torch/csrc/jit/codegen/cuda/mutator.cpp b/torch/csrc/jit/codegen/cuda/mutator.cpp
@@ -58,8 +58,14 @@ void OptOutMutator::mutate(NamedScalar* ns) {}
 void OptOutMutator::mutate(IterDomain* id) {
  Val* start = maybeMutated(id->start());
  Val* extent = maybeMutated(id->extent());
+ Val* expanded_extent = nullptr;
+ if (id->hasExpandedExtent()) {
+ expanded_extent = maybeMutated(id->expandedExtent());
+ }
  Val* stop_offset = maybeMutated(id->stopOffset());
  if (start->sameAs(id->start()) && extent->sameAs(id->extent()) &&
+ (!id->hasExpandedExtent() ||
+ expanded_extent->sameAs(id->expandedExtent())) &&
  stop_offset->sameAs(id->stopOffset())) {
  return;
  }
@@ -69,6 +75,7 @@ void OptOutMutator::mutate(IterDomain* id) {
  .start(start)
  .extent(extent)
  .stop_offset(stop_offset)
+ .expanded_extent(expanded_extent)
  .build());
 }
 

diff --git a/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp b/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp
@@ -23815,7 +23815,6 @@ TEST_F(NVFuserTest, FusionLoopSwizzleCheck1_CUDA) {
  // Make tv2 swizzled and half-inlined (unsupported).
  tv0->computeAt(tv3, -2);
 
- fusion.print();
  FusionExecutor fe;
  ASSERT_ANY_THROW(fe.compileFusion(&fusion));
 }