code: def test_shared_memory():
      def check_shared_memory(storage_scope, dtype):
          N = 1024
          M = 128
  
          tvm_type = tvm.runtime.DataType(dtype)
          type_size = tvm_type.bits // 8 * tvm_type.lanes
  
          A = te.placeholder((N,), name="A", dtype=dtype)
          B = te.compute((N,), lambda i: A[i], name="B")
  
          s = te.create_schedule([B.op])
          AA = s.cache_read(A, storage_scope, [B])
          o, i = s[B].split(s[B].op.axis[0], M)
          s[AA].compute_at(s[B], o)
          s[B].bind(o, te.thread_axis("blockIdx.x"))
          s[B].bind(i, te.thread_axis("threadIdx.x"))
  
          # shared memory usage: M * sizeof(dtype) Bytes
          # thread usage: M
  
          for target in ["opencl", "cuda"]:
              if not tvm.testing.device_enabled(target):
                  continue
              valid = [None]
              with tvm.transform.PassContext(
                  config={
                      "tir.add_lower_pass": [
                          (
                              2,
                              get_verify_pass(
                                  valid,
                                  max_shared_memory_per_block=type_size
                                  * M
                                  - 1,
                                  max_threads_per_block=M,
                              ),
                          )
                      ]
                  }
              ):
                  tvm.build(s, [A, B], target)
              assert not valid[0]
  
              with tvm.transform.PassContext(
                  config={
                      "tir.add_lower_pass": [
                          (
                              2,
                              get_verify_pass(
                                  valid,
                                  max_shared_memory_per_block=type_size
                                  * M,
                                  max_threads_per_block=M,
                              ),
                          )
                      ]
                  }
              ):
                  tvm.build(s, [A, B], target)
              assert valid[0]
  
      check_shared_memory("shared", "float32")
      check_shared_memory("shared", "int8x4")
      check_shared_memory("shared.dyn", "float32")
  
file_path: tvm/tests/python/tir-analysis/test_tir_analysis_verify_gpu_code.py
imports: import tvm
  from tvm import te
  from tvm import topi
  import tvm.testing
  import tvm.topi.testing
label_desc: Use image node_id 2 for calling a global function with example usage: test_shared_memory() without return types