[export] Add support for symbool to make it usable for torch.cond #138765

henrylhtsang · 2024-10-23T22:54:23Z

Why?

I want the following code to work.

minimal repro:

class M(torch.nn.Module):
    def forward(self, dilate_flag):
        return dilate_flag.item()

input1 = (torch.tensor([1], dtype=torch.bool, device="cuda"),)
model = M().cuda()

ep = torch.export.export(model, input1, strict=True)
path = torch._inductor.aot_compile(ep.module(), input1)
aot_model = torch._export.aot_load(path, device="cuda")
actual_output = aot_model(*input1)

error: AssertionError: Encountered an unsupported object of type <class 'torch.SymBool'> while writing the metadata for exported program

second error will be handled by #138760

Motivation

I could technically bypass it with a torch.int tensor. However, it doesn't work with torch.cond. I want the following to work. It would also require #138760 for aot compile to work.

class M(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.dilate_flag = 0

    def forward(self, dilate_flag):
        self.dilate_flag = dilate_flag.item()

        def true_fn(dilate_flag):
            return dilate_flag.clone()

        def false_fn(dilate_flag):
            return dilate_flag.clone()

        torch.cond(
            self.dilate_flag,
            true_fn,
            false_fn,
            (dilate_flag,),
        )
        return self.dilate_flag

input1 = (torch.tensor([1], dtype=torch.bool, device="cuda"),)
input2 = (torch.tensor([0], dtype=torch.bool, device="cuda"),)
inputs = (input1, input2)
model = M().cuda()

for input in inputs:
    expected_output = model(*input)

    ep = torch.export.export(model, input, strict=False)
    path = torch._inductor.aot_compile(ep.module(), input)
    aot_model = torch._export.aot_load(path, device="cuda")
    actual_output = aot_model(*input)

    assert (
        expected_output == actual_output
    ), f"henry they are not equal {expected_output} != {actual_output}"

Differential Revision: D64867504

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov

pytorch-bot · 2024-10-23T22:54:28Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/138765

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit dbd5906 with merge base 30a83ca ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

facebook-github-bot · 2024-10-23T22:54:31Z

This pull request was exported from Phabricator. Differential Revision: D64867504

Summary: todo Test Plan: ci Differential Revision: D64867504

facebook-github-bot · 2024-10-23T23:25:44Z

This pull request was exported from Phabricator. Differential Revision: D64867504

Summary: expect failure right now need pytorch#138765 pytorch#138760 Differential Revision: D64936442

ydwu4 · 2024-10-24T02:51:44Z

torch/export/graph_signature.py

    elif isinstance(val, SymInt):
        return SymIntArgument(name=node.name)
+    elif isinstance(val, SymBool):
+        return SymIntArgument(name=node.name)


Maye need to add a SymBoolArgument and some test.

cc @angelayi for help on the implementation

Summary: todo Test Plan: ci Differential Revision: D64867504

facebook-github-bot · 2024-10-25T22:28:28Z

This pull request was exported from Phabricator. Differential Revision: D64867504

henrylhtsang · 2024-10-25T22:29:23Z

@ydwu4 Added Symbool argument, added tests in export and aot inductor.

Summary: todo Test Plan: ci Differential Revision: D64867504

facebook-github-bot · 2024-10-28T18:42:34Z

This pull request was exported from Phabricator. Differential Revision: D64867504

henrylhtsang · 2024-10-28T20:06:10Z

@ydwu4 I am getting the following rst error. Any idea?


Testing of coverage in the sources finished, look at the results in build/coverage/python.txt.
/opt/conda/envs/py_3.9/lib/python3.9/tempfile.py:830: ResourceWarning: Implicitly cleaning up <TemporaryDirectory '/tmp/tmpjw7iy0ci'>
  _warnings.warn(warn_message, ResourceWarning)
++ wc -l build/coverage/python.txt
++ cut -f1 '-d '
+ lines=7
+ undocumented=5
+ '[' 5 -lt 0 ']'
+ '[' 5 -gt 0 ']'
+ echo undocumented objects found:
undocumented objects found:
+ cat build/coverage/python.txt
Undocumented Python objects
===========================
torch.export.graph_signature
----------------------------
Classes:
 * SymBoolArgument

+ echo 'Make sure you'\''ve updated relevant .rsts in docs/source!'
Make sure you've updated relevant .rsts in docs/source!
+ echo 'You can reproduce locally by running '\''cd docs && make coverage && cat build/coverage/python.txt'\'''
You can reproduce locally by running 'cd docs && make coverage && cat build/coverage/python.txt'
+ exit 1

henrylhtsang · 2024-10-29T21:14:54Z

@pytorchbot rebase -b viable/strict

pytorchmergebot · 2024-10-29T21:16:20Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2024-10-29T21:16:23Z

Successfully rebased export-D64867504 onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout export-D64867504 && git pull --rebase)

Summary: todo Test Plan: ci Differential Revision: D64867504

ydwu4 · 2024-10-29T22:31:10Z

test/inductor/test_aot_inductor_arrayref.py

    ),
    "test_size_from_multi_output": fail_stack_allocation(is_skip=True),
    "test_torchvision_transforms_functional_tensor_resize": fail_minimal_arrayref_interface(),
+    # TODO: AttributeError: 'ShapeAsConstantBuffer' object has no attribute 'dtype'


Can you paste the full error stack trace? cc @desertfire for review on this change.

stack trace: https://gist.github.com/henrylhtsang/a6ef1a85c56c9fafef1cbc188ca82e67

/torch/_inductor/codegen/cpp_wrapper_cpu_array_ref.py", line 134, in <genexpr> f"ArrayRefTensor<{DTYPE_TO_CPP[x.get_dtype()]}>" /torch/_inductor/ir.py", line 393, in get_dtype return self.dtype AttributeError: 'ShapeAsConstantBuffer' object has no attribute 'dtype'

ydwu4 · 2024-10-29T22:48:26Z

torch/_export/verifier.py


    for input_spec, node in zip(gs.input_specs, input_node_names):
-        if isinstance(input_spec.arg, (TensorArgument, SymIntArgument)):
+        if isinstance(input_spec.arg, (TensorArgument, SymIntArgument, SymBoolArgument)):


When we have SymBoolArgument inputs? Can add a test for it?

@ydwu4 I am not sure, I was just following the pattern of SymIntArgument.

Any SymIntArgument test that I can take a look?

may check test_lifted_constants::Foo

ydwu4 · 2024-10-29T22:55:07Z

the following rst error

Appears to be some doc issue. Not familiar with doc issues. Probably can follow https://github.com/pytorch/pytorch/blob/main/docs/source/export.rst?plain=1#L896 . But not sure why this is not triggered for SymIntArgument. Worth digging deeper.

Summary: todo Test Plan: ci Differential Revision: D64867504

facebook-github-bot · 2024-10-30T21:23:35Z

This pull request was exported from Phabricator. Differential Revision: D64867504

facebook-github-bot · 2024-10-31T23:48:32Z

This pull request was exported from Phabricator. Differential Revision: D64867504

Summary: Pull Request resolved: #138765 todo Test Plan: ci Differential Revision: D64867504

ydwu4

Add another test case for test_lifted_constants of boolean?

Summary: todo Test Plan: ci Differential Revision: D64867504

facebook-github-bot · 2024-11-04T19:36:51Z

This pull request was exported from Phabricator. Differential Revision: D64867504

henrylhtsang · 2024-11-04T19:37:37Z

Add another test case for test_lifted_constants of boolean?

my bad, added now

henrylhtsang · 2024-11-04T23:24:02Z

@pytorchbot merge

pytorchmergebot · 2024-11-04T23:26:09Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

henrylhtsang requested review from angelayi, avikchaudhuri, tugsbayasgalan, ydwu4 and zhxchen17 as code owners October 23, 2024 22:54

pytorch-bot bot added the release notes: export label Oct 23, 2024

facebook-github-bot added the fb-exported label Oct 23, 2024

henrylhtsang added the topic: not user facing topic category label Oct 23, 2024

henrylhtsang changed the title ~~[export] support symbool for torch.cond~~ [export] Add support for symbool to make torch.cond usable Oct 23, 2024

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 23, 2024

henrylhtsang changed the title ~~[export] Add support for symbool to make torch.cond usable~~ [export] Add support for symbool to make it usable for torch.cond Oct 23, 2024

henrylhtsang added a commit to henrylhtsang/pytorch that referenced this pull request Oct 23, 2024

[export] support symbool for torch.cond (pytorch#138765)

94742f0

Summary: todo Test Plan: ci Differential Revision: D64867504

henrylhtsang force-pushed the export-D64867504 branch from ec14551 to 94742f0 Compare October 23, 2024 23:25

This was referenced Oct 24, 2024

torch.cond works, but would still print a misleading warning #138782

Closed

[aoti] fix typo in codegen_dynamic_scalar #138760

Closed

[aoti] Add tests for symint and symbool .item() #138867

Closed

henrylhtsang added a commit to henrylhtsang/pytorch that referenced this pull request Oct 24, 2024

[aoti] Add tests for symint and symbool .item()

5bf3de5

Summary: expect failure right now need pytorch#138765 pytorch#138760 Differential Revision: D64936442

ydwu4 reviewed Oct 25, 2024

View reviewed changes

henrylhtsang force-pushed the export-D64867504 branch from 94742f0 to 4e53e47 Compare October 25, 2024 22:28

henrylhtsang added a commit to henrylhtsang/pytorch that referenced this pull request Oct 25, 2024

[export] support symbool for torch.cond (pytorch#138765)

4e53e47

Summary: todo Test Plan: ci Differential Revision: D64867504

pytorch-bot bot added ciflow/inductor module: inductor labels Oct 25, 2024

henrylhtsang added a commit to henrylhtsang/pytorch that referenced this pull request Oct 28, 2024

[export] support symbool for torch.cond (pytorch#138765)

76b2ecb

Summary: todo Test Plan: ci Differential Revision: D64867504

henrylhtsang force-pushed the export-D64867504 branch from 4e53e47 to 76b2ecb Compare October 28, 2024 18:42

pytorchmergebot force-pushed the export-D64867504 branch from 76b2ecb to b67372e Compare October 29, 2024 21:16

pytorch-bot bot pushed a commit that referenced this pull request Oct 29, 2024

[export] support symbool for torch.cond (#138765)

b67372e

Summary: todo Test Plan: ci Differential Revision: D64867504

ydwu4 reviewed Oct 29, 2024

View reviewed changes

henrylhtsang force-pushed the export-D64867504 branch from b67372e to ea17f08 Compare October 30, 2024 21:23

henrylhtsang added a commit to henrylhtsang/pytorch that referenced this pull request Oct 30, 2024

[export] support symbool for torch.cond (pytorch#138765)

ea17f08

Summary: todo Test Plan: ci Differential Revision: D64867504

henrylhtsang force-pushed the export-D64867504 branch from ea17f08 to fdc9f01 Compare October 31, 2024 23:48

pytorch-bot bot pushed a commit that referenced this pull request Oct 31, 2024

[export] support symbool for torch.cond (#138765)

fdc9f01

Summary: Pull Request resolved: #138765 todo Test Plan: ci Differential Revision: D64867504

ydwu4 approved these changes Nov 4, 2024

View reviewed changes

[export] support symbool for torch.cond (pytorch#138765)

dbd5906

Summary: todo Test Plan: ci Differential Revision: D64867504

henrylhtsang force-pushed the export-D64867504 branch from fdc9f01 to dbd5906 Compare November 4, 2024 19:36

pytorchmergebot added the merging label Nov 4, 2024

pytorchmergebot added the Merged label Nov 4, 2024

pytorchmergebot closed this in 350bc2a Nov 4, 2024

pytorchmergebot removed the merging label Nov 4, 2024

bhack mentioned this pull request Nov 5, 2024

Add truediv support in export serializer #136364

Closed

5 tasks

[export] Add support for symbool to make it usable for torch.cond #138765

[export] Add support for symbool to make it usable for torch.cond #138765

Uh oh!

Conversation

henrylhtsang commented Oct 23, 2024 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why?

Motivation

Uh oh!

pytorch-bot bot commented Oct 23, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/138765

✅ No Failures

Uh oh!

facebook-github-bot commented Oct 23, 2024

Uh oh!

facebook-github-bot commented Oct 23, 2024

Uh oh!

ydwu4 Oct 24, 2024

Choose a reason for hiding this comment

Uh oh!

ydwu4 Oct 25, 2024

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Oct 25, 2024

Uh oh!

henrylhtsang commented Oct 25, 2024

Uh oh!

facebook-github-bot commented Oct 28, 2024

Uh oh!

henrylhtsang commented Oct 28, 2024

Uh oh!

henrylhtsang commented Oct 29, 2024

Uh oh!

pytorchmergebot commented Oct 29, 2024

Uh oh!

pytorchmergebot commented Oct 29, 2024

Uh oh!

ydwu4 Oct 29, 2024

Choose a reason for hiding this comment

Uh oh!

henrylhtsang Oct 30, 2024

Choose a reason for hiding this comment

Uh oh!

ydwu4 Oct 29, 2024

Choose a reason for hiding this comment

Uh oh!

henrylhtsang Oct 29, 2024

Choose a reason for hiding this comment

Uh oh!

ydwu4 Oct 29, 2024

Choose a reason for hiding this comment

Uh oh!

ydwu4 commented Oct 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot commented Oct 30, 2024

Uh oh!

facebook-github-bot commented Oct 31, 2024

Uh oh!

ydwu4 left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Nov 4, 2024

Uh oh!

henrylhtsang commented Nov 4, 2024

Uh oh!

henrylhtsang commented Nov 4, 2024

Uh oh!

pytorchmergebot commented Nov 4, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

henrylhtsang commented Oct 23, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Oct 23, 2024 •

edited

Loading

ydwu4 commented Oct 29, 2024 •

edited

Loading