def make_function_strided_args(
    func_attrs,
    dim_info_dict,
    default_mm_info,
    is_permute=False,
):
    """
    Return a tuple of (problem_args, input_addr_calculator, output_addr_calculator)
    """
    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"]
    )

    input_a_batch_stride_dim = default_mm_info.a_batch_stride
    input_a_stride_lda_dim = default_mm_info.lda
    input_a_offset = 0
    input_b_batch_stride_dim = default_mm_info.b_batch_stride
    input_b_stride_ldb_dim = default_mm_info.ldb
    input_b_offset = 0

    has_bias = len(func_attrs["inputs"]) == 3

    if "input_accessors" in func_attrs:
        input_a_accessor = func_attrs["input_accessors"][0]
        input_b_accessor = func_attrs["input_accessors"][1]

        if input_a_accessor.is_from_strided_tensor:
            input_a_offset = input_a_accessor.offset
            if not input_a_accessor.is_contiguous:
                a_dims = reverse_dim_info_mapping(
                    dim_info_dict, gemm_common.Source.INPUT, 0
                )

                input_a_batch_stride_dim = input_a_accessor.gen_stride_str(0, a_dims)
                input_a_stride_lda_dim = input_a_accessor.stride(1)

        if input_b_accessor.is_from_strided_tensor:
            input_b_offset = input_b_accessor.offset
            if not input_b_accessor.is_contiguous:
                b_dims = reverse_dim_info_mapping(
                    dim_info_dict, gemm_common.Source.INPUT, 1
                )
                input_b_batch_stride_dim = input_b_accessor.gen_stride_str(0, b_dims)
                input_b_stride_ldb_dim = input_b_accessor.stride(1)

        if has_bias:
            # FIXME: we don't suppor strided bias yet. Will enable it once
            # we support it.
            input_bias_accessor = func_attrs["input_accessors"][2]
            assert (
                not input_bias_accessor.is_from_strided_tensor
            ), f'strided bias is not supported for op {func_attrs["name"]}'

    input_addr_calculator = common.INPUT_ADDR_CALCULATOR.render(
        input_a_batch_stride_dim=input_a_batch_stride_dim,
        input_a_stride_dim=input_a_stride_lda_dim,
        input_a_offset_val=input_a_offset,
        input_b_batch_stride_dim=input_b_batch_stride_dim,
        input_b_stride_dim=input_b_stride_ldb_dim,
        input_b_offset_val=input_b_offset,
    )

    # bmm_permute requires a slightly different c_batch_stride and
    # output_batch_stride_dim values
    if is_permute:
        output_batch_stride_dim = default_mm_info.bias_batch_stride
        c_batch_stride = default_mm_info.c_batch_stride
    else:
        output_batch_stride_dim = default_mm_info.c_batch_stride
        c_batch_stride = "output_batch_stride"
    output_stride_ldc_dim = default_mm_info.ldc
    output_offset = 0

    if "output_accessors" in func_attrs:
        output_accessor = func_attrs["output_accessors"][0]
        if output_accessor.is_from_strided_tensor:
            output_offset = output_accessor.offset
            if not output_accessor.is_contiguous:
                c_dims = reverse_dim_info_mapping(
                    dim_info_dict, gemm_common.Source.OUTPUT, 0
                )
                output_batch_stride_dim = output_accessor.gen_stride_str(0, c_dims)
                output_stride_ldc_dim = output_accessor.stride(1)

    output_addr_calculator = OUTPUT_ADDR_CALCULATOR.render(
        output_batch_stride_dim=output_batch_stride_dim,
        output_stride_dim=output_stride_ldc_dim,
        output_offset_val=output_offset,
    )

    bmm_problem_info = Bmm_problem_info(
        alpha_value=default_mm_info.alpha_value,
        beta_value=default_mm_info.beta_value,
        a_ptr=f"({elem_input_type}*)({default_mm_info.a_ptr}) + input_a_offset",
        b_ptr=f"({elem_input_type}*)({default_mm_info.b_ptr}) + input_b_offset",
        bias_ptr=f"({elem_output_type}*)({default_mm_info.bias_ptr})",
        c_ptr=f"({elem_output_type}*)({default_mm_info.c_ptr}) + output_offset",
        a_batch_stride="input_a_batch_stride",
        b_batch_stride="input_b_batch_stride",
        bias_batch_stride=f"{default_mm_info.bias_batch_stride}",
        c_batch_stride=c_batch_stride,
        lda="input_a_stride",
        ldb="input_b_stride",
        ldbias=f"{default_mm_info.ldbias}",
        ldc="output_stride",
        a_row_major=default_mm_info.a_row_major,
        b_row_major=default_mm_info.b_row_major,
        c_row_major=default_mm_info.c_row_major,
    )
    a_shapes = func_attrs["input_accessors"][0].original_shapes
    b_shapes = func_attrs["input_accessors"][1].original_shapes
    d_shapes = None
    if has_bias:
        d_shapes = func_attrs["input_accessors"][2].original_shapes
    _update_stride_info(bmm_problem_info, a_shapes, b_shapes, d_shapes)

    problem_args = PROBLEM_ARGS_TEMPLATE.render(
        mm_info=bmm_problem_info,
    )
    problem_args_cutlass_3x = PROBLEM_ARGS_TEMPLATE_CUTLASS_3X.render(
        mm_info=bmm_problem_info,
    )
    return (
        problem_args,
        problem_args_cutlass_3x,
        input_addr_calculator,
        output_addr_calculator,
    )
Return a tuple of (problem_args, input_addr_calculator, output_addr_calculator)
python/aitemplate/backend/cuda/gemm_universal/bmm_common.py