[Inductor] Pick ISA for inductor based on ATEN_CPU_CAPABILITY (#123514)

It is part of https://github.com/pytorch/pytorch/issues/123224. Pick ISA based on the environment ATEN_CPU_CAPABILITY to control CPU vec ISA level for Inductor like eager.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/123514
Approved by: https://github.com/jgong5, https://github.com/peterbell10
This commit is contained in:
CaoE
2024-09-23 18:40:21 -07:00
committed by PyTorch MergeBot
parent 9dbc6bacff
commit 6931c1644a
4 changed files with 232 additions and 25 deletions

View File

@ -148,7 +148,10 @@ class ExtensionBackendTests(TestCase):
metrics.reset()
opt_fn = torch.compile()(fn)
_, code = run_and_get_cpp_code(opt_fn, x, y, z)
if cpu_vec_isa.valid_vec_isa_list():
if (
cpu_vec_isa.valid_vec_isa_list()
and os.getenv("ATEN_CPU_CAPABILITY") != "default"
):
load_expr = "loadu"
else:
load_expr = " = in_ptr0[static_cast<long>(i0)];"