infra to load yaml from cwrap
diff --git a/aten/src/aten/CMakeLists.txt b/aten/src/aten/CMakeLists.txt
index 48da378..f01d113 100644
--- a/aten/src/aten/CMakeLists.txt
+++ b/aten/src/aten/CMakeLists.txt
@@ -83,7 +83,8 @@
EXECUTE_PROCESS(
COMMAND python ${CMAKE_CURRENT_SOURCE_DIR}/gen.py -s ${CMAKE_CURRENT_SOURCE_DIR} --print-dependencies
- OUTPUT_VARIABLE generated_cpp
+ # user stderr rather than stdout so we can still debug the script with print
+ ERROR_VARIABLE generated_cpp
RESULT_VARIABLE RETURN_VALUE
)
if (NOT RETURN_VALUE EQUAL 0)
@@ -94,9 +95,27 @@
FILE(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/TensorLib)
+SET(CWRAP_FILES_BASE ${CMAKE_CURRENT_SOURCE_DIR}/../../csrc )
+SET(cwrap_files
+# ${CWRAP_FILES_BASE}/cudnn/cuDNN.cwrap
+ ${CWRAP_FILES_BASE}/generic/TensorMethods.cwrap
+# ${CWRAP_FILES_BASE}/generic/methods/SparseTensor.cwrap
+ ${CWRAP_FILES_BASE}/generic/methods/Tensor.cwrap
+ ${CWRAP_FILES_BASE}/generic/methods/TensorApply.cwrap
+ ${CWRAP_FILES_BASE}/generic/methods/TensorCompare.cwrap
+ ${CWRAP_FILES_BASE}/generic/methods/TensorCuda.cwrap
+ ${CWRAP_FILES_BASE}/generic/methods/TensorMath.cwrap
+ ${CWRAP_FILES_BASE}/generic/methods/TensorRandom.cwrap
+# ${CWRAP_FILES_BASE}/generic/methods/TensorSerialization.cwrap
+# ${CWRAP_FILES_BASE}/nn/THNN.cwrap
+# ${CWRAP_FILES_BASE}/nn/THCUNN.cwrap
+# ${CWRAP_FILES_BASE}/nn/THNN_generic.cwrap
+)
+
+
ADD_CUSTOM_COMMAND(OUTPUT ${generated_cpp}
-COMMAND python ${CMAKE_CURRENT_SOURCE_DIR}/gen.py -s ${CMAKE_CURRENT_SOURCE_DIR}
-DEPENDS gen.py ${all_templates})
+COMMAND python ${CMAKE_CURRENT_SOURCE_DIR}/gen.py -s ${CMAKE_CURRENT_SOURCE_DIR} ${cwrap_files}
+DEPENDS gen.py cwrap_parser.py code_template.py ${all_templates} ${cwrap_files})
SET(all_cpp ${base_cpp} ${generated_cpp})
filter_list(generated_h generated_cpp "\\.h$")
diff --git a/aten/src/aten/CodeTemplate.py b/aten/src/aten/code_template.py
similarity index 100%
rename from aten/src/aten/CodeTemplate.py
rename to aten/src/aten/code_template.py
diff --git a/aten/src/aten/cwrap_parser.py b/aten/src/aten/cwrap_parser.py
new file mode 100644
index 0000000..76e9eec
--- /dev/null
+++ b/aten/src/aten/cwrap_parser.py
@@ -0,0 +1,20 @@
+import yaml
+
+# follows similar logic to cwrap, ignores !inc, and just looks for [[]]
+def parse(filename):
+ with open(filename,'r') as file:
+ declaration_lines = []
+ declarations = []
+ in_declaration = False
+ for line in file.readlines():
+ line = line.rstrip()
+ if line == '[[':
+ declaration_lines = []
+ in_declaration = True
+ elif line == ']]':
+ in_declaration = False
+ declaration = yaml.load('\n'.join(declaration_lines))
+ declarations.append(declaration)
+ elif in_declaration:
+ declaration_lines.append(line)
+ return declarations
diff --git a/aten/src/aten/gen.py b/aten/src/aten/gen.py
index d6010b1..dd6b93e 100644
--- a/aten/src/aten/gen.py
+++ b/aten/src/aten/gen.py
@@ -1,13 +1,18 @@
import os
-from optparse import OptionParser
-from CodeTemplate import CodeTemplate
import sys
+import yaml
+from optparse import OptionParser
+
+import cwrap_parser
+from code_template import CodeTemplate
+
+
parser = OptionParser()
parser.add_option('-s', '--source-path', help='path to source director for tensorlib',
action='store', default='.')
parser.add_option('-p', '--print-dependencies',
help='only output a list of dependencies', action='store_true')
-options,args = parser.parse_args()
+options,cwrap_files = parser.parse_args()
TEMPLATE_PATH = options.source_path+"/templates"
GENERATOR_DERIVED = CodeTemplate.from_file(TEMPLATE_PATH+"/GeneratorDerived.h")
@@ -58,7 +63,7 @@
def write(filename,s):
filename = "TensorLib/"+filename
if options.print_dependencies:
- sys.stdout.write(filename+";")
+ sys.stderr.write(filename+";")
return
with open(filename,"w") as f:
f.write(s)
@@ -116,6 +121,14 @@
top_env['type_registrations'].append(type_register)
top_env['type_headers'].append('#include "TensorLib/{}.h"'.format(env['Type']))
+
+
+declarations = [ d
+ for file in cwrap_files
+ for d in cwrap_parser.parse(file) ]
+
+print(yaml.dump(declarations))
+
for fname,env in generators.items():
write(fname,GENERATOR_DERIVED.substitute(env))
