code: def test_gpu():
      def _create_schedule():
          func = matmul
          sch = tir.Schedule(func, debug_mask="all")
          c = sch.get_block("C")
          c_local = sch.cache_write(c, 0, "local")
          i, j, k = sch.get_loops(c)
          # pylint: disable=invalid-name
          i0, i1, i2, i3, i4 = sch.split(
              i, factors=[None, 1, 16, 32, 1]
          )  # outer: 1
          j0, j1, j2, j3, j4 = sch.split(
              j, factors=[None, 4, 1, 1, 16]
          )  # outer: 8
          k0, k1, k2 = sch.split(k, factors=[None, 1, 2])  # outer: 256
          # pylint: enable=invalid-name
          # fmt: off
          sch.reorder(
              i0, j0,  # S
              i1, j1,  # S
              i2, j2,  # S
              k0,      # R
              k1,      # R
              i3, j3,  # S
              k2,      # R
              i4, j4,  # S
          )
          # fmt: on
          # thread binding
          i0_j0 = sch.fuse(i0, j0)
          i1_j1 = sch.fuse(i1, j1)
          i2_j2 = sch.fuse(i2, j2)
          sch.bind(i0_j0, "blockIdx.x")
          sch.bind(i1_j1, "vthread.x")
          sch.bind(i2_j2, "threadIdx.x")
          # fusion
          sch.reverse_compute_at(c_local, i2_j2)
          # cache read 'A'
          a_shared = sch.cache_read(c, 1, "shared")
          sch.compute_at(a_shared, k0)
          _, _, _, _, a_i, a_j = sch.get_loops(a_shared)
          a_ij = sch.fuse(a_i, a_j)
          _, a_j = sch.split(a_ij, factors=[None, 16])  # outer: 64
          sch.bind(a_j, "threadIdx.x")
          # cache read 'B'
          b_shared = sch.cache_read(c, 2, "shared")
          sch.compute_at(b_shared, k0)
          _, _, _, _, b_i, b_j = sch.get_loops(b_shared)
          b_ij = sch.fuse(b_i, b_j)
          _, b_j = sch.split(b_ij, factors=[None, 16])  # outer: 8
          sch.bind(b_j, "threadIdx.x")
          # auto unroll
          sch.annotate(
              i0_j0,
              "pragma_auto_unroll_max_step",
              tir.IntImm("int32", 1024),
          )
          sch.annotate(
              i0_j0, "pragma_unroll_explicit", tir.IntImm("int32", 1)
          )
          return sch
  
      extractor = ms.feature_extractor.PerStoreFeature()
      (feature,) = extractor.extract_from(
          _make_context(tvm.target.Target("cuda")),
          candidates=[_make_candidate(_create_schedule)],
      )
      feature = feature.numpy()
      assert feature.shape == (4, N_FEATURES)
      ### Check feature[0]: BufferStore(A_shared) <= A[...]
      f = feature[0]
      # Group 1.1: arith
      assert_allclose(
          actual=f[0:16],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              24.000000085991324,
              24.000000085991324,
              24.000000085991324,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.2: vectorize
      assert_allclose(
          actual=f[16:27],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.3: unroll
      assert_allclose(
          actual=f[27:38],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.4: parallel
      assert_allclose(
          actual=f[38:49],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.5: is_gpu, blockIdx.x/y/z, threadIdx.x/y/z, vthread
      assert_allclose(
          actual=f[49:57],
          desired=[
              1.0,
              3.169925001442312,
              1.0,
              1.0,
              4.087462841250339,
              1.0,
              1.0,
              2.321928094887362,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.1: Buffer A
      assert_allclose(
          actual=f[57:75],
          desired=[
              1.0,
              0.0,
              0.0,
              25.000000042995662,
              20.000001375860553,
              23.00000017198264,
              14.000088052430122,
              1.0,
              0.0,
              0.0,
              18.00000550343433,
              20.00562591970089,
              2.321928094887362,
              23.00000017198264,
              18.00000550343433,
              21.000000687930438,
              12.0003521774803,
              12.0003521774803,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.2: Buffer A.shared
      assert_allclose(
          actual=f[75:93],
          desired=[
              0.0,
              1.0,
              0.0,
              25.000000042995662,
              12.0003521774803,
              23.00000017198264,
              9.002815015607053,
              1.0,
              0.0,
              0.0,
              6.022367813028454,
              11.98049663618346,
              8.005624549193879,
              17.000011006847668,
              4.087462841250339,
              15.000044026886828,
              1.584962500721156,
              4.087462841250339,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.3: Dummy padding
      assert_allclose(
          actual=f[93:111],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.4: Dummy padding
      assert_allclose(
          actual=f[111:129],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.5: Dummy padding
      assert_allclose(
          actual=f[129:147],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 3: Arithmetic intensity
      assert_allclose(
          actual=f[147:157],
          desired=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 4 & 5
      assert_allclose(
          actual=f[157:164],
          desired=[
              12.0003521774803,
              27.000000010748916,
              17.000011006847668,
              6.022367813028454,
              23.00000017198264,
              2.584962500721156,
              10.001408,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      ### Check feature[1]: BufferStore(B_shared) <= B[...]
      f = feature[1]
      # Group 1.1: arith
      assert_allclose(
          actual=f[0:16],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              21.584962959341485,
              21.584962959341485,
              21.000000687930438,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.2: vectorize
      assert_allclose(
          actual=f[16:27],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.3: unroll
      assert_allclose(
          actual=f[27:38],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.4: parallel
      assert_allclose(
          actual=f[38:49],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.5: is_gpu, blockIdx.x/y/z, threadIdx.x/y/z, vthread
      assert_allclose(
          actual=f[49:57],
          desired=[
              1.0,
              3.169925001442312,
              1.0,
              1.0,
              4.087462841250339,
              1.0,
              1.0,
              2.321928094887362,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.1: Buffer B
      assert_allclose(
          actual=f[57:75],
          desired=[
              1.0,
              0.0,
              0.0,
              22.00000034396526,
              20.000001375860553,
              20.000001375860553,
              14.000088052430122,
              1.0,
              0.0,
              0.0,
              15.000044026886828,
              20.17555076886471,
              2.321928094887362,
              20.000001375860553,
              18.00000550343433,
              18.00000550343433,
              12.0003521774803,
              4.087462841250339,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.2: Buffer B.shared
      assert_allclose(
          actual=f[75:93],
          desired=[
              0.0,
              1.0,
              0.0,
              22.00000034396526,
              9.002815015607053,
              20.000001375860553,
              3.169925001442312,
              1.0,
              0.0,
              0.0,
              3.169925001442312,
              9.61654884377899,
              8.005624549193879,
              14.000088052430122,
              1.584962500721156,
              12.0003521774803,
              0.044394119358453436,
              4.087462841250339,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.3: Dummy padding
      assert_allclose(
          actual=f[93:111],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.4: Dummy padding
      assert_allclose(
          actual=f[111:129],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.5: Dummy padding
      assert_allclose(
          actual=f[129:147],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 3: Arithmetic intensity
      assert_allclose(
          actual=f[147:157],
          desired=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 4 & 5
      assert_allclose(
          actual=f[157:164],
          desired=[
              9.002815015607053,
              24.000000085991324,
              17.000011006847668,
              3.169925001442312,
              20.000001375860553,
              2.584962500721156,
              10.001408,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      ### Check feature[2]: BufferStore(C_local) <= C_local[...] + A_shared[...] * B_shared[...]
      f = feature[2]
      # Group 1.1: arith
      assert_allclose(
          actual=f[0:16],
          desired=[
              0.0,
              27.000000010748916,
              27.000000010748916,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              28.000000005374456,
              28.000000005374456,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.2: vectorize
      assert_allclose(
          actual=f[16:27],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.3: unroll
      assert_allclose(
          actual=f[27:38],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.4: parallel
      assert_allclose(
          actual=f[38:49],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.5: is_gpu, blockIdx.x/y/z, threadIdx.x/y/z, vthread
      assert_allclose(
          actual=f[49:57],
          desired=[
              1.0,
              3.169925001442312,
              1.0,
              1.0,
              4.087462841250339,
              1.0,
              1.0,
              2.321928094887362,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.1: Buffer B.shared
      assert_allclose(
          actual=f[57:75],
          desired=[
              1.0,
              0.0,
              0.0,
              29.00000000268723,
              9.002815015607053,
              23.00000017198264,
              3.169925001442312,
              1.0,
              0.0,
              0.0,
              5.044394119358453,
              7.651051691178929,
              5.044394119358453,
              24.000000085991324,
              4.087462841250339,
              18.00000550343433,
              0.32192809488736235,
              1.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.2: Buffer C.local
      assert_allclose(
          actual=f[75:93],
          desired=[
              0.0,
              0.0,
              1.0,
              29.00000000268723,
              11.000704269011246,
              23.00000017198264,
              5.044394119358453,
              1.0,
              0.0,
              0.0,
              4.087462841250339,
              7.05528243550119,
              1.584962500721156,
              28.000000005374456,
              10.001408194392809,
              22.00000034396526,
              4.087462841250339,
              1.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.3: Buffer A.shared
      assert_allclose(
          actual=f[93:111],
          desired=[
              1.0,
              0.0,
              0.0,
              29.00000000268723,
              12.0003521774803,
              19.00000275171979,
              9.002815015607053,
              1.0,
              0.0,
              0.0,
              1.0,
              3.700439718141092,
              4.087462841250339,
              25.000000042995662,
              8.005624549193879,
              15.000044026886828,
              5.044394119358453,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.4: Dummy padding
      assert_allclose(
          actual=f[111:129],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.5: Dummy padding
      assert_allclose(
          actual=f[129:147],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 3: Arithmetic intensity
      # Arithmetic intensity is high here because of repeated use of a shared
      # buffer. Multiple accesses to the same memory location are counted as a
      # single byte, skewing these numbers towards higher intensity.
      assert_allclose(
          actual=f[147:157],
          desired=[
              11.98533,
              12.977811,
              13.562714,
              13.977722,
              14.299632,
              14.562654,
              14.785038,
              14.977677,
              15.147597,
              15.299596,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 4 & 5
      assert_allclose(
          actual=f[157:164],
          desired=[
              11.000704269011246,
              18.00000550343433,
              9.002815015607053,
              18.00000550343433,
              27.000000010748916,
              3.0,
              10.001408,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      ### Check feature[3]: BufferStore(C) <= C_local[...]
      f = feature[3]
      # Group 1.1: arith
      assert_allclose(
          actual=f[0:16],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.2: vectorize
      assert_allclose(
          actual=f[16:27],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.3: unroll
      assert_allclose(
          actual=f[27:38],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.4: parallel
      assert_allclose(
          actual=f[38:49],
          desired=[
              0.0,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 1.5: is_gpu, blockIdx.x/y/z, threadIdx.x/y/z, vthread
      assert_allclose(
          actual=f[49:57],
          desired=[
              1.0,
              3.169925001442312,
              1.0,
              1.0,
              4.087462841250339,
              1.0,
              1.0,
              2.321928094887362,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.1: Buffer C
      assert_allclose(
          actual=f[57:75],
          desired=[
              0.0,
              1.0,
              0.0,
              20.000001375860553,
              20.000001375860553,
              14.000088052430122,
              14.000088052430122,
              0.0,
              0.0,
              1.0,
              0.0,
              0.0,
              0.0,
              21.000000687930438,
              21.000000687930438,
              15.000044026886828,
              15.000044026886828,
              1.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.2: Buffer C.local
      assert_allclose(
          actual=f[75:93],
          desired=[
              1.0,
              0.0,
              0.0,
              20.000001375860553,
              11.000704269011246,
              14.000088052430122,
              5.044394119358453,
              1.0,
              0.0,
              0.0,
              9.002815015607053,
              12.0003521774803,
              4.087462841250339,
              16.00002201361136,
              7.011227255423254,
              10.001408194392809,
              1.584962500721156,
              1.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.3: Dummy padding
      assert_allclose(
          actual=f[93:111],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.4: Dummy padding
      assert_allclose(
          actual=f[111:129],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 2.5: Dummy padding
      assert_allclose(
          actual=f[129:147],
          desired=[
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
              0.0,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 3: Arithmetic intensity
      assert_allclose(
          actual=f[147:157],
          desired=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
          rtol=1e-5,
          atol=1e-5,
      )
      # Group 4 & 5
      assert_allclose(
          actual=f[157:164],
          desired=[
              20.000001375860553,
              18.00000550343433,
              1.0,
              18.00000550343433,
              18.00000550343433,
              2.584962500721156,
              10.001408,
          ],
          rtol=1e-5,
          atol=1e-5,
      )
  
imports: import sys
  from typing import Callable, List
  import pytest
  import tvm
  import tvm.testing
  from numpy.testing import assert_allclose
  from tvm import meta_schedule
  from tvm import te, tir
  from tvm.script import tir
label_desc: Use image node_id 10 for calling a global function with example usage: test_gpu() and returns: sch