Parallelize TensorMethods.cpp builds (pytorch#1400)

.gitignore

@@ -19,6 +19,7 @@ torch/csrc/nn/THCUNN.cpp
 torch/csrc/nn/THNN_generic.cwrap
 torch/csrc/nn/THNN_generic.cpp
 torch/csrc/nn/THNN_generic.h
+torch/csrc/generated
 docs/src/**/*
 test/data/legacy_modules.t7
 test/data/gpu_tensors.pt

setup.py

@@ -15,6 +15,7 @@
 from tools.setup_helpers.env import check_env_flag
 from tools.setup_helpers.cuda import WITH_CUDA, CUDA_HOME
 from tools.setup_helpers.cudnn import WITH_CUDNN, CUDNN_LIB_DIR, CUDNN_INCLUDE_DIR
+from tools.setup_helpers.split_types import split_types
 DEBUG = check_env_flag('DEBUG')
 WITH_DISTRIBUTED = check_env_flag('WITH_DISTRIBUTED')
 WITH_DISTRIBUTED_MW = WITH_DISTRIBUTED and check_env_flag('WITH_DISTRIBUTED_MW')
@@ -259,7 +260,6 @@ def run(self):
     "torch/csrc/Generator.cpp",
     "torch/csrc/Size.cpp",
     "torch/csrc/Exceptions.cpp",
-    "torch/csrc/Tensor.cpp",
     "torch/csrc/Storage.cpp",
     "torch/csrc/DynamicTypes.cpp",
     "torch/csrc/byte_order.cpp",
@@ -283,6 +283,7 @@ def run(self):
     "torch/csrc/autograd/functions/init.cpp",
     "torch/csrc/nn/THNN_generic.cpp",
 ]
+main_sources += split_types("torch/csrc/Tensor.cpp")
 
 try:
     import numpy as np
@@ -325,11 +326,11 @@ def run(self):
         "torch/csrc/cuda/Module.cpp",
         "torch/csrc/cuda/Storage.cpp",
         "torch/csrc/cuda/Stream.cpp",
-        "torch/csrc/cuda/Tensor.cpp",
         "torch/csrc/cuda/AutoGPU.cpp",
         "torch/csrc/cuda/utils.cpp",
         "torch/csrc/cuda/serialization.cpp",
     ]
+    main_sources += split_types("torch/csrc/cuda/Tensor.cpp")
 
 if WITH_NCCL:
     if SYSTEM_NCCL:

tools/setup_helpers/split_types.py

@@ -0,0 +1,58 @@
+import os
+
+this_file = os.path.dirname(os.path.abspath(__file__))
+generated_dir = os.path.abspath(os.path.join(this_file, '..', '..', 'torch', 'csrc', 'generated'))
+
+line_start = '//generic_include '
+
+types = [
+    'Double',
+    'Float',
+    'Half',
+    'Long',
+    'Int',
+    'Short',
+    'Char',
+    'Byte'
+]
+
+generic_include = '#define {lib}_GENERIC_FILE "{path}"'
+generate_include = '#include "{lib}/{lib}Generate{type}Type.h"'
+
+
+def split_types(file_name):
+    assert file_name.startswith('torch/csrc/')
+    if not os.path.exists(generated_dir):
+        os.makedirs(generated_dir)
+
+    with open(file_name, 'r') as f:
+        lines = f.read().split('\n')
+
+    # Find //generic_include
+    for i, l in enumerate(lines):
+        if l.startswith(line_start):
+            args = l[len(line_start):]
+            lib_prefix, generic_file = filter(bool, args.split())
+            break
+    else:
+        raise RuntimeError("generic include not found")
+
+    gen_name_prefix = file_name[len('torch/csrc/'):].replace('/', '_').replace('.cpp', '')
+    gen_path_prefix = os.path.join(generated_dir, gen_name_prefix)
+
+    prefix = '\n'.join(lines[:i])
+    suffix = '\n'.join(lines[i + 1:])
+
+    to_build = []
+
+    g_include = generic_include.format(lib=lib_prefix, path=generic_file)
+    for t in types:
+        t_include = generate_include.format(lib=lib_prefix, type=t)
+        gen_path = gen_path_prefix + t + '.cpp'
+        to_build.append(gen_path)
+        with open(gen_path, 'w') as f:
+            f.write(prefix + '\n' +
+                    g_include + '\n' +
+                    t_include + '\n' +
+                    suffix)
+    return to_build

torch/csrc/Tensor.cpp

@@ -9,12 +9,8 @@
 #include <tuple>
 #include <TH/THMath.h>
 
-#include "THP.h"
-#include "copy_utils.h"
-#include "DynamicTypes.h"
+#include "torch/csrc/THP.h"
+#include "torch/csrc/copy_utils.h"
+#include "torch/csrc/DynamicTypes.h"
 
-#include "generic/Tensor.cpp"
-#include <TH/THGenerateAllTypes.h>
-
-#include "generic/Tensor.cpp"
-#include <TH/THGenerateHalfType.h>
+//generic_include TH torch/csrc/generic/Tensor.cpp

torch/csrc/cuda/Tensor.cpp

@@ -6,14 +6,13 @@
 #include <vector>
 #include <stack>
 #include <tuple>
-#include "THCP.h"
+#include "torch/csrc/cuda/THCP.h"
 
-#include "override_macros.h"
+#include "torch/csrc/cuda/override_macros.h"
 #include "torch/csrc/copy_utils.h"
 #include "DynamicTypes.h"
 
-#define THC_GENERIC_FILE "torch/csrc/generic/Tensor.cpp"
-#include <THC/THCGenerateAllTypes.h>
+//generic_include THC torch/csrc/generic/Tensor.cpp
 
-#include "undef_macros.h"
-#include "restore_macros.h"
+#include "torch/csrc/cuda/undef_macros.h"
+#include "torch/csrc/cuda/restore_macros.h"

torch/csrc/generic/methods/TensorMath.cwrap

@@ -1625,8 +1625,7 @@
     - THTensor* tensor2
 ]]
 
-#ifndef THP_LAPACK_CONSTANTS
-#define THP_LAPACK_CONSTANTS
+#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE) || CUDA_FLOAT || CUDA_DOUBLE
 // We need to pass pointers to chars to tensor lapack functions...
 static const char __U = 'U';
 static const char __L = 'L';