Check that CPUInductor selects valid SIMD ISA

malfet · malfet · commit 954bf675a7bb · 2024-08-08T10:13:20.000+01:00
I.e. for every platform we support, it should be something but invalid
AVX512 on Intel Linux, and NEON on Apple silicon and Linux aarch64
diff --git a/test/smoke_test/smoke_test.py b/test/smoke_test/smoke_test.py
@@ -258,6 +258,14 @@ def foo(x: torch.Tensor) -> torch.Tensor:
         x_pt2 = torch.compile(foo)(x)
         torch.testing.assert_close(x_eager, x_pt2)
 
+    # Check that SIMD were detected for the architecture
+    if device == "cpu":
+        from torch._inductor.cpu_vec_isa import invalid_vec_isa, pick_vec_isa
+        isa = pick_vec_isa()
+        if isa == invalid_vec_isa:
+            raise RuntimeError("Can't detect vectorized ISA for CPU")
+        print(f"Picked CPU ISA {type(isa).__name__} bit width {isa.bit_width()}")
+
     # Reset torch dynamo since we are changing mode
     torch._dynamo.reset()
     dtype = torch.float32
