Skip to content

[Relax] Legalize shape_to_tensor to device kernel#19957

Open
guan404ming wants to merge 1 commit into
apache:mainfrom
guan404ming:fix-19925-shape-to-tensor-device
Open

[Relax] Legalize shape_to_tensor to device kernel#19957
guan404ming wants to merge 1 commit into
apache:mainfrom
guan404ming:fix-19925-shape-to-tensor-device

Conversation

@guan404ming

@guan404ming guan404ming commented Jul 6, 2026

Copy link
Copy Markdown
Member

Why

Fixes #19925. relax.shape_to_tensor had no legalization, so it always lowered to the host packed func relax.run.shape_to_tensor, producing a CPU tensor regardless of the target device.

How

  • Register a legalization that emits the shape values as a call_tir TE kernel, passing symbolic dims via tir_vars.
  • Fall back to the packed func when the shape values are unknown (ShapeStructInfo without values).

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

@guan404ming guan404ming marked this pull request as ready for review July 6, 2026 12:17
@guan404ming

Copy link
Copy Markdown
Member Author

/gemini review

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

@guan404ming guan404ming force-pushed the fix-19925-shape-to-tensor-device branch from ea19c4d to f34ed21 Compare July 7, 2026 03:02
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.

[Bug] [Relax ONNX frontend] Relax ONNX import can produce CPU shape tensor consumed by CUDA call_tir

1 participant