Strength reduction: fold transpose into a subsequent GEMM call

🚀 Feature

The program:

  class DynamoModule(torch.nn.Module):
    def forward(self, L_intermediate_parallel_ : torch.Tensor, L_self_modules_dense_4h_to_h_parameters_weight_ : torch.nn.
parameter.Parameter):
        l_intermediate_parallel_ = L_intermediate_parallel_
        l_self_modules_dense_4h_to_h_parameters_weight_ = L_self_modules_dense_4h_to_h_parameters_weight_
        function_ctx = torch.autograd.function.FunctionCtx();  function_ctx = None
        t = l_self_modules_dense_4h_to_h_parameters_weight_.t();  l_self_modules_dense_4h_to_h_parameters_weight_ = None
        output = torch.matmul(l_intermediate_parallel_, t);  l_intermediate_parallel_ = t = None
        function_ctx_1 = torch.autograd.function.FunctionCtx();  function_ctx_1 = None
        return (output,)

results in the trace:

def computation(L_intermediate_parallel_, L_self_modules_dense_4h_to_h_parameters_weight_):
  # L_intermediate_parallel_: "cuda:0 bf16[384, 2, 13824]"
  # L_self_modules_dense_4h_to_h_parameters_weight_: "cuda:0 bf16[5120, 13824]"
  [t] = nvFusion0(L_self_modules_dense_4h_to_h_parameters_weight_)
    # t = prims.transpose(L_self_modules_dense_4h_to_h_parameters_weight_, (1, 0))  # t: "cuda:0 bf16[13824, 5120]"
  output = torch.matmul(L_intermediate_parallel_, t)  # output: "cuda:0 bf16[384, 2, 5120]"
    # output = ltorch.matmul(L_intermediate_parallel_, t)  # output: "cuda:0 bf16[384, 2, 5120]"
      # t7 = ltorch.reshape(L_intermediate_parallel_, -1, 13824)  # t7: "cuda:0 bf16[768, 13824]"
        # t7 = prims.reshape(L_intermediate_parallel_, (768, 13824))  # t7: "cuda:0 bf16[768, 13824]"
      # t8 = prims.matmul(t7, t)  # t8: "cuda:0 bf16[768, 5120]"
      # output = ltorch.reshape(t8, 384, -1, 5120)  # output: "cuda:0 bf16[384, 2, 5120]"
        # output = prims.reshape(t8, (384, 2, 5120))  # output: "cuda:0 bf16[384, 2, 5120]"
  del t
  return (output,)

g23.py.txt -- full program showing the above.

This single-op transpose nvFusion region before the matmul is unfortunate. Since matmul under the hood is implemented with gemm, we should just be able to use a so-called "NT" GEMM and skip the transpose as a separate op beforehand. Libraries such as cuBLAS often have specialized kernels for such cases.

Motivation

This is coming up in NeVA through the ThunderFX path, and I suspect will come up in other situations as well.

In this particular case we get hit with a double whammy:

• The additional op induces a CPU delay before we can get to launching the GPU work.
• we end up doing a "TN" GEMM, which generally means we access the second array in exactly the wrong memory order. (I imagine in practice BLAS routines would just pre-transpose things rather than pay this cost during the computation itself, but of course then that would mean both Thunder and BLAS transpose this tensor.)

Pitch

Other than "fuse the transpose into the gemm" I am not entirely sure.

• Does aten expose the ability to specify the transpose args on its gemms? Ideally we'd just use that.
• Can nvFuser's in-progress work to accept gemms take this on and do the right thing?
• Do we need a special "transposable gemm" executor that just takes transpose+matmuls and skips aten?

Discussion welcome.

cc @tfogal

Start by reproducing the attached g23.py.txt program and examining the shown nvFusion trace around transpose and matmul. Investigate the aten, nvFuser and ThunderFX paths mentioned in the issue, then determine whether the transpose can be represented as an NT GEMM; done means the transpose is no longer emitted as a separate operation while results remain correct.