[Refactor]:Move device_assert from extern_call to intrin_call #1134
Conversation
WalkthroughAdds two TL intrinsics ( Changes
Sequence Diagram(s)sequenceDiagram
participant TLPrint as tilelang/print.py
participant IR as TVM IR
participant Codegen as CodeGenTileLangCUDA
participant GPU as CUDA device
alt no message
TLPrint->>IR: T.call_intrin('tl.device_assert', cond)
else with message
TLPrint->>IR: T.call_intrin('tl.device_assert_with_msg', cond, msg)
end
IR->>IR: builtin ops registered (src/op/builtin.cc)
IR->>Codegen: EvaluateNode visits CallNode
Note over Codegen: matches tl.device_assert / tl.device_assert_with_msg
Codegen->>Codegen: PrintDeviceAssert(call) %%# (fill:#f7fbff)
Codegen->>GPU: emit optional printf(msg) and assert(cond)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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. Comment |
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
src/op/builtin.cc(1 hunks)src/op/builtin.h(2 hunks)src/target/codegen_cuda.cc(2 hunks)src/target/codegen_cuda.h(1 hunks)src/tl_templates/cuda/debug.h(0 hunks)tilelang/language/print.py(2 hunks)
💤 Files with no reviewable changes (1)
- src/tl_templates/cuda/debug.h
🚧 Files skipped from review as they are similar to previous changes (1)
- src/target/codegen_cuda.h
🧰 Additional context used
🧬 Code graph analysis (3)
src/op/builtin.h (1)
tilelang/language/print.py (1)
device_assert(144-155)
tilelang/language/print.py (1)
tilelang/language/tir/op.py (1)
call_intrin(120-145)
src/op/builtin.cc (1)
tilelang/language/print.py (1)
device_assert(144-155)
⏰ 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 Metal (on macos-latest)
- GitHub Check: Test for Python 3.12 with ROCm-6.3 (on self-hosted-amd)
🔇 Additional comments (5)
tilelang/language/print.py (2)
8-8: LGTM: Import added for intrinsic call support.The import of
tilelang.language as Tis necessary for the intrinsic-based device assertion calls.
151-155: LGTM: Correct migration to intrinsic-based assertions.The switch from
tir.call_externtoT.call_intrinwith registered intrinsics (tl.device_assertandtl.device_assert_with_msg) is correct and aligns with the PR objective.src/op/builtin.h (1)
517-529: LGTM: Intrinsic declarations are correct.The declarations for
device_assert()anddevice_assert_with_msg()follow the established pattern for TL builtin intrinsics. Documentation is clear and concise.src/target/codegen_cuda.cc (1)
2313-2320: LGTM: Correct interception of device assert intrinsics.The modification to
VisitStmt_correctly identifies device assertion intrinsics by name and routes them to the specializedPrintDeviceAsserthandler, preventing them from being processed by the standard codegen path.src/op/builtin.cc (1)
304-312: LGTM: Intrinsic registrations are correct.Both
device_assertanddevice_assert_with_msgare properly registered with the correct number of inputs (1 and 2 respectively) and appropriatekOpaquecall effect kind, which is correct for operations with observable side effects.
| void CodeGenTileLangCUDA::PrintDeviceAssert(const CallNode *call) { | ||
| std::string cond = PrintExpr(call->args[0]); | ||
| if (call->args.size() == 2) { | ||
| if (auto str = call->args[1].as<StringImmNode>()) { | ||
| this->PrintIndent(); | ||
| stream << "if (!(" << cond << ")) printf(\"" << str->value << "\\n\");\n"; | ||
| } | ||
| } | ||
| this->PrintIndent(); | ||
| stream << "assert(" << cond << ");\n"; | ||
| } |
There was a problem hiding this comment.
Verify printf format string safety.
The implementation correctly emits device-side assertions. However, on line 2292, the message is passed directly to printf without format specifiers. If the message contains % characters, this could cause undefined behavior or crashes.
Consider using a format specifier to safely print the message:
- stream << "if (!(" << cond << ")) printf(\"" << str->value << "\\n\");\n";
+ stream << "if (!(" << cond << ")) printf(\"%s\\n\", \"" << str->value << "\");\n";Or escape any % characters in the message string before emitting it.
🤖 Prompt for AI Agents
In src/target/codegen_cuda.cc around lines 2287-2297, the emitted device printf
uses the message string directly which is unsafe if the message contains '%'
characters; change the emission to use a format specifier (e.g. printf("%s\n",
...)) or escape any '%' characters in the StringImmNode value before embedding
it so the runtime printf never interprets stray '%' sequences; update the branch
that handles call->args[1] to either print with "%s\\n" and pass the literal as
the second argument or preprocess the literal to replace each '%' with "%%"
prior to emitting the printf call.
1 participant
Summary by CodeRabbit