[MPS] Add index.Tensor and aten.logical_not

[DenisVieriu97]

Add missing llama ops for MPS delegate:

• index.Tensor
• logical_not

index.put works correctly for generating 1 token, but gives incorrect results on 2nd token. This remains disabled.

Summary of changes:

• Adds missing llama2 ops
• Adds support for launching Metal kernels instead of MPSGraph ops (if MPSGraph doesn't have the support)

cc @cccclai , @shoumikhin

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Unrelated Failure

As of commit 2b833bd with merge base 1eaed2b ():

NEW FAILURE - The following job has failed:

• Apple / test-demo-ios / macos-job (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 65

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Apple / upload-frameworks-ios (gh)

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

[cccclai]

Thank you! The logical op is actually from decomposing sdpa - if the sdpa is not ready yet, maybe can use the simpler version #3165 which has fewer ops than the decomposing F.scaled_dot_product_attention

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[DenisVieriu97]

Apple / test-demo-ios / macos-job (pull_request) is failing because it uses the prebuilt mpsdelegate static lib stored in AWS (https://ossci-ios.s3.amazonaws.com/executorch/, "mps_backend": ["sha256": "97db0fd2b458ff4dae3f4e927d417b4ce88ef3bd4114759abe8372a05bac84ad"] which doesn't match anymore with the runtime changes made in this PR. This PR is using the new AOT changes but the older runtime, which are not compatible anymore. I've checked locally, and Apple / test-demo-ios / macos-job passes with the new runtime.
Any ideas how to fix this ? (cc @cccclai , @shoumikhin )

[DenisVieriu97]

Thank you! The logical op is actually from decomposing sdpa - if the sdpa is not ready yet, maybe can use the simpler version #3165 which has fewer ops than the decomposing F.scaled_dot_product_attention

Thank you @cccclai . I'll take a look and create a new PR for those changes

[cccclai]

Thank you! The logical op is actually from decomposing sdpa - if the sdpa is not ready yet, maybe can use the simpler version #3165 which has fewer ops than the decomposing F.scaled_dot_product_attention

Thank you @cccclai . I'll take a look and create a new PR for those changes

oh sorry pointed to a wrong pr...should be this one #3037. Regarding the test failure, probably @shoumikhin knows more. Will ping him tomorrow

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@cccclai merged this pull request in 02a6b66.

[cccclai]

@pytorchbot cherry-pick --onto release/0.2 -c regression

[pytorchbot]

Cherry picking #3221

The cherry pick PR is at #3267 and it is recommended to link a regression cherry pick PR with an issue

Details for Dev Infra team

Raised by workflow job