gemini-code-assist[bot] commented on code in PR #19735:
URL: https://github.com/apache/tvm/pull/19735#discussion_r3398724950
##########
src/s_tir/transform/inject_permuted_layout.cc:
##########
@@ -246,6 +246,17 @@ class PermutedLayoutInjector : private
IRMutatorWithAnalyzer {
return access_ptr_call;
}
+ // Device intrinsics are registered under both a flat name (the builtin Op)
+ // and a canonical dotted name (emitted by TVMScript and the tensor
+ // intrinsics), so compare against both.
+ static bool IsOp(const Call& call, const Op& compat_op, const char*
canonical_name) {
+ if (call->op.same_as(compat_op)) {
+ return true;
+ }
+ const auto* op_node = call->op.as<OpNode>();
+ return op_node != nullptr && op_node->name == canonical_name;
+ }
+
PrimExpr VisitExpr_(const CallNode* op) final {
// Rewrite from/to shared or shared.dyn to/from local
auto call = Downcast<Call>(IRMutatorWithAnalyzer::VisitExpr_(op));
Review Comment:
Comparing `Op` pointers via `same_as` is significantly faster than
performing dynamic casts (`as<OpNode>()`) and string comparisons on every
visited `CallNode`. We can look up the canonical `Op`s once using `Op::Get` and
store them as static local variables, then perform direct pointer comparisons.
```suggestion
// Device intrinsics are registered under both a flat name (the builtin Op)
// and a canonical dotted name (emitted by TVMScript and the tensor
// intrinsics), so compare against both.
static bool IsOp(const Call& call, const Op& compat_op, const Op&
canonical_op) {
return call->op.same_as(compat_op) || call->op.same_as(canonical_op);
}
PrimExpr VisitExpr_(const CallNode* op) final {
// Rewrite from/to shared or shared.dyn to/from local
auto call = Downcast<Call>(IRMutatorWithAnalyzer::VisitExpr_(op));
static const Op& ptx_ldmatrix_legacy_op =
Op::Get("tirx.ptx.ldmatrix_legacy");
static const Op& mma_store_legacy_op = Op::Get("tirx.mma_store_legacy");
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
