def add_elem_types_to_mm_info(mm_info, func_attrs):
    """
    CUTLASS 3.x problem args require explicit I/O pointer types
    (not void*). This function arugments the input and output
    pointers in the mm_info with the appropriate elem_input_type
    and elem_output_type casts.
    """
    backend_spec = CUDASpec()
    elem_input_type = backend_spec.dtype_to_lib_type(
        func_attrs["inputs"][0]._attrs["dtype"]
    )
    elem_output_type = backend_spec.dtype_to_lib_type(
        func_attrs["outputs"][0]._attrs["dtype"]
    )

    return dataclasses.replace(
        mm_info,
        a_ptr=f"({elem_input_type}*)({mm_info.a_ptr})",
        b_ptr=f"({elem_input_type}*)({mm_info.b_ptr})",
        bias_ptr=f"({elem_output_type}*)({mm_info.bias_ptr})",
        c_ptr=f"({elem_output_type}*)({mm_info.c_ptr})",
    )
CUTLASS 3.x problem args require explicit I/O pointer types
(not void*). This function arugments the input and output
pointers in the mm_info with the appropriate elem_input_type
and elem_output_type casts.
python/aitemplate/backend/cuda/gemm_universal/bmm_common.py