[ROCm] Scalarize vectorizable calls in vector Select lowering

[morluto]

Follow-up to #2889.

Problem

When a fixed-length vector Select branch contains a vectorizable call, the HIP lane scalarizer falls back to extracting a lane from the original vector expression. For a four-lane result, this can emit the complete vector load and operation four times, then select one different lane from each temporary.

For example, the base source generated for a fixed-size four-element bitwise input contains four separate *(int4*)(source + 0) loads, one before each lane assignment.

Change

Use the existing TVectorizable operation attribute to rebuild vectorizable calls from their scalarized arguments. This is the same operation contract used by the TIR vectorizer, and avoids maintaining a separate HIP-specific list. Each Select assignment then computes only the requested lane instead of rematerializing the complete vector expression.

Target-specific intrinsic lowering still runs before HIP source generation. For example, supported popcount operations become __popc or __popcll pure-extern calls and continue through the existing pure-extern scalarization path.

This does not add a lazy-evaluation guarantee to TIR Select. The regression uses fixed-size source and mask buffers, so both branches are valid for every lane regardless of the condition.

Regression coverage

The parameterized source-generation test lowers bitwise_and, bitwise_or, bitwise_xor, bitwise_not, shift_left, shift_right, and unsigned popcount through the full HIP pipeline. On the affected base the bitwise and shift expressions rematerialize packed source vectors. With this change, every generated assignment computes its corresponding scalar source lane and no packed source load is materialized. The popcount case also verifies the established uint32 to __popc intrinsic-lowering path.

Validation

• cmake --build build -j$(nproc) passed.
• python -S -m pytest -q testing/python/amd/test_tilelang_hip_vector_select_codegen.py passed: 36 passed, 12 hardware-gated tests skipped.
• pre-commit run --files src/rocm/codegen/codegen_hip.cc testing/python/amd/test_tilelang_hip_vector_select_codegen.py passed.

Summary

• Fixed HIP vector Select lowering for fixed-length vector branches with vectorizable calls.
• Rebuilds vectorizable calls from lane-scalarized arguments.
• Prevents repeated rematerialization of full vector loads and operations.
• Preserves target-specific intrinsic lowering.
• Added regression coverage for bitwise and shift operations across the HIP pipeline.
• Tests verify per-lane computation and reject packed int4 source loads.

Validation

• ROCm build passed.
• HIP vector Select tests passed: 35 passed, 12 hardware-gated skips.
• Pre-commit checks passed.
• Supported uint32 and uint64 popcount lowering remains handled by hip.FLowerIntrinsic.

C++ style / lint notes

• The PR changes C++ code but does not change rules documented in docs/developer_guide/cpp_style.md.
• The C++ API Style Audit is warning-only. No correctness or build issue is reported.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

HIP vector scalarization now rebuilds vectorizable calls with scalar element return types. New HIP code-generation tests cover masked four-lane int32 and uint32 selection for bitwise operations, shifts, and popcount, and reject packed vector loads.

Changes

HIP vector scalarization

Layer / File(s)	Summary
Scalarize vector integer operations src/rocm/codegen/codegen_hip.cc	VectorLaneScalarizer reads TVectorizable attributes and rebuilds vector calls with scalar element return types and lane-scalarized arguments.
Validate masked vector selection testing/python/amd/test_tilelang_hip_vector_select_codegen.py	Parameterized tests lower masked four-lane selections for bitwise operations, shifts, and popcount across int32 and uint32. Tests verify scalar lane assignments and reject packed int4 and uint4 loads.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related issues

• [ROCm] Vector Select does not scalarize reinterpret before lane extraction #2911 — Both changes update VectorLaneScalarizer for lane-wise HIP vector expression handling.

Possibly related PRs

• tile-ai/tilelang#2889 — This change extends the VectorLaneScalarizer path to rebuild scalarized TVectorizable calls.
• tile-ai/tilelang#2843 — Both changes implement lane-wise scalarization of vectorized operations in GPU code generation.

Sequence Diagram(s)

sequenceDiagram
  participant HIPVectorSelectTest
  participant TileLangBuilder
  participant VectorLaneScalarizer
  participant HIPKernelSource
  HIPVectorSelectTest->>TileLangBuilder: Build masked four-lane integer selection
  TileLangBuilder->>VectorLaneScalarizer: Lower vectorizable integer calls
  VectorLaneScalarizer->>HIPKernelSource: Emit lane-specific scalar ternary assignments
  HIPKernelSource-->>HIPVectorSelectTest: Return generated HIP source
  HIPVectorSelectTest->>HIPKernelSource: Check scalar operators and reject packed vector loads

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: scalarizing vectorizable calls during vector Select lowering on ROCm.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

testing/python/amd/test_tilelang_hip_vector_select_codegen.py (1)

230-239: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Cover every newly scalarized builtin.

This test exercises only bitwise_and. A regression in bitwise_or, bitwise_xor, bitwise_not, shift_left, shift_right, or popcount would pass. Add parameterized expressions or focused source checks for each operation.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@testing/python/amd/test_tilelang_hip_vector_select_codegen.py` around lines
230 - 239, The test test_bitwise_call_keeps_masked_load_inside_select_branch
currently covers only bitwise_and; extend it with parameterized or focused cases
for bitwise_or, bitwise_xor, bitwise_not, shift_left, shift_right, and popcount,
verifying each scalarized builtin keeps masked loads inside the select branch
and avoids vectorized source loads.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/rocm/codegen/codegen_hip.cc`:
- Around line 192-203: The HIP codegen change around the vectorizable operation
handling must not make generic TIR Select semantics lazy or depend on C++
ternary short-circuiting for masked loads. Update
src/rocm/codegen/codegen_hip.cc:192-203 to use an explicitly documented lazy
conditional/backend contract, or ensure both branches are valid for every lane;
update testing/python/amd/test_tilelang_hip_vector_select_codegen.py:199
accordingly so the guarded masked_bitwise_select test does not rely on
out-of-range branch suppression.

---

Nitpick comments:
In `@testing/python/amd/test_tilelang_hip_vector_select_codegen.py`:
- Around line 230-239: The test
test_bitwise_call_keeps_masked_load_inside_select_branch currently covers only
bitwise_and; extend it with parameterized or focused cases for bitwise_or,
bitwise_xor, bitwise_not, shift_left, shift_right, and popcount, verifying each
scalarized builtin keeps masked loads inside the select branch and avoids
vectorized source loads.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: fd693d63-abd6-40cd-badd-d67790742b8e

📥 Commits

Reviewing files that changed from the base of the PR and between c7fabc4 and c99a742.

📒 Files selected for processing (2)

• src/rocm/codegen/codegen_hip.cc
• testing/python/amd/test_tilelang_hip_vector_select_codegen.py

[morluto]

Addressed the remaining builtin-coverage note in c4f265ff.

The full-pipeline regression is now parameterized across bitwise_and, bitwise_or, bitwise_xor, bitwise_not, shift_left, and shift_right, with per-lane source checks and rejection of packed int4 source loads.

Trying to add the requested raw popcount case exposed a separate flaw in the original allowlist: after late scalarization it reaches HIP source codegen as unresolved tirx.popcount and aborts. Supported uint32 and uint64 popcount is already lowered earlier by hip.FLowerIntrinsic to __popc or __popcll, which the existing pure-extern path scalarizes. I removed the redundant and invalid late popcount entry instead of adding code that bypasses that established lowering contract.

The exact-tree ROCm build passes, the HIP vector Select suite is now 35 passed with 12 hardware-gated skips, and pre-commit passes on both changed files.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

testing/python/amd/test_tilelang_hip_vector_select_codegen.py (1)

243-247: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Do not depend on generated assignment order.

enumerate(assignments) assumes that HIP source emits lane 0 first and lane 3 last. A valid backend reordering would fail this test. Check exactly one assignment for each lane instead.

Suggested assertion

-    for lane, assignment in enumerate(assignments):
-        assert f"source[{lane}]" in assignment
-        assert source_pattern in assignment
+    for lane in range(4):
+        lane_assignments = [assignment for assignment in assignments if f"source[{lane}]" in assignment]
+        assert len(lane_assignments) == 1
+        assert source_pattern in lane_assignments[0]

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@testing/python/amd/test_tilelang_hip_vector_select_codegen.py` around lines
243 - 247, Update the assignment assertions in the test to avoid
enumerate(assignments) and generated ordering assumptions. Validate that exactly
one assignment exists for each lane 0 through 3, and apply the existing
source_pattern check to each lane-specific assignment regardless of their order.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@testing/python/amd/test_tilelang_hip_vector_select_codegen.py`:
- Around line 229-237: Extend the parameterized cases in
test_integer_call_is_scalarized_per_select_lane with T.popcount, using the
_popcount source pattern and popcount ID. Ensure the test asserts generated
select branches contain scalar source[0] through source[3] calls for this
operation.

---

Nitpick comments:
In `@testing/python/amd/test_tilelang_hip_vector_select_codegen.py`:
- Around line 243-247: Update the assignment assertions in the test to avoid
enumerate(assignments) and generated ordering assumptions. Validate that exactly
one assignment exists for each lane 0 through 3, and apply the existing
source_pattern check to each lane-specific assignment regardless of their order.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 4c19a2ba-7ecb-4eb3-bd72-15e8804300b8

📥 Commits

Reviewing files that changed from the base of the PR and between 3ea62e5 and c4f265f.

📒 Files selected for processing (2)

• src/rocm/codegen/codegen_hip.cc
• testing/python/amd/test_tilelang_hip_vector_select_codegen.py

🚧 Files skipped from review as they are similar to previous changes (1)

• src/rocm/codegen/codegen_hip.cc

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.