def METHOD_NAME(op, sch, is_idx_reduce=False):
    if is_idx_reduce:
        data_out = op.input_tensors[0]
    else:
        data_in = op.input_tensors[0]
        data_out = op.output(0)
    if not sch[data_out].op.reduce_axis:
        return schedule_injective_from_existing(sch, op.output(0))
    if len(sch[data_out].op.axis) > 0:
        all_reduce = False
        num_thread = 32
        target = tvm.target.Target.current()
        if target and (target.kind.name == "opencl" or target.kind.name == "metal"):
            # without it, CL_INVALID_WORK_GROUP_SIZE occurred when running test_topi_reduce.py
            # don't know why
            num_thread = 16
        block_x = te.thread_axis("blockIdx.x")
        thread_x = te.thread_axis((0, num_thread), "threadIdx.x")
        thread_y = te.thread_axis((0, num_thread), "threadIdx.y")
    else:
        all_reduce = True
        num_thread = tvm.target.Target.current(allow_none=False).max_num_threads
        thread_x = te.thread_axis((0, num_thread), "threadIdx.x")
    # Fuse and refactor the reduce axis
    fused_reduce = sch[data_out].fuse(
        *[sch[data_out].op.reduce_axis[i] for i in range(len(sch[data_out].op.reduce_axis))]
    )
    ko, ki = sch[data_out].split(fused_reduce, factor=num_thread)
    if is_idx_reduce:
        data_out_rf, _ = sch.rfactor(data_out, ki)
    else:
        data_out_rf = sch.rfactor(data_out, ki)
    tx = sch[data_out].op.reduce_axis[0]
    sch[data_out].bind(tx, thread_x)
    sch[data_out_rf].compute_at(sch[data_out], tx)
    if is_idx_reduce:
        real_output = op.output(0)
        temp_idx_input = data_out.op.output(0)
        temp_val_input = data_out.op.output(1)
    else:
        real_output = data_out
    if not all_reduce:
        # Fuse and split the axis
        fused_outer = sch[real_output].fuse(
            *[sch[real_output].op.axis[i] for i in range(len(sch[real_output].op.axis))]
        )
        bx, outer_in = sch[real_output].split(fused_outer, factor=num_thread)
        # Bind the axes to threads and blocks
        sch[real_output].bind(outer_in, thread_y)
        sch[real_output].bind(bx, block_x)
        if is_idx_reduce:
            sch[temp_idx_input].compute_at(sch[real_output], outer_in)
            sch[temp_val_input].compute_at(sch[real_output], outer_in)
        sch[real_output].set_store_predicate(
            tvm.tir.all(
                thread_x.equal(0), block_x * num_thread + thread_y < reduce(mul, real_output.shape)
            )
        )
    else:
        if is_idx_reduce:
            spatial_axis = sch[real_output].fuse(*(sch[real_output].op.axis))
            sch[real_output].bind(spatial_axis, te.thread_axis("blockIdx.x"))
            sch[temp_idx_input].compute_at(sch[real_output], spatial_axis)
            sch[temp_val_input].compute_at(sch[real_output], spatial_axis)
        sch[real_output].set_store_predicate(thread_x.equal(0))
    return sch