code: def test_matmul_dyn_shared():
      n = 1024
      block = 16
      A = te.placeholder((n, n), name="A", dtype="float16")
      B = te.placeholder((n, n), name="B", dtype="float16")
  
      def syncthread():
          return tvm.tir.Call(
              None,
              "tir.tvm_storage_sync",
              tvm.runtime.convert(["shared"]),
          )
  
      def test_matmul_ir(A, B, C):
          ib = tvm.tir.ir_builder.create()
  
          tx = te.thread_axis("threadIdx.x")
          ty = te.thread_axis("threadIdx.y")
          bx = te.thread_axis("blockIdx.x")
          by = te.thread_axis("blockIdx.y")
          ib.scope_attr(tx, "thread_extent", block)
          ib.scope_attr(ty, "thread_extent", block)
          ib.scope_attr(bx, "thread_extent", n // block)
          ib.scope_attr(by, "thread_extent", n // block)
  
          A_sh = ib.allocate(
              A.dtype, (block, block), scope="shared.dyn", name="A_sh"
          )  # fp16
          B_sh = ib.allocate(
              B.dtype, (block, block), scope="shared.dyn", name="B_sh"
          )  # fp16
          # Create a dynamic shared memory for the accumulation.
          # This is for testing merging dynamic shared memory alloctions with different data type.
          # In practice, there is no need to allocate a shared memory for C.
          C_local = ib.allocate(
              C.dtype, (1,), scope="local", name="C_local"
          )
          C_sh = ib.allocate(
              C.dtype, (block, block), scope="shared.dyn", name="C_sh"
          )  # fp32
  
          A_ptr = ib.buffer_ptr(A)
          B_ptr = ib.buffer_ptr(B)
          C_ptr = ib.buffer_ptr(C)
  
          C_local[0] = 0.0
  
          with ib.for_range(0, n // block, name="i") as i:
              A_sh[ty, tx] = A_ptr[by * block + ty, i * block + tx]
              B_sh[ty, tx] = B_ptr[i * block + ty, bx * block + tx]
              ib.emit(syncthread())
  
              with ib.for_range(0, block, name="k") as k:
                  C_local[0] += cast(
                      A_sh[ty, k] * B_sh[k, tx], "float32"
                  )
              ib.emit(syncthread())
  
          C_sh[ty, tx] = C_local[0]
          C_ptr[by * block + ty, bx * block + tx] = C_sh[ty, tx]
  
          return ib.get()
  
      C = te.extern(
          A.shape,
          [A, B],
          lambda ins, outs: test_matmul_ir(ins[0], ins[1], outs[0]),
          name="matmul",
          dtype="float32",
      )
      s = te.create_schedule(C.op)
      mod = run_passes(s, [A, B, C])
      # C can be allocated at the start of A, so we only need to allocate 2 block * block memory with dtype = float16
      expected_alloc_size = block * block * 4
      verify_single_allocation(mod["main"].body, expected_alloc_size)
  
      def check_target(target):
          if not tvm.testing.device_enabled(target):
              return
  
          fmatmul = tvm.build(s, [A, B, C], target)
          dev = tvm.device(target, 0)
  
          size = (n, n)
          a_np = np.random.uniform(size=size).astype(A.dtype)
          b_np = np.random.uniform(size=size).astype(B.dtype)
          a = tvm.nd.array(a_np, dev)
          b = tvm.nd.array(b_np, dev)
          c = tvm.nd.array(np.zeros(size, dtype=C.dtype), dev)
          fmatmul(a, b, c)
          np_ref = np.dot(
              a_np.astype("float32"), b_np.astype("float32")
          )
          tvm.testing.assert_allclose(c.numpy(), np_ref, 1e-4, 1e-4)
  
      for target in ["cuda", "nvptx"]:
          check_target(target)
  
imports: import numpy
  import tvm
  import tvm.testing
  from tvm import te
  from tvm.driver.build_module import schedule_to_module
  from tvm.topi.math import cast
  from tvm.script import tir
label_desc: Use image node_id 3 for calling a global function with example usage: test_matmul_dyn_shared() and returns: tvm, ib