Skip to content

Use expanded extent for empty check in RemoveEmptyPass #604

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
jacobhinkle merged 6 commits into from
Jul 19, 2023
Merged

Use expanded extent for empty check in RemoveEmptyPass #604

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
6a6a57e
Add frontend test reproing #603
jacobhinkle Jul 18, 2023
d491f47
Use maybeExpandedExtent always in remove_empty.cpp
jacobhinkle Jul 18, 2023
5b908a7
Add double broadcast case to test_expand_to_zero
jacobhinkle Jul 18, 2023
aeaf910
Check getMaybeExpandedExtent at concretization
jacobhinkle Jul 18, 2023
14ab252
Merge branch 'main' into zero_expanded_extent
jacobhinkle Jul 18, 2023
289721a
Merge branch 'main' into zero_expanded_extent
jacobhinkle Jul 18, 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
11 changes: 6 additions & 5 deletions csrc/dynamic_transform.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}
}
Expand All @@ -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;
Expand Down
8 changes: 5 additions & 3 deletions csrc/optimization/remove_empty.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@ std::vector<int64_t> emptyAxes(const std::vector<IterDomain*>& domain) {
std::vector<int64_t> 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);
}
}
Expand Down Expand Up @@ -112,7 +113,7 @@ class EmptyTensorRemover : public DeadCodeRemover {
static std::vector<Val*> noReductionShape(TensorView* tv) {
std::vector<Val*> shape;
for (auto id : TensorDomain::noReductions(tv->getMaybeRFactorDomain())) {
shape.push_back(id->extent());
shape.push_back(id->getMaybeExpandedExtent());
}
return shape;
}
Expand Down Expand Up @@ -252,7 +253,8 @@ class EmptyTensorRemover : public DeadCodeRemover {
auto tv = inp->definition()->as<PadOp>()->in()->as<TensorView>();
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);
Expand Down
23 changes: 23 additions & 0 deletions python_tests/test_python_frontend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2422,6 +2422,29 @@ 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 = [
# 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.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__":
run_tests()