code: def test_vthread_simplified():
      """Indices resulting from vthread injection should simplified
  
      This ensures that downstream passes that check for Ramp nodes do
      not need to each simplify the indices.
      """
  
      @T.prim_func
      def before_func():
          vthread = T.env_thread("vthread")
          T.launch_thread(vthread, 4)
          B_data = T.allocate([4], "int32", scope="shared")
          B = T.Buffer([4], "int32", data=B_data, scope="shared")
          B[0:4] = T.broadcast(vthread, 4)
  
      @T.prim_func
      def expected_func():
          B_data = T.allocate([16], "int32", scope="shared")
          B = T.Buffer([16], "int32", data=B_data, scope="shared")
          # The indices for B should each be a single Ramp node, and
          # should not be the sum of a Ramp and Broadcast node.
          B[T.Mul(0, 4) : T.Mul(0, 4) + 4] = T.broadcast(0, 4)
          B[T.Mul(1, 4) : T.Mul(1, 4) + 4] = T.broadcast(1, 4)
          B[T.Mul(2, 4) : T.Mul(2, 4) + 4] = T.broadcast(2, 4)
          B[T.Mul(3, 4) : T.Mul(3, 4) + 4] = T.broadcast(3, 4)
  
      before_mod = tvm.IRModule.from_expr(
          before_func.with_attr("global_symbol", "main")
      )
      after_mod = tvm.tir.transform.InjectVirtualThread()(before_mod)
      after_func = after_mod["main"]
  
      tvm.ir.assert_structural_equal(
          after_func, expected_func.with_attr("global_symbol", "main")
      )
  
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 4 for calling a global function with example usage: test_vthread_simplified() without return types