fix(codegen): remove unsafe name-based extra_tile_buf_types_ fallback in GetExprTypeAnnotation

[zhangqi-chen]

Summary

Remove the var_name_to_mlir_ → extra_tile_buf_types_ name-based fallback in GetExprTypeAnnotation, which produced incorrect type annotations when multiple variables shared the same name_hint_.

Bug Analysis

After the auto-naming refactor (#619), the ConvertToSSA pass generates variables with a base__qualifier_vN naming scheme. When a program has both a tile.reshape result and a tile.row_expand_mul result derived from the same source variable, they can both receive name_hint_ = "t__tile" despite having completely different shapes ([1, 64] vs [64, 64]).

The lookup chain in GetExprTypeAnnotation was:

1. var_to_mlir_[ptr] → extra_tile_buf_types_[mlir_name] (pointer-based, correct)
2. var_name_to_mlir_[name_hint_] → extra_tile_buf_types_[mlir_name] (name-based, BUGGY)
3. var_to_memref_[ptr] → GetTileBufTypeString(memref) (pointer-based, correct)
4. var_name_to_memref_[name_hint_] → GetTileBufTypeString(memref) (name-based, safe — MemRef carries its own shape)

When the row_expand_mul result t__tile (shape [64, 64]) was looked up:

• Step 1 missed (different Var pointer, not in var_to_mlir_)
• Step 2 hit: var_name_to_mlir_["t__tile"] returned %reshape_buf, and extra_tile_buf_types_["%reshape_buf"] returned the reshape output type [1, 64]

This caused tcolexpandmul to be emitted with a [1, 64] type annotation for its [64, 64] input, which ptoas rejected as a type mismatch.

The fix removes step 2. The remaining fallback paths (steps 3 and 4) are sufficient and correct — they resolve through MemRef which carries its own shape information.

Testing

• All 2577 unit tests pass
• examples/rms_norm.py --sim passes (previously failed with type mismatch)
• examples/softmax.py --sim passes

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical bug in the code generation process by removing an unsafe name-based fallback mechanism that led to incorrect type annotations. The problematic fallback caused issues when different variables shared the same name_hint_, resulting in type mismatches and compilation failures. By eliminating this faulty logic, the system now relies on more robust pointer-based and MemRef-based lookups, ensuring accurate type propagation and improving the overall stability of the code generation.

Highlights

• Removed unsafe type annotation fallback: Eliminated the var_name_to_mlir_ → extra_tile_buf_types_ name-based fallback in GetExprTypeAnnotation to prevent incorrect type annotations when multiple variables shared the same name_hint_.
• Resolved type mismatch bug: Fixed a bug where tcolexpandmul could receive an incorrect type annotation (e.g., [1, 64] instead of [64, 64]), leading to compilation failures.
• Improved type lookup robustness: Ensured that type annotation relies solely on robust pointer-based and MemRef-based lookups, which inherently carry shape information, preventing future name-collision issues.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: e70d88e0-b435-4594-8a89-36bff6b73b43

📥 Commits

Reviewing files that changed from the base of the PR and between e7271b2 and da1185e.

📒 Files selected for processing (1)

• src/codegen/pto/pto_codegen.cpp

💤 Files with no reviewable changes (1)

• src/codegen/pto/pto_codegen.cpp

---

📝 Walkthrough

Walkthrough

A cleanup to PTOCodegen::GetExprTypeAnnotation that removes a name-based compatibility fallback mechanism. The function previously checked var_name_to_mlir_ using name_hint_ as a key; this code path is now removed, and tile-buffer type annotations are derived exclusively from the remaining three resolution mechanisms.

Changes

Cohort / File(s)	Summary
PTOCodegen Type Annotation Simplification src/codegen/pto/pto_codegen.cpp	Removed name-based compatibility fallback in GetExprTypeAnnotation that queried var_name_to_mlir_ via name_hint_ and returned types from extra_tile_buf_types_; type resolution now exclusively uses var_to_mlir_, var_to_memref_, and var_name_to_memref_ paths.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Possibly related PRs

• fix(codegen): Usa SSA format for pto.treshape and infer tile blayout #384: Both PRs modify tile-buffer type resolution in PTOCodegen, specifically the extra_tile_buf_types_ and GetExprTypeAnnotation mechanisms.
• feat(codegen): Add PTOAS-aware type annotations and update PTO naming #267: Both PRs alter GetExprTypeAnnotation and tile-buffer/memref type-resolution logic with implications for how tile types are derived.

Suggested reviewers

• Hzfengsy
• lyfne123

Poem

🐰 A fallback path once stood so tall,
But cleaner code heeds the call,
Three paths remain, tried and true,
The name-based way bid its adieu! ✨

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The pull request title clearly and specifically summarizes the main change: removing an unsafe name-based fallback in GetExprTypeAnnotation.
Description check	✅ Passed	The pull request description provides comprehensive context about the bug, the root cause of the issue, the fix applied, and test results confirming the fix works.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

📝 Coding Plan

• Generate coding plan for human review comments

---

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 and usage tips.}

[gemini-code-assist]

Code Review

This pull request addresses a bug in the code generator's type annotation logic. It removes an unsafe name-based fallback in GetExprTypeAnnotation that could lead to incorrect type lookups when multiple variables shared the same name_hint_. The change is a direct removal of the problematic code block. The analysis in the pull request description accurately details the cause of the bug, and the fix appears correct and sufficient. No issues were found with this change.