Skip to content

Add support for XLA_DISABLE_FUNCTIONALIZATION flag#4792

Merged
wonjoo-wj merged 12 commits intomasterfrom
disable-func
Mar 16, 2023
Merged

Add support for XLA_DISABLE_FUNCTIONALIZATION flag#4792
wonjoo-wj merged 12 commits intomasterfrom
disable-func

Conversation

@wonjoo-wj
Copy link
Copy Markdown
Collaborator

@wonjoo-wj wonjoo-wj commented Mar 16, 2023 •
edited
Loading

Add support for XLA_DISABLE_FUNCTIONALIZATION flag

This change adds support for XLA_DISABLE_FUNCTIONALIZATION. This involves restoring some of the deleted ops when we enabled functionalization and updating some ops to fall back to pre-functionalization state. This change also adds a new test suite that runs test_operations.py with functionalization disabled.

Test plan: source run_tests.sh and CI

Ops that are completely restored:

  • as_strided
  • as_strided_
  • diagonal
  • expand
  • view

Ops that are modified by checking XLA_DISABLE_FUNCTIONALIZATION flag

  • convolution_backward
  • new_empty_strided_symint
  • select_symint
  • slice
  • as_strided

Skipped tests (these two tests were added as part of the funtionalization commit, skipping when XLA_DISABLE_FUNCTIONALIZATION is set to true)

  • test_replace_xla_tensor
  • test_set

Fixed tests

  • test_ailing_slice (slice)
  • test_as_strided (as_strided, as_strided_)
  • test_conv2d_backward (convlution_backward)
  • test_diagonal (diagonal)
  • test_expand (expand)
  • test_flip (view)
  • test_index_bool (select_symint)
  • test_inplace_view (view, new_empty_strided_symint)
  • test_matmul_integer_types (view)
  • test_negative_slice (slice)
  • test_pred_one_hot (select_symint)
  • test_rrelu_module (select_symint)
  • test_slice (slice)
  • test_spooky_ailing (select_symint)
  • test_squeeze_nonzero (new_empty_strided_symint)
  • test_view_1718 (view)
  • test_view_and_copy_ (view)
  • test_view_out_computation (view)
  • TestParallelTensorMNIST.test (convlution_backward)
  • TestParallelTensorResnet18.test (convlution_backward)

cc @alanwaketan

wonjoo-wj added 11 commits March 16, 2023 01:50
@wonjoo-wj
Add support for XLA_DISABLE_FUNCTIONALIZATION flag
bd297d0 
This change adds support for XLA_DISABLE_FUNCTIONALIZATION. This involves restoring some of the deleted ops when we enabled functionalization and updating some ops to fall back to pre-functionalization state. This change also adds a new test suite that runs test_operations.py with functionalization disabled. Also, see note: [Disabling functionalization].
@wonjoo-wj
Restore view op
ee8e6f3
@wonjoo-wj
Restore expand op
1032b2c
@wonjoo-wj
Add CPU fallback for convolution_backward
e14f04a
@wonjoo-wj
Update select_symtint and slice.Tensor ops to support XLA_DISABLE_FUN…
6543e9b 
…CTIONALIZATION case
@wonjoo-wj
Remove duplicate select.int declaration in xla_native_functions.yaml
47153d2
@wonjoo-wj
Update new_empty_strided_symint to fallback properly when functionali…
129a5d0 
…zation is disabled
@wonjoo-wj
Remove debugging print lines
ac7c386
@wonjoo-wj
Run linter
04bb0d8
@wonjoo-wj
Add test_operation test suite with functionalization disabled
37d5498
@wonjoo-wj
Add comments
3faf089
@wonjoo-wj
Copy link
Copy Markdown
Collaborator Author

wonjoo-wj commented Mar 16, 2023 •
edited
Loading

Locally I can see that test_operations.py succeeds when XLA_DISABLE_FUNCTIONALIZATION=1:

----------------------------------------------------------------------
Ran 151 tests in 170.539s

OK (skipped=3)
(base) jenkins@26d7adccbc26:/workspace/pytorch/xla/test$

And also succeeds without the flag:

----------------------------------------------------------------------
Ran 151 tests in 130.983s

OK (skipped=1)
(base) jenkins@26d7adccbc26:/workspace/pytorch/xla/test$

I'll let the CI verify other tests.

@wonjoo-wj wonjoo-wj requested a review from alanwaketan March 16, 2023 02:22
@wonjoo-wj wonjoo-wj self-assigned this Mar 16, 2023
alanwaketan
alanwaketan reviewed Mar 16, 2023
View reviewed changes
Copy link
Copy Markdown
Collaborator

@alanwaketan alanwaketan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Generally LGTM. Please address my comments. Thanks for coming out with this PR so quickly.

Comment thread test/test_operations.py
b = torch.ones([2, 2])
self.runAtenTest((a, b), func)

@unittest.skipIf(
Copy link
Copy Markdown
Collaborator

@alanwaketan alanwaketan Mar 16, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of skipping the tests here, maybe you can utilize the env and have different assertEquals.

Copy link
Copy Markdown
Collaborator Author

@wonjoo-wj wonjoo-wj Mar 16, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These two tests didn't exist before functionalization was enabled (they were newly added as of #4158), should we still even run these when functionalization is disabled?

Copy link
Copy Markdown
Collaborator

@alanwaketan alanwaketan Mar 16, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Forgot the details. Thanks for pointing it out. Then, let's skip.

@wonjoo-wj
Copy link
Copy Markdown
Collaborator Author

wonjoo-wj commented Mar 16, 2023

Thanks for the review, Jiewen. Merging this now as CIs are green.

@wonjoo-wj wonjoo-wj merged commit b44325f into master Mar 16, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@alanwaketan alanwaketan alanwaketan approved these changes

Assignees

@wonjoo-wj wonjoo-wj

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@wonjoo-wj @alanwaketan