code: def test_multiple_cache_write():
      """Test multiple cache writes."""
      # graph
      arr_length = 1024
      arr_length_tvm = tvm.runtime.convert(arr_length)
      placeholder_a0 = te.placeholder(
          (arr_length_tvm,), name="A0", dtype="float32"
      )
      placeholder_a1 = te.placeholder(
          (arr_length_tvm,), name="A1", dtype="float32"
      )
      result_b0, result_b1 = te.compute(
          (arr_length_tvm,),
          lambda *i: (
              placeholder_a0(*i) + placeholder_a1(*i),
              placeholder_a0(*i) * placeholder_a1(*i),
          ),
          name="B",
      )
      result_c = te.compute(
          (arr_length_tvm,),
          lambda *i: result_b0(*i) + result_b1(*i),
          name="C",
      )
      schedule = te.create_schedule(result_c.op)
      # create iter var and assign them tags.
      num_thread = 8
      cache_b0, _ = schedule.cache_write(
          [result_b0, result_b1], "local"
      )
      axis0, axis1 = schedule[result_c].split(
          result_c.op.axis[0], factor=num_thread
      )
      schedule[result_b0].compute_at(schedule[result_c], axis0)
      schedule[cache_b0].compute_at(schedule[result_c], axis0)
      schedule[result_c].bind(axis0, te.thread_axis("blockIdx.x"))
      schedule[result_c].bind(axis1, te.thread_axis("threadIdx.x"))
  
      # one line to build the function.
      def check_device(device, host="stackvm"):
          if not tvm.testing.device_enabled(host):
              return
          dev = tvm.device(device, 0)
          if not tvm.testing.device_enabled(device):
              return
          func = tvm.build(
              schedule,
              [placeholder_a0, placeholder_a1, result_c],
              device,
              host,
              name="multiple_cache_write",
          )
          dev = tvm.device(device, 0)
          # launch the kernel.
          buff_a0 = tvm.nd.array(
              np.random.uniform(size=arr_length).astype(
                  placeholder_a0.dtype
              ),
              dev,
          )
          buff_a1 = tvm.nd.array(
              np.random.uniform(size=arr_length).astype(
                  placeholder_a1.dtype
              ),
              dev,
          )
          buff_c = tvm.nd.array(
              np.zeros(arr_length, dtype=result_c.dtype), dev
          )
          func(buff_a0, buff_a1, buff_c)
          tvm.testing.assert_allclose(
              buff_c.numpy(),
              buff_a0.numpy()
              + buff_a1.numpy()
              + (buff_a0.numpy() * buff_a1.numpy()),
              rtol=1e-5,
          )
  
      check_device("cuda", "llvm")
      check_device("vulkan")
      check_device("opencl")
  
file_path: tvm/tests/python/integration/test_ewise.py
imports: import numpy
  import tvm
  import tvm.testing
  from tvm import te
  from tvm.contrib import nvcc
label_desc: Use image node_id 3 for calling a global function with example usage: test_multiple_cache_write() without return types