update cmake args

CMakeLists.txt

@@ -16,7 +16,7 @@ set(PYTHON_SUPPORTED_VERSIONS "3.8" "3.9" "3.10" "3.11")
 set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.9;9.0")
 
 # Supported AMD GPU architectures.
-set(HIP_SUPPORTED_ARCHS "gfx908;gfx90a;gfx942;gfx926;gfx928;gfx1100")
+set(HIP_SUPPORTED_ARCHS "gfx908;gfx90a;gfx942;gfx1100;gfx906;gfx926;gfx928")
 
 #
 # Supported/expected torch versions for CUDA/ROCm.

cmake/utils.cmake

@@ -114,7 +114,8 @@ function (get_torch_gpu_compiler_flags OUT_GPU_FLAGS GPU_LANG)
       "-DUSE_ROCM"
       "-U__HIP_NO_HALF_CONVERSIONS__"
       "-U__HIP_NO_HALF_OPERATORS__"
-      "-fno-gpu-rdc")
+      "-fno-gpu-rdc"
+      "--gpu-max-threads-per-block=1024")
 
   endif()
   set(${OUT_GPU_FLAGS} ${GPU_FLAGS} PARENT_SCOPE)

setup.py

@@ -118,9 +118,6 @@ class cmake_build_ext(build_ext):
             '-DCMAKE_ARCHIVE_OUTPUT_DIRECTORY={}'.format(self.build_temp),
         ]
 
-        if _is_hip:
-            cmake_args += ['--gpu-max-threads-per-block=1024']
-
         verbose = bool(int(os.getenv('VERBOSE', '0')))
         if verbose:
             cmake_args += ['-DCMAKE_VERBOSE_MAKEFILE=ON']