Skip to content

[TensorDescriptor] Fallback to TritonRewriteTensorDescriptorToPointer #4367

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.

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
merged 1 commit into from
May 30, 2025

Conversation

whitneywhtsang
Copy link
Contributor

@whitneywhtsang whitneywhtsang commented May 29, 2025

This PR adds fallback logic for tensor descriptor gather and scatter in Triton, integrates the upstream rewrite pass in the pipeline, and updates tests accordingly.

Fixes #4267

@whitneywhtsang whitneywhtsang requested a review from Copilot May 29, 2025 20:53
@whitneywhtsang whitneywhtsang self-assigned this May 29, 2025
Copy link

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull Request Overview

This PR adds fallback logic for tensor descriptor rewrites in Triton, integrates a new rewrite pass in the Python compiler pipeline, and updates tests accordingly.

  • Detect and skip modules containing unsupported descriptor_gather/descriptor_scatter ops before converting descriptors.
  • Invoke rewrite_tensor_descriptor_to_pointer in compiler.py.
  • Add an MLIR test to ensure descriptors aren’t converted when used by descriptor_gather.
  • Remove XPU-specific skips in Python unit tests.

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File Description
third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp Add unsupported-op prewalk and skip logic for fallback case
third_party/intel/backend/compiler.py Register the new rewrite pass in the Triton compiler pipeline
test/Triton/Intel/TensorDescToBlockPointer/invalid.mlir New invalid.mlir to verify descriptor not rewritten on gather
python/test/unit/language/test_tensor_descriptor.py Removed XPU skips in several test cases
Comments suppressed due to low confidence (2)

test/Triton/Intel/TensorDescToBlockPointer/invalid.mlir:3

  • Consider adding a parallel test for descriptor_scatter to verify the fallback logic handles that unsupported op similarly.
// COM: Test make_tensor_descriptor is not rewritten when it is used by descriptor_gather.

python/test/unit/language/test_tensor_descriptor.py:1332

  • Removing the XPU-specific skip may cause failures on XPU devices; either update these tests to support XPU or reintroduce a conditional skip based on actual unsupported features.
if is_xpu():

@whitneywhtsang whitneywhtsang force-pushed the whitneywhtsang/wi4267 branch from 38fe265 to 4dcfe32 Compare May 29, 2025 20:58
@whitneywhtsang whitneywhtsang requested review from alexbaden, etiotto, chengjunlu and a team May 29, 2025 21:01
@etiotto etiotto merged commit 3de88ac into main May 30, 2025
36 of 38 checks passed
@etiotto etiotto deleted the whitneywhtsang/wi4267 branch May 30, 2025 14:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[TEST] Fix test_tensor_descriptor.py::test_tma_[gather|scatter|gather_dot_pipeline]
4 participants