[Refactor] Optimize debug message for parallel inference

[LeiWang1999]

This pull request improves error handling and debugging in the layout inference logic, making it easier to diagnose issues when layout inference fails during parallel operations. The most important changes are:

Error reporting improvements:

• Enhanced the error message in get_unused_iters (in utils.cc) to include details about the operands when divisibility cannot be proven, aiding in debugging layout-related issues.

• Wrapped the layout inference logic in ParallelOpNode::InferLayout (in parallel.cc) with a try-catch block to catch TVM runtime errors. The new error message provides context about the failed buffer, the underlying TVM error, the problematic loop AST, and a hint for resolving the issue, before logging a fatal error.

Now:

for i in T.Parallel(16):
     A_fragment[i, i] = A_fragment[i, i] + 1.0

will throw errs with it's for stmt:

InternalError: Check failed: (CanProveDivisible(splits[lowest]->lower_factor, expected_lower_factor)) is false:  Cannot prove divisible for 2 and 16

  Problematic loop AST:
 for i in T.parallel(16):
    A_fragment = T.Buffer((16, 16), "bfloat16", scope="local.fragment")
    A_fragment[i, i] = T.Cast("bfloat16", T.Cast("float32", A_fragment[i, i]) + T.float32(1.0))
Hint: ensure the loop extent divides the thread binding or adjust the fragment mapping.

Summary by CodeRabbit

• Bug Fixes
  • Improved error reporting with more detailed diagnostic messages when parallel computation operations encounter issues, providing clearer information for troubleshooting.

[coderabbitai]

Walkthrough

Two changes enhance error handling and diagnostics. The first adds user-facing diagnostic messages to an ICHECK assertion in layout utilities. The second wraps thread binding operations with exception handling that logs detailed context before terminating.

Changes

Cohort / File(s)	Summary
Layout utilities diagnostics src/layout/utils.cc	Enhanced ICHECK in get_unused_iters with user-facing error message when divisibility cannot be proven
Parallel operations error handling src/op/parallel.cc	Wrapped Fragment(...)->BindThreadRange call in try-catch to handle tvm::runtime::Error exceptions with detailed diagnostic logging

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

The changes are focused and localized: one is a straightforward assertion message enhancement, and the other adds exception handling with error logging. Both follow clear patterns without introducing complex logic or requiring deep architectural understanding.

Possibly related PRs

• [Bugfix] Recover code for flexible parallel #1032: Modifies compute_loop_layout_from_buffer in src/op/parallel.cc, adding layout-conflict checks and exception throwing in the same function where exception handling is now being enhanced.

Poem

🐰 When threads bind and errors might creep,
A rabbit catches exceptions deep,
With logs so clear, diagnostics bright,
We turn the failures into light! 🌟

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check	✅ Passed	The pull request title "[Refactor] Optimize debug message for parallel inference" is clearly related to the main changes in the changeset. The PR fundamentally improves error handling and diagnostic capabilities in the layout inference logic for parallel operations. The first change adds diagnostic messages to the ICHECK in utils.cc when divisibility cannot be proven, while the second wraps layout inference in parallel.cc with a try-catch that constructs detailed error messages. Both changes directly align with the title's focus on optimizing debug messages. The title is concise, clear, and directly summarizes the primary objective of the changeset without unnecessary details or vague phrasing.

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

[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! 🚀