Fix bugs in PyTorch codegen. (halide#7443)

ardier · Mar 3, 2024 · 902f764 · 902f764
1 parent 4d97824
commit 902f764
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/CodeGen_PyTorch.cpp b/src/CodeGen_PyTorch.cpp
@@ -25,6 +25,7 @@ void CodeGen_PyTorch::compile(const Module &module) {
                           "Please add \"-user_context\" to the generator's target options.\n";
         }
         stream << "#include \"ATen/cuda/CUDAContext.h\"\n";
+        stream << "#include \"HalidePyTorchCudaHelpers.h\"\n";
     }
     stream << "#include \"HalideBuffer.h\"\n";
     stream << "#include \"HalidePyTorchHelpers.h\"\n";
@@ -43,6 +44,11 @@ void CodeGen_PyTorch::compile(const Module &module) {
     }
 
     for (const auto &f : module.functions()) {
+        // Don't put non-external function declarations in headers.
+        // We need to be consistent with CodeGen_C::compile.
+        if (f.linkage == LinkageType::Internal) {
+            continue;
+        }
         if (target.has_feature(Target::CUDA)) {
             compile(f, true);
         } else {

diff --git a/src/runtime/HalidePyTorchHelpers.h b/src/runtime/HalidePyTorchHelpers.h
@@ -16,7 +16,7 @@
 #include "HalideBuffer.h"
 
 // Forward declare the cuda_device_interface, for tensor wrapper.
-const halide_device_interface_t *halide_cuda_device_interface();
+extern "C" const halide_device_interface_t *halide_cuda_device_interface();
 
 #define HLPT_CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
 #define HLPT_CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")