Allow different Microsoft products for the search of VC.Tools.x86.x64

python/setup.py

@@ -119,13 +119,13 @@ def find_visual_studio(version_ranges):
     for version_range in version_ranges:
         command = [
             str(vswhere), "-version", version_range, "-requires", "Microsoft.VisualStudio.Component.VC.Tools.x86.x64",
-            "-property", "installationPath", "-prerelease"
+            "-products", "*", "-property", "installationPath", "-prerelease"
         ]
 
         try:
             output = subprocess.check_output(command, text=True).strip()
             if output:
-                return output
+                return output.split("\n")[0]
         except subprocess.CalledProcessError:
             continue
 
@@ -146,6 +146,13 @@ def set_env_vars(vs_path, arch="x64"):
             os.environ[var] = value
 
 
+def initialize_visual_studio_env(version_ranges, arch="x64"):
+    vs_path = find_visual_studio(version_ranges)
+    if not vs_path:
+        raise EnvironmentError("Visual Studio not found in specified version ranges.")
+    set_env_vars(vs_path, arch)
+
+
 # Taken from https://github.com/pytorch/pytorch/blob/master/tools/setup_helpers/env.py
 def check_env_flag(name: str, default: str = "") -> bool:
     return os.getenv(name, default).upper() in ["ON", "1", "YES", "TRUE", "Y"]

python/triton/runtime/CLFinder.py

@@ -19,13 +19,13 @@ def find_visual_studio(version_ranges):
     for version_range in version_ranges:
         command = [
             str(vswhere), "-version", version_range, "-requires", "Microsoft.VisualStudio.Component.VC.Tools.x86.x64",
-            "-property", "installationPath", "-prerelease"
+            "-products", "*", "-property", "installationPath", "-prerelease"
         ]
 
         try:
             output = subprocess.check_output(command, text=True).strip()
             if output:
-                return output
+                return output.split("\n")[0]
         except subprocess.CalledProcessError:
             continue
 
@@ -37,7 +37,7 @@ def set_env_vars(vs_path, arch="x64"):
     if not vcvarsall_path.exists():
         raise FileNotFoundError(f"vcvarsall.bat not found in expected path: {vcvarsall_path}")
 
-    command = f'call "{vcvarsall_path}" {arch} && set'
+    command = ["call", vcvarsall_path, arch, "&&", "set"]
     output = subprocess.check_output(command, shell=True, text=True)
 
     for line in output.splitlines():
@@ -47,9 +47,7 @@ def set_env_vars(vs_path, arch="x64"):
 
 
 def initialize_visual_studio_env(version_ranges, arch="x64"):
-    # Check if the environment variable that vcvarsall.bat sets is present
-    if os.environ.get('VSCMD_ARG_TGT_ARCH') != arch:
-        vs_path = find_visual_studio(version_ranges)
-        if not vs_path:
-            raise EnvironmentError("Visual Studio not found in specified version ranges.")
-        set_env_vars(vs_path, arch)
+    vs_path = find_visual_studio(version_ranges)
+    if not vs_path:
+        raise EnvironmentError("Visual Studio not found in specified version ranges.")
+    set_env_vars(vs_path, arch)