[Fix] fix type imcompatible error in #1115

[kurisu6912]

This pr fix error of #1115

Summary by CodeRabbit

• Bug Fixes

  • Fixed type consistency in modulo operations for shape vectorization.

• Tests

  • Added test coverage for cache kernel behavior on CUDA.

[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]

Walkthrough

This PR includes a type-consistency fix to the FloorMod operation in the packed API transformation, where the second argument is now typed as an IntImm matching the result's dtype, and adds a new test for cache kernel behavior with int64 and int32 addressing on CUDA.

Changes

Cohort / File(s)	Summary
Type-consistent FloorMod correction src/transform/make_packed_api.cc	Changed FloorMod second argument from plain integer constant to IntImm typed with result's dtype, ensuring type-consistent modulo operations during vectorization shape checks
Cache kernel addressing test testing/python/issue/test_tilelang_issue_1115.py	Added new test verifying set_cache_kernel behavior with int64 and int32 address space handling on CUDA, building kernels, executing them, and asserting cache correctness

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~15 minutes

• Verify the IntImm type matches the result's dtype in all execution paths
• Confirm the test adequately covers int64/int32 addressing edge cases and CUDA-specific behavior

Suggested reviewers

• LeiWang1999

Poem

🐰 A type was loose, now snugly bound,
IntImm wraps the modulo round,
And caches tested, int by int,
No addressing bugs shall make us squint!

Pre-merge checks and finishing touches

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check	✅ Passed	The PR title "[Fix] fix type imcompatible error in #1115" is directly related to the main change in the changeset. The primary modification in src/transform/make_packed_api.cc involves replacing the FloorMod second argument with a properly typed IntImm to ensure type-consistent modulo operations, which directly addresses a type incompatibility issue. The addition of the test file further validates this type-related fix. The title is clear and specific enough that it references the issue number and describes the nature of the fix (type incompatibility), allowing teammates to quickly understand the changeset's purpose without needing to examine the full details.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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.}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

testing/python/issue/test_tilelang_issue_1115.py (1)

6-32: Function name doesn't reflect full test scope.

The function test_int64_address actually tests both int64 and int32 addressing modes (see lines 40-45). Consider renaming to test_address_types or test_int64_and_int32_address for clarity.

Apply this diff:

-def test_int64_address():
+def test_address_types():

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 7c61d31 and f7475d1.

📒 Files selected for processing (2)

• src/transform/make_packed_api.cc (1 hunks)
• testing/python/issue/test_tilelang_issue_1115.py (1 hunks)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

• GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
• GitHub Check: Test for Python 3.12 with Metal (on macos-latest)

🔇 Additional comments (2)

src/transform/make_packed_api.cc (1)

436-436: Type-consistent FloorMod operands - fix looks good.

The change correctly creates an IntImm with the same dtype as result to ensure type compatibility in the FloorMod operation. This resolves the type mismatch error where the buffer's default index type (which can be int64 or int32) was not matching the plain integer constant.

testing/python/issue/test_tilelang_issue_1115.py (1)

34-45: Test coverage validates the fix for both address types.

The test correctly verifies that the type compatibility fix works for both int64 and int32 position types, directly addressing the issue mentioned in the PR. The test structure is sound and the assertions confirm correct behavior.