enable cpu adam op on powerpc architectures (#1213)

* enable cpu adam operation on powerpc

* fix formatting

Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Jeff Rasley <jerasley@microsoft.com>

csrc/includes/cpu_adam.h

@@ -1,10 +1,13 @@
 #pragma once
 
+#if (__x86_64__ || __i386__)
 #include <cpuid.h>
+#include <x86intrin.h>
+#endif
+
 #include <cuda_fp16.h>
 #include <cuda_runtime_api.h>
 #include <stdio.h>
-#include <x86intrin.h>
 #include <cassert>
 #include "context.h"
 #include "cublas_v2.h"

op_builder/cpu_adam.py

@@ -29,6 +29,21 @@ def include_paths(self):
         CUDA_INCLUDE = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")
         return ['csrc/includes', CUDA_INCLUDE]
 
+    def cpu_arch(self):
+        if not self.command_exists('lscpu'):
+            self.warning(
+                "CPUAdam attempted to query 'lscpu' to detect the CPU architecture. "
+                "However, 'lscpu' does not appear to exist on "
+                "your system, will fall back to use -march=native.")
+            return ''
+
+        result = subprocess.check_output('lscpu', shell=True)
+        result = result.decode('utf-8').strip().lower()
+        if 'ppc64le' in result:
+            # gcc does not provide -march on PowerPC, use -mcpu instead
+            return '-mcpu=native'
+        return '-march=native'
+
     def simd_width(self):
         if not self.command_exists('lscpu'):
             self.warning(
@@ -49,6 +64,7 @@ def simd_width(self):
     def cxx_args(self):
         import torch
         CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib64")
+        CPU_ARCH = self.cpu_arch()
         SIMD_WIDTH = self.simd_width()
 
         return [
@@ -59,7 +75,7 @@ def cxx_args(self):
             '-lcublas',
             '-g',
             '-Wno-reorder',
-            '-march=native',
+            CPU_ARCH,
             '-fopenmp',
             SIMD_WIDTH
         ]