add nvcc support (apache#7668)

neo-ai · May 11, 2021 · 060105f · 060105f
1 parent 2f29ba8
commit 060105f
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 7 deletions.
diff --git a/python/tvm/contrib/cc.py b/python/tvm/contrib/cc.py
@@ -192,12 +192,16 @@ def _fcompile(outputs, objects, options=None):
 
 def _linux_compile(output, objects, options, compile_cmd="g++", compile_shared=False):
     cmd = [compile_cmd]
-    if compile_shared or output.endswith(".so") or output.endswith(".dylib"):
-        cmd += ["-shared", "-fPIC"]
-        if sys.platform == "darwin":
-            cmd += ["-undefined", "dynamic_lookup"]
-    elif output.endswith(".obj"):
-        cmd += ["-c"]
+    if compile_cmd != "nvcc":
+        if compile_shared or output.endswith(".so") or output.endswith(".dylib"):
+            cmd += ["-shared", "-fPIC"]
+            if sys.platform == "darwin":
+                cmd += ["-undefined", "dynamic_lookup"]
+        elif output.endswith(".obj"):
+            cmd += ["-c"]
+    else:
+        if compile_shared or output.endswith(".so") or output.endswith(".dylib"):
+            cmd += ["--shared"]
     cmd += ["-o", output]
     if isinstance(objects, str):
         cmd += [objects]

diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
@@ -340,6 +340,9 @@ def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=No
                 else:
                     assert module.type_key == "c"
                     object_format = "c"
+                    if "cc" in kwargs:
+                        if kwargs["cc"] == "nvcc":
+                            object_format = "cu"
                     has_c_module = True
             path_obj = os.path.join(workspace_dir, f"lib{index}.{object_format}")
             module.save(path_obj)

diff --git a/src/relay/backend/contrib/codegen_c/codegen_c.h b/src/relay/backend/contrib/codegen_c/codegen_c.h
@@ -343,6 +343,8 @@ class CodegenCBase {
     std::string dtype;
     if (runtime::TypeMatch(ttype->dtype, kDLFloat, 32)) {
       dtype = "float";
+    } else if (runtime::TypeMatch(ttype->dtype, kDLFloat, 16)) {
+      dtype = "half";
     } else if (runtime::TypeMatch(ttype->dtype, kDLInt, 32)) {
       dtype = "int";
     } else if (runtime::TypeMatch(ttype->dtype, kDLInt, 64)) {

diff --git a/src/target/source/source_module.cc b/src/target/source/source_module.cc
@@ -104,7 +104,7 @@ class CSourceModuleNode : public runtime::ModuleNode {
   void SaveToFile(const std::string& file_name, const std::string& format) final {
     std::string fmt = GetFileFormat(file_name, format);
     std::string meta_file = GetMetaFilePath(file_name);
-    if (fmt == "c") {
+    if (fmt == "c" || fmt == "cu") {
       ICHECK_NE(code_.length(), 0);
       SaveBinaryToFile(file_name, code_);
     } else {