From 6a6a57e844e7f5246c9573da47b1335645833b61 Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Tue, 18 Jul 2023 08:32:52 -0400 Subject: [PATCH 1/4] Add frontend test reproing #603 --- python_tests/test_python_frontend.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/python_tests/test_python_frontend.py b/python_tests/test_python_frontend.py index 60b03813d32..82156e11ca6 100644 --- a/python_tests/test_python_frontend.py +++ b/python_tests/test_python_frontend.py @@ -2422,6 +2422,22 @@ def fusion_func(fd: FusionDefinition, *, deterministic) -> None: print(e) break + # Test expand to zero is replaced with expanded extent and not 1 + # see https://github.com/NVIDIA/Fuser/issues/603 + def test_expand_to_zero(self): + inputs = [ + torch.zeros((1, 0), dtype=torch.float32, device="cuda:0"), + ] + + def fusion_func(fd: FusionDefinition) -> None: + T0 = fd.from_pytorch(inputs[0]) + T1 = fd.ops.broadcast_in_dim(T0, output_shape=[0, 0], broadcast_dims=[0, 1]) + fd.add_output(T1) + + nvf_out, _ = self.exec_nvfuser(fusion_func, inputs) + + self.assertEqual(nvf_out[0].shape, (0, 0)) + if __name__ == "__main__": run_tests() From d491f47f3a7190569b4ea57214f30116822a2879 Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Tue, 18 Jul 2023 08:36:08 -0400 Subject: [PATCH 2/4] Use maybeExpandedExtent always in remove_empty.cpp --- csrc/optimization/remove_empty.cpp | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/csrc/optimization/remove_empty.cpp b/csrc/optimization/remove_empty.cpp index a582984f5de..01b426ca916 100644 --- a/csrc/optimization/remove_empty.cpp +++ b/csrc/optimization/remove_empty.cpp @@ -28,7 +28,8 @@ std::vector emptyAxes(const std::vector& domain) { std::vector empty_axes; for (auto ax : c10::irange(domain.size())) { auto id = domain.at(ax); - if (id->extent()->isConst() && id->extent()->evaluateInt() == 0) { + if (id->getMaybeExpandedExtent()->isConst() && + id->getMaybeExpandedExtent()->evaluateInt() == 0) { empty_axes.push_back((int64_t)ax); } } @@ -112,7 +113,7 @@ class EmptyTensorRemover : public DeadCodeRemover { static std::vector noReductionShape(TensorView* tv) { std::vector shape; for (auto id : TensorDomain::noReductions(tv->getMaybeRFactorDomain())) { - shape.push_back(id->extent()); + shape.push_back(id->getMaybeExpandedExtent()); } return shape; } @@ -252,7 +253,8 @@ class EmptyTensorRemover : public DeadCodeRemover { auto tv = inp->definition()->as()->in()->as(); auto cat_id = TensorDomain::noReductions(tv->getMaybeRFactorDomain()).at(dim); - if (cat_id->extent()->isConst() && cat_id->extent()->evaluateInt() == 0) { + if (cat_id->getMaybeExpandedExtent()->isConst() && + cat_id->getMaybeExpandedExtent()->evaluateInt() == 0) { continue; } non_empty_inputs.push_back(tv); From 5b908a73c8229c4baf12c72dfda0f435d9a30879 Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Tue, 18 Jul 2023 08:51:44 -0400 Subject: [PATCH 3/4] Add double broadcast case to test_expand_to_zero --- python_tests/test_python_frontend.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/python_tests/test_python_frontend.py b/python_tests/test_python_frontend.py index 82156e11ca6..959a83edfe6 100644 --- a/python_tests/test_python_frontend.py +++ b/python_tests/test_python_frontend.py @@ -2426,17 +2426,24 @@ def fusion_func(fd: FusionDefinition, *, deterministic) -> None: # see https://github.com/NVIDIA/Fuser/issues/603 def test_expand_to_zero(self): inputs = [ + # This is an actually empty tensor torch.zeros((1, 0), dtype=torch.float32, device="cuda:0"), + # This one is not actually empty, but should appear to be empty due to expand + torch.zeros((1, 1), dtype=torch.float32, device="cuda:0"), ] def fusion_func(fd: FusionDefinition) -> None: T0 = fd.from_pytorch(inputs[0]) - T1 = fd.ops.broadcast_in_dim(T0, output_shape=[0, 0], broadcast_dims=[0, 1]) - fd.add_output(T1) + T1 = fd.from_pytorch(inputs[1]) + T2 = fd.ops.broadcast_in_dim(T0, output_shape=[0, 0], broadcast_dims=[0, 1]) + T3 = fd.ops.broadcast_in_dim(T1, output_shape=[0, 0], broadcast_dims=[0, 1]) + fd.add_output(T2) + fd.add_output(T3) nvf_out, _ = self.exec_nvfuser(fusion_func, inputs) self.assertEqual(nvf_out[0].shape, (0, 0)) + self.assertEqual(nvf_out[1].shape, (0, 0)) if __name__ == "__main__": From aeaf9102e96dc7bff56456f32803b85af886ac7c Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Tue, 18 Jul 2023 08:51:58 -0400 Subject: [PATCH 4/4] Check getMaybeExpandedExtent at concretization --- csrc/dynamic_transform.cpp | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/csrc/dynamic_transform.cpp b/csrc/dynamic_transform.cpp index 6889e35ebf5..6ce6f592e9d 100644 --- a/csrc/dynamic_transform.cpp +++ b/csrc/dynamic_transform.cpp @@ -115,10 +115,10 @@ class DynamicTransformInitialInfoBuilder : public IterVisitor { // Input and output extent expressions both affect concretization for (const auto& id : TensorDomain::noReductions(inp_tv->getMaybeRFactorDomain())) { - leaf_dynamic_vals_.push_back(id->extent()); + leaf_dynamic_vals_.push_back(id->getMaybeExpandedExtent()); } for (const auto& id : out_tv->getMaybeRFactorDomain()) { - leaf_dynamic_vals_.push_back(id->extent()); + leaf_dynamic_vals_.push_back(id->getMaybeExpandedExtent()); } } } @@ -127,9 +127,10 @@ class DynamicTransformInitialInfoBuilder : public IterVisitor { void handle(TensorView* tv) override { const auto& rfd = tv->getMaybeRFactorDomain(); for (auto id : rfd) { - if (!id->extent()->isConstScalar() || id->extent()->evaluateInt() == 0) { - info_.maybe_zero_extents_set_.insert(id->extent()); - leaf_dynamic_vals_.push_back(id->extent()); + if (!id->getMaybeExpandedExtent()->isConstScalar() || + id->getMaybeExpandedExtent()->evaluateInt() == 0) { + info_.maybe_zero_extents_set_.insert(id->getMaybeExpandedExtent()); + leaf_dynamic_vals_.push_back(id->getMaybeExpandedExtent()); } if (!id->definition() || id->getIterType() != IterType::Symbolic) { continue;