Skip to content

[PyTorch] Add fused addmm path in linear for contiguous 3D input #72728

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
swolchok wants to merge 11 commits into from
Closed

[PyTorch] Add fused addmm path in linear for contiguous 3D input #72728

Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
6b2a078
[PyTorch] Add fused addmm path in linear for contiguous 3D input
swolchok Feb 11, 2022
d5f3665
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Feb 11, 2022
173b9d6
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Feb 19, 2022
e3a2f5e
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Feb 22, 2022
350677c
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Feb 24, 2022
40d0838
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Feb 25, 2022
aa0e8af
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Feb 26, 2022
744edb5
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Mar 1, 2022
0b17745
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Mar 7, 2022
d7239fd
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Mar 8, 2022
693aad7
Update on "[PyTorch] Add fused addmm path in linear for contiguous 3D…
swolchok Mar 25, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions aten/src/ATen/native/Linear.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,12 @@ Tensor linear(const Tensor& input, const Tensor& weight, const c10::optional<Ten
// Fused op is marginally faster.
return at::addmm(*bias, input, weight.t());
}
if (input.dim() == 3 && bias->defined() && input.is_contiguous()) {
// Also hit the fused path for contiguous 3D input.
const auto input_sizes = input.sizes();
const auto result = at::addmm(*bias, input.view({input_sizes[0] * input_sizes[1], input_sizes[2]}), weight.t());
return result.view({input_sizes[0], input_sizes[1], result.size(1)});
}
auto output = at::matmul(input, weight.t());
if (bias->defined()) {
output.add_(*bias);
Expand Down