[GitHub] [tvm] junrushao1994 commented on a change in pull request #10405: [TensorIR] Update VerifyGPU

[GitBox <gi...@apache.org>]

junrushao1994 commented on a change in pull request #10405:
URL: https://github.com/apache/tvm/pull/10405#discussion_r815627197



##########
File path: src/meta_schedule/postproc/verify_gpu_code.cc
##########
@@ -21,6 +21,51 @@
 #include "../utils.h"
 
 namespace tvm {
+namespace tir {
+class ThreadExtentChecker : private StmtVisitor {
+ public:
+  static bool Check(const Stmt& stmt) {
+    try {
+      ThreadExtentChecker().VisitStmt(stmt);
+      return true;
+    } catch (const dmlc::Error& e) {
+      return false;
+    }
+  }
+
+ private:
+  void VisitStmt_(const ForNode* loop) {
+    if (IsThreadIdx(GetThreadScope(loop))) {
+      if (const int64_t* p_ext = GetLoopIntExtent(loop)) {
+        thread_extent_product *= *p_ext;
+        StmtVisitor::VisitStmt_(loop);
+        thread_extent_product /= *p_ext;
+        return;
+      } else {
+        throw dmlc::Error("Dynamic thread extent");
+      }
+    }
+    StmtVisitor::VisitStmt_(loop);
+  }
+
+  void VisitStmt_(const BlockNode* block) {
+    if (Optional<Integer> low_inclusive =
+            GetAnn<Integer>(block, attr::meta_schedule_thread_extent_low_inclusive)) {
+      if (Optional<Integer> high_inclusive =
+              GetAnn<Integer>(block, attr::meta_schedule_thread_extent_high_inclusive)) {
+        int64_t low = low_inclusive.value()->value;
+        int64_t high = high_inclusive.value()->value;
+        if (!(low <= thread_extent_product && thread_extent_product <= high)) {
+          throw dmlc::Error("Thread extent");
+        }
+      }
+    }
+    StmtVisitor::VisitStmt_(block);
+  }
+
+  int64_t thread_extent_product = 1;

Review comment:
       @Hzfengsy A potential issue of this approach is that we might count threads twice for cooperative fetching blocks. For example,
   
   ```
   for tx in T.thread_binding(0, 128, "threadIdx.x"):
     for tx_2 in T.thread_binding(0, 128, "threadIdx.x"):
       // thread_extent_product counted as 128 * 128
       ...
   ```
   
   Let's fix this :-)




-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org