[PHI] Align paddle.inner with torch in matmul logic

[lshpku]

PR Category

Operator Mechanism

PR Types

Bug fixes

Description

对齐paddle.inner(x, y)与torch在matmul上的调用逻辑

原理

paddle动态图原来是先调用 y.transpose()，再调用 matmul，这会导致inner的执行变成2个kernel：transpose(y) + matmul(x, y^T)：

void phi::funcs::TilingSwapDim1And2<phi::dtype::float16, (int)256, (int)32, (int)32, long>
void cutlass::Kernel<cutlass_75_wmma_tensorop_s161616gemm_f16_32x32_128x2_nn_align1>

本PR将其改为直接调用 matmul(x, y, transpose_y=True)，这样底层的分发逻辑就会只执行1个kernel（外加一个小reduce
kernel）：

void cutlass::Kernel<cutlass_75_wmma_tensorop_s161616gemm_f16_32x32_128x2_tn_align1>
void splitKreduce_kernel<(int)32, (int)16, int, float, __half, float, __half, (bool)1, (bool)0, ...>

注意到，前一个cutlass的签名是nn，后一个是tn

参考torch

1. torch.inner(x, y) 先调用了 tensordot(x, y, -1, -1)：https://github.com/pytorch/pytorch/blob/afd7a13bca2ce518b7c32f868c8dba610a538e22/aten/src/ATen/native/LinearAlgebra.cpp#L1308
2. 然后 tensordot(x, y) 调用 y.permute()，再进一步调用 mm(x, y^T)：https://github.com/pytorch/pytorch/blob/afd7a13bca2ce518b7c32f868c8dba610a538e22/aten/src/ATen/native/Linear.cpp#L764
3. tensordot虽然也对y进行permute了，但torch的permute是lazy的，实际上调用到cutlass层是只调用了一个matmul kernel

对齐情况

本PR可以做到与torch完全对齐（结果完全相等，nsys trace完全相同）
测试范围：x[M, C]，y[N, C]；其中 1 <= N, M <= 2²⁸，1 <= C <= 2³⁰（在显存能装下的范围内）

性能比原先的 transpose + matmul 版本提升一倍左右，因为省了transpose的成本

另外，静态图无需修改，因为静态图会自动把 matmul_v2(x, y.T) 变成 matmul(x, y, transpose_y: true)

Pcard-85711

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Please upload report for BASE (develop@7869937). Learn more about missing BASE report.

Additional details and impacted files

@@             Coverage Diff             @@
##             develop    #72843   +/-   ##
===========================================
  Coverage           ?   100.00%           
===========================================
  Files              ?         1           
  Lines              ?         1           
  Branches           ?         0           
===========================================
  Hits               ?         1           
  Misses             ?         0           
  Partials           ?         0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.