Add aten mkldnn linear operator

[bddppq]

Stack:
    :white_circle:  #19633 Add is_mkldnn to at::Tensor  💚
    :white_circle:  #19204 Add aten mkldnn conv2d operator  💚
    :white_circle:  #19205 Add aten mkldnn ops: relu, max_pool2d and avg_pool2d  💚
    :white_circle:  #19206 Add aten mkldnn batch_norm operator  💚
    :white_circle:  #19207 Add aten mkldnn add operator  💚
    :white_circle:  #19209 Add aten mkldnn view operator  💚
    :black_circle:  #19210 Add aten mkldnn linear operator  💚
    :white_circle:  #19648 Adjust resnext run script  💚

Pull Request resolved: #19210

Differential Revision: D14901641

[bddppq]

@zdevito @suo @ZolotukhinM In this PR I'm changing nn.Linear to directly call c++ aten::linear (which underlyingly dispatches to the same addmm/matmul call), but the difference is nn.Linear's torchscript now will directly capture linear instead of addmm/matmul. (In this PR only scripting get changed, tracing is not unaffected, but I would like to change tracing as well in a follow up PR) I think this is the right change since the torchscript now get simplified from

graph(%x : Tensor,
      %23 : Tensor,
      %24 : Tensor):
  %2 : int = prim::Constant[value=1]()
  %3 : None = prim::Constant()
  %4 : bool = prim::Constant[value=0]()
  %5 : int = prim::Constant[value=2]()
  %9 : int = aten::dim(%x)
  %10 : bool = aten::eq(%9, %5)
  %11 : bool = prim::If(%10)
    block0():
      %12 : bool = aten::__isnot__(%23, %3)
      -> (%12)
    block1():
      -> (%4)
  %ret : Tensor = prim::If(%11)
    block0():
      %bias.2 : Tensor = prim::unchecked_unwrap_optional(%23)
      %15 : Tensor = aten::t(%24)
      %ret.1 : Tensor = aten::addmm(%bias.2, %x, %15, %2, %2)
      -> (%ret.1)
    block1():
      %17 : Tensor = aten::t(%24)
      %output.1 : Tensor = aten::matmul(%x, %17)
      %19 : bool = aten::__isnot__(%23, %3)
      %output : Tensor = prim::If(%19)
        block0():
          %bias.3 : Tensor = prim::unchecked_unwrap_optional(%23)
          %output.2 : Tensor = aten::add_(%output.1, %bias.3, %2)
          -> (%output.2)
        block1():
          -> (%output.1)
      -> (%output)
  return (%ret)

to

graph(%x : Tensor,
      %6 : Tensor,
      %7 : Tensor):
  %5 : Tensor = aten::linear(%x, %7, %6)
  return (%5)

And it makes other backends like mkldnn/conversion works like work's life easier to support linear.
However our concern is this change might break some compiler passes that specifically pattern match addmm/matmul, could you give me some pointers to where I can check them?

[eellison]

I'm not sure that we should do this. If you run the scripted function you get the following result:

@torch.jit.script
def test_linear(x, y):
    return F.linear(x, y)

test_linear(torch.rand([2, 2]), torch.rand([2, 2]))
print(test_linear.graph_for(torch.rand([2, 2]), torch.rand([2, 2])))

Which outputs the graph

graph(%x : Double(*, *),
      %y : Double(*, *)):
  %15 : Double(*, *) = aten::t(%y)
  %output.1 : Double(*, *) = aten::matmul(%x, %15)
  return (%output.1)

What is the end goal here ? If we continually create new ops which were previously composed of existing ones, we'll be putting more and more stress on various parts of the system -shape analysis, the graph fuser, TVM, etc.

I don't have a good context on mkldnn interop so there may be valid reasons there. But as far as torchscript I'm not really convinced.

[zdevito]

It is totally reasonable that there would be a built in linear operator. Tons of libraries already have it as optimized fused thing. We should add it.

However, @eellison is right: anytime a fused op is added, it is also necessary to change all of the JIT analysis and optimization passes so that they still work in this world. In the case of Linear, this is change is almost certainly breaking matrix multiply optimizations that the JIT does because of missing formulas for shape propagation, differentiation, and others.

This is an invasive change on a very important operator, so it deserves careful consideration about what possibly is breaking by hiding its implementation from optimization passes.

[dzhulgakov]

@zdevito - can you please help to come up with change plan for this change in JIT. As you said, switching to at:linear does make sense. I guess grepping for aten::addmm should be a good starting point.

One option is to land this PR first as it switches only scripting and follow up with a separate PR for doing tracing and proper JIT changes. Would it be reasonable?

[bddppq]

I added a workaround that overrides the forward method of nn.Linear in case in to_mkldnn. Will do the linear dispatch change in a separate diff to unblock mkldnn.

[dzhulgakov]

As a hack with the follow up later - it looks good

[dzhulgakov]

this does make a copy, right? (it's supposed to)

[bddppq]

yes it's a copy ctor https://bddppq.github.io/codebrowser/pytorch/pytorch/third_party/ideep/include/ideep/tensor.hpp.html#1198

[facebook-github-bot]

This pull request has been merged in c9f380d.