8000 use statically known true instead of guard size oblivious in bmm and mm inductor decompositions . by laithsakka · Pull Request #148893 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

use statically known true instead of guard size oblivious in bmm and mm inductor decompositions . #148893

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
wants to merge 4 commits into from
Closed

use statically known true instead of guard size oblivious in bmm and mm inductor decompositions . #148893

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0ff6b44
use statically known true instead of guard size oblivous
laithsakka Mar 10, 2025
90f0943
Update on "[DRAFT] use statically known true instead of guard size ob…
laithsakka Mar 16, 2025
2e26d07
Update on "[DRAFT] use statically known true instead of guard size ob…
laithsakka Apr 21, 2025
b68c0ef
Update on "use statically known true instead of guard size oblivious …
laithsakka Apr 21, 2025
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
26 changes: 14 additions & 12 deletions torch/_inductor/decomposition.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,11 @@
ELEMENTWISE_TYPE_PROMOTION_KIND,
type_to_dtype,
)
from torch.fx.experimental.symbolic_shapes import definitely_true, guard_size_oblivious
from torch.fx.experimental.symbolic_shapes import (
definitely_true,
guard_size_oblivious,
statically_known_true,
)

from . import config, inductor_prims
from .utils import (
Expand Down Expand Up @@ -265,13 +269,13 @@ def bmm(
# TODO: Re-enable for mps once our reductions are performant enough
# (https://github.com/pytorch/pytorch/issues/150121)
if config.coordinate_descent_tuning and self.device.type not in ["cpu", "mps"]:
if guard_size_oblivious(self.shape[1] == 1) or guard_size_oblivious(
if statically_known_true(self.shape[1] == 1) or statically_known_true(
batch2.shape[2] == 1
):
out = (self.unsqueeze(-1) * batch2.unsqueeze(1)).sum(dim=2)
return out
if self.device.type == "cpu":
if guard_size_oblivious(self.size(1) == 1) and guard_size_oblivious(
if statically_known_true(self.size(1) == 1) and statically_known_true(
batch2.size(-1) == 1
):
counters["inductor"]["decompose_bmm"] += 1
Expand All @@ -291,7 +295,7 @@ def addmm(
alpha: torch.types.Number = 1,
) -> torch.Tensor:
if self.device.type == "cpu":
if guard_size_oblivious(mat1.size(0) == 1) and guard_size_oblivious(
if statically_known_true(mat1.size(0) == 1) and statically_known_true(
mat2.size(-1) == 1
):
counters["inductor"]["decompose_addmm"] += 1
Expand All @@ -300,7 +304,7 @@ def addmm(
).unsqueeze(0)
return alpha * out + beta * self
if (
guard_size_oblivious(mat1.size(0) == 1)
statically_known_true(mat1.size(0) == 1)
and definitely_true(mat2.size(0) <= 16)
and definitely_true(mat2.size(1) <= 16)
):
Expand All @@ -322,21 +326,21 @@ def mm(
# TODO: Re-enable for mps once our reductions are performant enough
# (https://github.com/pytorch/pytorch/issues/150121)
if config.coordinate_descent_tuning and self.device.type not in ["cpu", "mps"]:
if guard_size_oblivious(self.shape[0] == 1) or guard_size_oblivious(
if statically_known_true(self.shape[0] == 1) or statically_known_true(
input2.shape[1] == 1
):
return (self.unsqueeze(2) * input2.unsqueeze(0)).sum(dim=1)
if self.device.type == "cpu":
if (
guard_size_oblivious(self.size(-1) == 1)
and guard_size_oblivious(self.size(0) > 0)
and guard_size_oblivious(input2.size(0) == 1)
statically_known_true(self.size(-1) == 1)
and statically_known_true(self.size(0) > 0)
and statically_known_true(input2.size(0) == 1)
and (self.dtype == input2.dtype)
and definitely_true((torch.numel(self) + torch.numel(input2)) <= 32)
):
counters["inductor"]["decompose_mm"] += 1
return torch.cat([self[i, :] * input2 for i in range(self.size(0))])
if guard_size_oblivious(self.size(0) == 1) and guard_size_oblivious(
if statically_known_true(self.size(0) == 1) and statically_known_true(
input2.size(-1) == 1
):
counters["inductor"]["decompose_mm"] += 1
Expand All @@ -355,8 +359,6 @@ def cat(
tensors: list[torch.Tensor],
dim: int = 0,
) -> torch.Tensor:
from torch.fx.experimental.symbolic_shapes import guard_size_oblivious

def non_empty_tensor(x: torch.Tensor) -> bool:
# For better or worse, this is a valid cat:
#
Expand Down
Loading
0