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
)