code: def test_three_stage_compute_two_stage_async():
      mod = tvm.IRModule.from_expr(
          three_stage_compute.with_attr("global_symbol", "main")
      )
      sch = tvm.tir.Schedule(mod)
  
      _, loop = sch.get_loops(sch.get_block("compute"))
      sch.annotate(
          loop, ann_key="software_pipeline_async_stages", ann_val=[0, 1]
      )
  
      mod = tvm.tir.transform.InjectSoftwarePipeline()(sch.mod)
  
      @T.prim_func
      def ref(
          A: T.Buffer((16, 16), "float32"),
          D: T.Buffer((16, 16), "float32"),
      ) -> None:
          for tx in T.thread_binding(16, thread="threadIdx.x"):
              with T.block():
                  T.reads(A[tx, 0:16])
                  T.writes(D[tx, 0:16])
                  B = T.alloc_buffer(
                      [2, 16, 1], dtype="float32", scope="shared"
                  )
                  C = T.alloc_buffer(
                      [2, 16, 1], dtype="float32", scope="shared"
                  )
                  with T.block():
                      T.reads(A[tx, 0:2], B[0:2, tx, 0])
                      T.writes(B[0:2, tx, 0], C[0:2, tx, 0])
                      for i in T.unroll(2):
                          with T.block():
                              T.where(i < 16)
                              T.reads(A[tx, i])
                              T.writes(B[i % 2, tx, 0])
                              with T.attr(
                                  0, "async_commit_queue_scope", 0
                              ):
                                  with T.attr(0, "async_scope", 1):
                                      B[i % 2, tx, 0] = A[
                                          tx, i
                                      ] * T.float32(2)
                          with T.block():
                              T.where(i == 1 and i - 1 < 16)
                              T.reads(B[(i - 1) % 2, tx, 0])
                              T.writes(C[(i - 1) % 2, tx, 0])
                              with T.attr(
                                  0, "async_commit_queue_scope", 1
                              ):
                                  with T.attr(
                                      0, "async_wait_queue_scope", 0
                                  ):
                                      with T.attr(
                                          0,
                                          "async_wait_inflight_count",
                                          1,
                                      ):
                                          with T.attr(
                                              0, "async_scope", 1
                                          ):
                                              C[(i - 1) % 2, tx, 0] = B[
                                                  (i - 1) % 2, tx, 0
                                              ] + T.float32(2)
                  with T.block():
                      T.reads(A[tx, 2:16], B[0:2, tx, 0], C[0:2, tx, 0])
                      T.writes(
                          B[0:2, tx, 0], C[0:2, tx, 0], D[tx, 0:14]
                      )
                      for i in T.serial(14):
                          with T.block():
                              T.where(i + 2 < 16)
                              T.reads(A[tx, i + 2])
                              T.writes(B[(i + 2) % 2, tx, 0])
                              with T.attr(
                                  0, "async_commit_queue_scope", 0
                              ):
                                  with T.attr(0, "async_scope", 1):
                                      B[(i + 2) % 2, tx, 0] = A[
                                          tx, i + 2
                                      ] * T.float32(2)
                          with T.block():
                              T.where(i + 2 - 1 < 16)
                              T.reads(B[(i - 1 + 2) % 2, tx, 0])
                              T.writes(C[(i - 1 + 2) % 2, tx, 0])
                              with T.attr(
                                  0, "async_commit_queue_scope", 1
                              ):
                                  with T.attr(
                                      0, "async_wait_queue_scope", 0
                                  ):
                                      with T.attr(
                                          0,
                                          "async_wait_inflight_count",
                                          1,
                                      ):
                                          with T.attr(
                                              0, "async_scope", 1
                                          ):
                                              C[
                                                  (i - 1 + 2) % 2, tx, 0
                                              ] = B[
                                                  (i - 1 + 2) % 2, tx, 0
                                              ] + T.float32(
                                                  2
                                              )
                          with T.block():
                              T.where(i + 2 - 2 < 16)
                              T.reads(C[0:2, tx, 0])
                              T.writes(D[tx, i - 2 + 2])
                              with T.attr(
                                  0, "async_wait_queue_scope", 1
                              ):
                                  with T.attr(
                                      0, "async_wait_inflight_count", 1
                                  ):
                                      D[tx, i - 2 + 2] = C[
                                          (i - 2 + 2) % 2, tx, 0
                                      ] + T.float32(1)
                  with T.block():
                      T.reads(B[0:2, tx, 0], C[0:2, tx, 0])
                      T.writes(C[0:2, tx, 0], D[tx, 14:16])
                      for i in T.unroll(2):
                          with T.block():
                              T.where(i + 16 - 1 < 16)
                              T.reads(B[(i - 1 + 16) % 2, tx, 0])
                              T.writes(C[(i - 1 + 16) % 2, tx, 0])
                              with T.attr(
                                  0, "async_commit_queue_scope", 1
                              ):
                                  with T.attr(
                                      0, "async_wait_queue_scope", 0
                                  ):
                                      with T.attr(
                                          0,
                                          "async_wait_inflight_count",
                                          0 - i,
                                      ):
                                          with T.attr(
                                              0, "async_scope", 1
                                          ):
                                              C[
                                                  (i - 1 + 16) % 2,
                                                  tx,
                                                  0,
                                              ] = B[
                                                  (i - 1 + 16) % 2,
                                                  tx,
                                                  0,
                                              ] + T.float32(
                                                  2
                                              )
                          with T.block():
                              T.where(i + 16 - 2 < 16)
                              T.reads(C[0:2, tx, 0])
                              T.writes(D[tx, i - 2 + 16])
                              with T.attr(
                                  0, "async_wait_queue_scope", 1
                              ):
                                  with T.attr(
                                      0,
                                      "async_wait_inflight_count",
                                      T.if_then_else(
                                          i + 16 - 1 < 16,
                                          1,
                                          0,
                                          dtype="int32",
                                      ),
                                  ):
                                      D[tx, i - 2 + 16] = C[
                                          (i - 2 + 16) % 2, tx, 0
                                      ] + T.float32(1)
  
      tvm.ir.assert_structural_equal(
          mod["main"], ref.with_attr("global_symbol", "main"), True
      )
  
imports: import sys
  import numpy
  import pytest
  import tvm
  import tvm.testing
  import tvm.tir.tensor_intrin.cuda
  from tvm import TVMError, te, tir
  from tvm.meta_schedule.testing import te_workload
  from tvm.script import tir
  from tvm.testing.tir import mma_schedule
  from tvm.tir.tensor_intrin.cuda import LDMATRIX_f16_A_DYN_INTRIN, LDMATRIX_f16_B_DYN_INTRIN, MMA_f16f16f32_INTRIN, MMA_fill_16x16_f32_INTRIN, MMA_store_16x16_f32_global_INTRIN, shared_16x16_to_ldmatrix_32x8_layout
label_desc: Use image node_id 41 for calling a global function with example usage: test_three_stage_compute_two_stage_async() without return types