code: def tir_broadcast_schedule(
      out_m,
      input_a,
      input_b,
      output_layout: str,
      input_a_layout: str,
      input_b_layout: str,
      op_name: str,
  ):
      """Schedule for input and output layout nhwc-8h2w32c2w-2d considering broadcast"""
      func = te.create_prim_func([input_a, input_b, out_m])
  
      s = tir.Schedule(func)
  
      block_dict = {
          "add": "T_add",
          "subtract": "T_subtract",
          "multiply": "T_multiply",
      }
  
      block = s.get_block(block_dict[op_name])
  
      if input_a_layout == "nhwc-8h2w32c2w-2d":
          input_a_transformed_layout = get_layout_transform_fn(
              input_a_layout
          )
          s.transform_layout(
              block,
              buffer=("read", 0),
              index_map=input_a_transformed_layout,
          )
  
      if input_b_layout == "nhwc-8h2w32c2w-2d":
          input_b_transformed_layout = get_layout_transform_fn(
              input_b_layout
          )
          s.transform_layout(
              block,
              buffer=("read", 1),
              index_map=input_b_transformed_layout,
          )
  
      output_transformed_layout = get_layout_transform_fn(output_layout)
      s.transform_layout(
          block,
          buffer=("write", 0),
          index_map=output_transformed_layout,
      )
  
      n, h, w, c = s.get_loops(block)
  
      h_o, h_i = s.split(h, [None, 8])
      w_o, w_i = s.split(w, [None, 4])
      c_o, c_i = s.split(c, [None, 32])
      wio, wii = s.split(w_i, [None, 2])
  
      s.reorder(n, h_o, w_o, c_o, h_i, wio, c_i, wii)
  
      fused = s.fuse(c_i, wii)
      s.vectorize(fused)
  
      return s
  
file_path: tvm/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
imports: from tvm import te
  from tvm import tir
  from tvm import topi
  from ..utils import get_layout_transform_fn
label_desc: Use image node_id 4 for calling a global function with example usage: tir_broadcast_schedule(out_m, input_a, input_b, output_layout, input_a_layout, input_b_layout, op_name) and returns: s