code: def test_vthread(vthread_name):
      dtype = "int64"
      n = 100
      m = 4
      nthread = 2
  
      def get_vthread(name):
          tx = te.thread_axis(name)
          ty = te.thread_axis(name)
          ib = tvm.tir.ir_builder.create()
          A = ib.pointer("float32", name="A")
          C = ib.pointer("float32", name="C")
          with ib.for_range(0, n) as i:
              ib.scope_attr(tx, "virtual_thread", nthread)
              ib.scope_attr(ty, "virtual_thread", nthread)
              B = ib.allocate("float32", m, name="B", scope="shared")
              B[i] = A[i * nthread + tx]
              bbuffer = B.asobject()
              ib.emit(
                  tvm.tir.call_extern(
                      "int32",
                      "Run",
                      bbuffer.access_ptr("r"),
                      tvm.tir.call_intrin(
                          "int32", "tir.tvm_context_id"
                      ),
                  )
              )
              C[i * nthread + tx] = B[i] + 1
          return ib.get()
  
      if vthread_name == "vthread":
          B_expected_alloc = m * nthread
      elif vthread_name == "cthread":
          B_expected_alloc = m * nthread * nthread
  
      stmt = tvm.tir.transform.InjectVirtualThread()(
          tvm.IRModule.from_expr(
              tvm.tir.PrimFunc([], get_vthread(vthread_name))
          )
      )["main"]
  
      assert list(stmt.body.body.extents) == [B_expected_alloc]
  
file_path: tvm/tests/python/tir-transform/test_tir_transform_inject_virtual_thread.py
imports: import tvm
  import tvm.testing
  from tvm import te
  from tvm.script import tir
label_desc: Use image node_id 1 for calling a global function with example usage: test_vthread(vthread_name) and returns: ib