MCPcopy Create free account
hub / github.com/apache/tvm / test_thread_info

Function test_thread_info

tests/python/tirx/codegen/test_codegen_nvshmem.py:78–88  ·  view source on GitHub ↗
(sess)

Source from the content-addressed store, hash-verified

76 sess.sync_worker_0()
77
78 def test_thread_info(sess):
79 @T.prim_func
80 def main(res: T.Buffer((2,), "int32")):
81 T.device_entry()
82 cta_id = T.cta_id([1])
83 tid = T.thread_id([nwarps * 32])
84 res[0] = T.nvshmem.my_pe()
85 res[1] = T.nvshmem.n_pes()
86
87 res_array = sess.empty((2,), "int32")
88 run_prim_func(sess, main, res_array)
89
90 def test_transfer(sess, scope, shape, nwarps, nelems, op_name):
91 """Tests data transfer operations (get/put) at thread, warp, and block scopes."""

Callers 1

_test_funcFunction · 0.85

Calls 2

run_prim_funcFunction · 0.85
emptyMethod · 0.45

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…