def gen_function(
    func_attrs,
    reduce_op,
    prologue_template=DEFAULT_PROLOGUE_TEMPLATE,
    epilogue_scalar_template=DEFAULT_EPILOGUE_SCALAR_TEMPLATE,
    extra_code_str="",
    accumulation_type=None,
    reduction_identity="ElementCompute()",
) -> str:
    """a common function for generating a reduce-family kernel

    Parameters
    ----------
    func_attrs : Dict[str, Any]
        holds attributes of this reduce op
    reduce_op : str
        the reduce op's string representation such as cutlass::plus
    prologue_template : str, optional
        a Template that will be rendered to process input before reduction
    epilogue_scalar_template : str, optional
        a Template that will be rendered to process each final reduced element
    epilogue_scalar_template : str, optional
        a Template that will be rendered to hold extra code for reduction
    accmulation_type : str, optional
        specifies the data type for accumulation
        (default is None so that we will use output's type for accumulation)

    Returns
    -------
    str
        returns the rendered code for the complete implementation of the reduce op
    """
    x = func_attrs["inputs"][0]
    y = func_attrs["outputs"][0]
    backend_spec = CUDASpec()
    input_type = backend_spec.dtype_to_lib_type(x._attrs["dtype"])
    output_type = backend_spec.dtype_to_lib_type(y._attrs["dtype"])
    if accumulation_type is None:
        # follow pytorch's semantics
        if (
            Target.current()._kwargs.get("use_fp16_acc", False)
            and y._attrs["dtype"] == "float16"
        ):
            acc_type = output_type
        else:
            acc_type = "float"
    else:
        acc_type = accumulation_type

    # FIXME: Our kernel launch configs are determined by the amount of shared
    # memory assigned to each block. Currently, we set up an approximate size for
    # shared memroy based on the reduction dim value. We should consider to
    # determine this value by profiling later.
    axes = func_attrs["reduction_axes"]
    if not len(axes) == 1:
        raise NotImplementedError("Multiple reduction axes are not supported yet")
    reduction_axis = axes[0]
    x_shape = x._attrs["shape"]
    reduction_dim_val = max(x_shape[reduction_axis]._attrs["values"])
    col_size_buckets = [8, 16, 32, 64, 128]
    # If reduction axis is not the last dim, let's increase the shared memory a bit.
    # Note that this heuristic is observed with some test inputs, and we need a
    # more systematic approach to determining the size of shared memory.
    if reduction_axis < len(x_shape) - 1:
        col_size_buckets.extend([256, 512])
    col_pos = bisect.bisect_right(col_size_buckets, reduction_dim_val - 1)
    col_pos = min(col_pos, len(col_size_buckets) - 1)
    shared_col_size = col_size_buckets[col_pos]

    row_layout = "RowMajor"
    col_layout = "ColumnMajor"

    # FIXME: these alignments values are only for half_t type.
    # make it adjustable to other types such as float.
    alignments = [8, 4, 2, 1]
    if x._attrs["dtype"] in ("float16", "bfloat16"):
        alignments.append(16)
    # This is ugly. Ideally, we should have templated code like below:
    # template <typename Alignment>
    # reduce_launcher(...) {
    #   using ReductionKernel = ReductionKernel3D<..., 4, ...>
    #   ...
    #   typename ReductionKernel::Layout output_layout(lda_output);
    #   ...
    # }
    #
    # However, this dependent template pattern caused cicc (at least 11.4) to
    # segfault. To workaround this cicc issue, we manually "instantiate" template.
    reduce_instances = [
        REDUCE_KERNEL_INSTANCE.render(
            indent="  ",
            elem_output_type=output_type,
            elem_input_type=input_type,
            elem_compute_type=acc_type,
            align=align,
            shared_col_size=shared_col_size,
            layout=row_layout,
        )
        for align in alignments
    ]
    reduce_instances.append(
        REDUCE_KERNEL_INSTANCE.render(
            indent="  ",
            elem_output_type=output_type,
            elem_input_type=input_type,
            elem_compute_type=acc_type,
            align=1,
            shared_col_size=shared_col_size,
            layout=col_layout,
        )
    )
    reduce_instance = "\n".join(reduce_instances)

    prologue_code = prologue_template.render(indent=" " * 8)
    epilogue_scalar_code = epilogue_scalar_template.render(indent=" " * 12)

    output_accessors = func_attrs["output_accessors"]
    assert (
        len(output_accessors) == 1
    ), f"expected the length of output_accessors to be one but got {len(output_accessors)}"
    dtype = func_attrs["inputs"][0].dtype()
    output_alignment = tensor_accessor_codegen.find_max_alignment_for_accessors(
        dtype, output_accessors
    )
    special_exec_path, special_kernel = reduce_small_axis.get_exec_cond_and_kernel(
        func_attrs,
        reduce_op,
        reduction_axis,
        prologue_code,
        epilogue_scalar_code,
        input_type,
        output_type,
        acc_type,
        output_accessors,
        output_alignment,
        reduction_identity,
    )
    exec_paths = EXEC_COND_TEMPLATE.render(
        indent="  ",
        func_name=func_attrs["name"],
        elem_output_type=output_type,
        elem_input_type=input_type,
        elem_compute_type=acc_type,
        alignments=alignments,
        special_exec_cond=special_exec_path,
    )

    strided_address_func_str = (
        tensor_accessor_codegen.STRIDED_ADDRESS_AT_IDX_FUNC_TEMPLATE.render(
            output_accessor=output_accessors[0],
        )
    )
    tensor_accessor_libs = tensor_accessor_codegen.get_libs()
    extra_code_str += "\n\n" + tensor_accessor_libs
    extra_code_str += "\n\nnamespace {\n" + strided_address_func_str + "\n}\n\n"

    kernel_src = KERNEL_SRC_TEMPLATE.render(
        extra_code=extra_code_str,
        reduce_op=reduce_op,
        reduction_identity=reduction_identity,
        reduce_kernel_instance=reduce_instance,
        alignments=alignments,
        prologue_code=prologue_code,
        epilogue_scalar_code=epilogue_scalar_code,
        special_kernel=special_kernel,
    )

    return SRC_TEMPLATE.render(
        func_name=func_attrs["name"],
        kernel_source=kernel_src,
        exec_paths=exec_paths,
        output_accessor=output_accessors[0],
    )
a common function for generating a reduce-family kernel

Parameters
----------
func_attrs : Dict[str, Any]
holds attributes of this reduce op
reduce_op : str
the reduce op's string representation such as cutlass::plus
prologue_template : str, optional
a Template that will be rendered to process input before reduction
epilogue_scalar_template : str, optional
a Template that will be rendered to process each final reduced element
epilogue_scalar_template : str, optional
a Template that will be rendered to hold extra code for reduction
accmulation_type : str, optional
specifies the data type for accumulation
(default is None so that we will use output's type for accumulation)

Returns
-------
str
returns the rendered code for the complete implementation of the reduce op
python/aitemplate/backend/cuda/reduce/reduce_3d.py