已开启
多项式 #17
zzll创建于  25 天前
zzll成员
25 天前 创建
    def addmm_args(
        mat1,
        mat2,
        inp,
        *others,
        layout=None,
        out_dtype=None,
        use_4x2_dim=False,
        mat2_transposed=False,
    ):
        from torch._inductor.select_algorithm import realize_inputs
        """
        Common arg processing for mm,bmm,addmm,etc
        """
        mat1, mat2, inp = realize_inputs(mat1, mat2, inp)
        *b1, m, k1 = mat1.get_size()
        if mat2_transposed:
            *b2, n, k2 = mat2.get_size()
        else:
            *b2, k2, n = mat2.get_size()
        b = [V.graph.sizevars.check_equals_and_simplify(a, b) for a, b in zip(b1, b2)]
        if use_4x2_dim:
            k2 = k2 * 2
        k = V.graph.sizevars.check_equals_and_simplify(k1, k2)
        if layout is None:
            from torch._inductor.ir import FixedLayout

            if out_dtype is None:
                out_dtype = mat1.get_dtype()

            layout = FixedLayout(
                mat1.get_device(),
                out_dtype,
                [*b, m, n],
            )
        else:
            assert out_dtype is None, "out_dtype is ignored if layout is specified."
        from torch._inductor.lowering import expand

        others = [realize_inputs(expand(x, layout.size)) for x in others]
        # inp_expand = realize_inputs(expand(inp, layout.size))

        return [m, n, k, layout, mat1, mat2, inp, *others]
        
        
        m, n, k, layout, mat1, mat2, inp_origin, inp_expanded = addmm_args(
            mat1, mat2, inp, inp, layout=layout
        )
likedislike
Zzzll成员
25 天前 修改了issue 的描述
Zzzll成员
25 天前 修改了issue 的描述
Zzzll成员
10 天前 修改了issue 的描述
Zzzll成员
10 天前 修改了issue 的描述
Zzzll成员
10 天前 修改了issue 的描述
Zzzll成员
9 天前 修改了issue 的描述
Zzzll成员
9 天前 修改了issue 的描述