warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/native/clang/clang.cpp

@@ -25,6 +25,8 @@
 #include <llvm/PassRegistry.h>
 #include <llvm/InitializePasses.h>
 #include <llvm/IR/LegacyPassManager.h>
+#include <llvm/IRReader/IRReader.h>
+#include <llvm/Linker/Linker.h>
 
 #include <llvm/ExecutionEngine/Orc/LLJIT.h>
 #include <llvm/ExecutionEngine/JITEventListener.h>
@@ -45,6 +47,7 @@
 #elif defined(__APPLE__)
     extern "C" void __bzero(void*, size_t);
     extern "C" __double2 __sincos_stret(double);
+    extern "C" __float2 __sincosf_stret(float);
 #endif
 
 extern "C" {
@@ -54,21 +57,20 @@ extern "C" {
 // On Linux it suffices for these symbols not to be stripped out, while for Windows a .pdb has to contain
 // their information. LLVM defines them, but we don't want a huge .pdb with all LLVM source code's debug
 // info. By forward-declaring them here it suffices to compile this file with /Zi.
-struct jit_descriptor;
-extern jit_descriptor __jit_debug_descriptor;
+extern struct jit_descriptor __jit_debug_descriptor;
 extern void __jit_debug_register_code();
 
 }
 
 namespace wp {
-	
+
 #if defined (_WIN32)
-	// Windows defaults to using the COFF binary format (aka. "msvc" in the target triple).
-	// Override it to use the ELF format to support DWARF debug info, but keep using the
-	// Microsoft calling convention (see also https://llvm.org/docs/DebuggingJITedCode.html).
-	static const char* target_triple = "x86_64-pc-windows-elf";
+    // Windows defaults to using the COFF binary format (aka. "msvc" in the target triple).
+    // Override it to use the ELF format to support DWARF debug info, but keep using the
+    // Microsoft calling convention (see also https://llvm.org/docs/DebuggingJITedCode.html).
+    static const char* target_triple = "x86_64-pc-windows-elf";
 #else
-	static const char* target_triple = LLVM_DEFAULT_TARGET_TRIPLE;
+    static const char* target_triple = LLVM_DEFAULT_TARGET_TRIPLE;
 #endif
 
 static void initialize_llvm()
@@ -93,6 +95,11 @@ static std::unique_ptr<llvm::Module> cpp_to_llvm(const std::string& input_file,
     args.push_back("-triple");
     args.push_back(target_triple);
 
+    #if defined(__x86_64__) || defined(_M_X64)
+        args.push_back("-target-feature");
+        args.push_back("+f16c");  // Enables support for _Float16
+    #endif
+
     clang::IntrusiveRefCntPtr<clang::DiagnosticOptions> diagnostic_options = new clang::DiagnosticOptions();
     std::unique_ptr<clang::TextDiagnosticPrinter> text_diagnostic_printer =
             std::make_unique<clang::TextDiagnosticPrinter>(llvm::errs(), &*diagnostic_options);
@@ -114,8 +121,6 @@ static std::unique_ptr<llvm::Module> cpp_to_llvm(const std::string& input_file,
     std::unique_ptr<llvm::MemoryBuffer> buffer = llvm::MemoryBuffer::getMemBufferCopy(cpp_src);
     compiler_invocation.getPreprocessorOpts().addRemappedFile(input_file.c_str(), buffer.get());
 
-    compiler_instance.getPreprocessorOpts().addMacroDef("WP_CPU");
-
     if(!debug)
     {
         compiler_instance.getPreprocessorOpts().addMacroDef("NDEBUG");
@@ -133,18 +138,71 @@ static std::unique_ptr<llvm::Module> cpp_to_llvm(const std::string& input_file,
     return success ? std::move(emit_llvm_only_action.takeModule()) : nullptr;
 }
 
-extern "C" {
-
-WP_API int compile_cpp(const char* cpp_src, const char* include_dir, const char* output_file, bool debug)
+static std::unique_ptr<llvm::Module> cuda_to_llvm(const std::string& input_file, const char* cpp_src, const char* include_dir, bool debug, llvm::LLVMContext& context)
 {
-    #if defined (_WIN32)
-        const char* obj_ext = ".obj";
-    #else
-        const char* obj_ext = ".o";
-    #endif
+    // Compilation arguments
+    std::vector<const char*> args;
+    args.push_back(input_file.c_str());
+
+    args.push_back("-I");
+    args.push_back(include_dir);
+
+    args.push_back(debug ? "-O0" : "-O2");
+
+    args.push_back("-triple");
+    args.push_back("nvptx64-nvidia-cuda");
+
+    args.push_back("-target-cpu");
+    args.push_back("sm_70");
+
+    clang::IntrusiveRefCntPtr<clang::DiagnosticOptions> diagnostic_options = new clang::DiagnosticOptions();
+    std::unique_ptr<clang::TextDiagnosticPrinter> text_diagnostic_printer =
+            std::make_unique<clang::TextDiagnosticPrinter>(llvm::errs(), &*diagnostic_options);
+    clang::IntrusiveRefCntPtr<clang::DiagnosticIDs> diagnostic_ids;
+    std::unique_ptr<clang::DiagnosticsEngine> diagnostic_engine =
+            std::make_unique<clang::DiagnosticsEngine>(diagnostic_ids, &*diagnostic_options, text_diagnostic_printer.release());
+
+    clang::CompilerInstance compiler_instance;
+
+    auto& compiler_invocation = compiler_instance.getInvocation();
+    clang::CompilerInvocation::CreateFromArgs(compiler_invocation, args, *diagnostic_engine.release());
+
+    if(debug)
+    {
+        compiler_invocation.getCodeGenOpts().setDebugInfo(clang::codegenoptions::FullDebugInfo);
+    }
+
+    // Map code to a MemoryBuffer
+    std::unique_ptr<llvm::MemoryBuffer> buffer = llvm::MemoryBuffer::getMemBufferCopy(cpp_src);
+    compiler_invocation.getPreprocessorOpts().addRemappedFile(input_file.c_str(), buffer.get());
+
+    // According to https://llvm.org/docs/CompileCudaWithLLVM.html, "Both clang and nvcc define `__CUDACC__` during CUDA compilation."
+    // But this normally happens in the __clang_cuda_runtime_wrapper.h header, which we don't include.
+    // The __CUDA__ and __CUDA_ARCH__ macros are internally defined by llvm-project/clang/lib/Frontend/InitPreprocessor.cpp
+    compiler_instance.getPreprocessorOpts().addMacroDef("__CUDACC__");
+
+    if(!debug)
+    {
+        compiler_instance.getPreprocessorOpts().addMacroDef("NDEBUG");
+    }
+
+    compiler_instance.getLangOpts().CUDA = 1;
+    compiler_instance.getLangOpts().CUDAIsDevice = 1;
+    compiler_instance.getLangOpts().CUDAAllowVariadicFunctions = 1;
+
+    compiler_instance.createDiagnostics(text_diagnostic_printer.get(), false);
 
-    std::string input_file = std::string(output_file).substr(0, std::strlen(output_file) - std::strlen(obj_ext));
+    clang::EmitLLVMOnlyAction emit_llvm_only_action(&context);
+    bool success = compiler_instance.ExecuteAction(emit_llvm_only_action);
+    buffer.release();
+
+    return success ? std::move(emit_llvm_only_action.takeModule()) : nullptr;
+}
 
+extern "C" {
+
+WP_API int compile_cpp(const char* cpp_src, const char *input_file, const char* include_dir, const char* output_file, bool debug)
+{
     initialize_llvm();
 
     llvm::LLVMContext context;
@@ -155,13 +213,13 @@ WP_API int compile_cpp(const char* cpp_src, const char* include_dir, const char*
         return -1;
     }
 
-    std::string Error;
-    const llvm::Target* target = llvm::TargetRegistry::lookupTarget(target_triple, Error);
+    std::string error;
+    const llvm::Target* target = llvm::TargetRegistry::lookupTarget(target_triple, error);
 
     const char* CPU = "generic";
     const char* features = "";
     llvm::TargetOptions target_options;
-    llvm::Reloc::Model relocation_model = llvm::Reloc::PIC_;  // DLLs need Position Independent Code
+    llvm::Reloc::Model relocation_model = llvm::Reloc::PIC_;  // Position Independent Code
     llvm::CodeModel::Model code_model = llvm::CodeModel::Large;  // Don't make assumptions about displacement sizes
     llvm::TargetMachine* target_machine = target->createTargetMachine(target_triple, CPU, features, target_options, relocation_model, code_model);
 
@@ -182,6 +240,59 @@ WP_API int compile_cpp(const char* cpp_src, const char* include_dir, const char*
     return 0;
 }
 
+WP_API int compile_cuda(const char* cpp_src, const char *input_file, const char* include_dir, const char* output_file, bool debug)
+{
+    initialize_llvm();
+
+    llvm::LLVMContext context;
+    std::unique_ptr<llvm::Module> module = cuda_to_llvm(input_file, cpp_src, include_dir, debug, context);
+
+    if(!module)
+    {
+        return -1;
+    }
+
+    std::string error;
+    const llvm::Target* target = llvm::TargetRegistry::lookupTarget("nvptx64-nvidia-cuda", error);
+
+    const char* CPU = "sm_70";
+    const char* features = "+ptx75";  // Warp requires CUDA 11.5, which supports PTX ISA 7.5
+    llvm::TargetOptions target_options;
+    llvm::Reloc::Model relocation_model = llvm::Reloc::PIC_;
+    llvm::TargetMachine* target_machine = target->createTargetMachine("nvptx64-nvidia-cuda", CPU, features, target_options, relocation_model);
+
+    module->setDataLayout(target_machine->createDataLayout());
+
+    // Link libdevice
+    llvm::SMDiagnostic diagnostic;
+    std::string libdevice_path = std::string(include_dir) + "/libdevice/libdevice.10.bc";
+    std::unique_ptr<llvm::Module> libdevice(llvm::parseIRFile(libdevice_path, diagnostic, context));
+    if(!libdevice)
+    {
+        return -1;
+    }
+
+    llvm::Linker linker(*module.get());
+    if(linker.linkInModule(std::move(libdevice), llvm::Linker::Flags::LinkOnlyNeeded) == true)
+    {
+        return -1;
+    }
+
+    std::error_code error_code;
+    llvm::raw_fd_ostream output(output_file, error_code, llvm::sys::fs::OF_None);
+
+    llvm::legacy::PassManager pass_manager;
+    llvm::CodeGenFileType file_type = llvm::CGFT_AssemblyFile;
+    target_machine->addPassesToEmitFile(pass_manager, output, nullptr, file_type);
+
+    pass_manager.run(*module);
+    output.flush();
+
+    delete target_machine;
+
+    return 0;
+}
+
 // Global JIT instance
 static llvm::orc::LLJIT* jit = nullptr;
 
@@ -248,6 +359,7 @@ WP_API int load_obj(const char* object_file, const char* module_name)
             SYMBOL(log10f), SYMBOL_T(log10, double(*)(double)),
             SYMBOL(expf), SYMBOL_T(exp, double(*)(double)),
             SYMBOL(sqrtf), SYMBOL_T(sqrt, double(*)(double)),
+            SYMBOL(cbrtf), SYMBOL_T(cbrt, double(*)(double)),
             SYMBOL(powf), SYMBOL_T(pow, double(*)(double, double)),
             SYMBOL(floorf), SYMBOL_T(floor, double(*)(double)),
             SYMBOL(ceilf), SYMBOL_T(ceil, double(*)(double)),
@@ -276,7 +388,7 @@ WP_API int load_obj(const char* object_file, const char* module_name)
             SYMBOL(__chkstk),
         #elif defined(__APPLE__)
             SYMBOL(__bzero),
-            SYMBOL(__sincos_stret),
+            SYMBOL(__sincos_stret), SYMBOL(__sincosf_stret),
         #else
             SYMBOL(sincosf), SYMBOL_T(sincos, void(*)(double,double*,double*)),
         #endif
@@ -335,7 +447,7 @@ WP_API uint64_t lookup(const char* dll_name, const char* function_name)
     if(!func)
     {
         std::cerr << "Failed to lookup symbol: " << llvm::toString(func.takeError()) << std::endl;
-        return -1;
+        return 0;
     }
 
     return func->getValue();

warp/native/crt.cpp

@@ -29,79 +29,4 @@ extern "C" WP_API void _wp_assert(const char* expression, const char* file, unsi
     // Now invoke the standard assert(), which may abort the program or break
     // into the debugger as decided by the runtime environment.
     assert(false && "assert() failed");
-}
-
-// Export CRT symbols from warp.dll for use by compute kernel DLLs
-// These are declared in crt.h
-#if defined(_MSC_VER)
-
-#pragma comment(linker,"/export:printf")
-
-#pragma comment(linker,"/export:abs")
-#pragma comment(linker,"/export:llabs")
-
-#pragma comment(linker,"/export:fmodf")
-#pragma comment(linker,"/export:fmod")
-#pragma comment(linker,"/export:logf")
-#pragma comment(linker,"/export:log")
-#pragma comment(linker,"/export:log2f")
-#pragma comment(linker,"/export:log2")
-#pragma comment(linker,"/export:log10f")
-#pragma comment(linker,"/export:log10")
-#pragma comment(linker,"/export:expf")
-#pragma comment(linker,"/export:exp")
-#pragma comment(linker,"/export:sqrtf")
-#pragma comment(linker,"/export:sqrt")
-#pragma comment(linker,"/export:powf")
-#pragma comment(linker,"/export:pow")
-#pragma comment(linker,"/export:floorf")
-#pragma comment(linker,"/export:floor")
-#pragma comment(linker,"/export:ceilf")
-#pragma comment(linker,"/export:ceil")
-#pragma comment(linker,"/export:fabsf")
-#pragma comment(linker,"/export:fabs")
-#pragma comment(linker,"/export:roundf")
-#pragma comment(linker,"/export:round")
-#pragma comment(linker,"/export:truncf")
-#pragma comment(linker,"/export:trunc")
-#pragma comment(linker,"/export:rintf")
-#pragma comment(linker,"/export:rint")
-#pragma comment(linker,"/export:acosf")
-#pragma comment(linker,"/export:acos")
-#pragma comment(linker,"/export:asinf")
-#pragma comment(linker,"/export:asin")
-#pragma comment(linker,"/export:atanf")
-#pragma comment(linker,"/export:atan")
-#pragma comment(linker,"/export:atan2f")
-#pragma comment(linker,"/export:atan2")
-#pragma comment(linker,"/export:cosf")
-#pragma comment(linker,"/export:cos")
-#pragma comment(linker,"/export:sinf")
-#pragma comment(linker,"/export:sin")
-#pragma comment(linker,"/export:tanf")
-#pragma comment(linker,"/export:tan")
-#pragma comment(linker,"/export:sinhf")
-#pragma comment(linker,"/export:sinh")
-#pragma comment(linker,"/export:coshf")
-#pragma comment(linker,"/export:cosh")
-#pragma comment(linker,"/export:tanhf")
-#pragma comment(linker,"/export:tanh")
-#pragma comment(linker,"/export:fmaf")
-
-#pragma comment(linker,"/export:memset")
-#pragma comment(linker,"/export:memcpy")
-
-#pragma comment(linker,"/export:_wp_isfinite")
-#pragma comment(linker,"/export:_wp_assert")
-
-// For functions with large stack frames the MSVC compiler will emit a call to
-// __chkstk() to linearly touch each memory page. This grows the stack without
-// triggering the stack overflow guards.
-#pragma comment(linker,"/export:__chkstk")
-
-// The MSVC linker checks for the _fltused symbol if any floating-point
-// functionality is used. It's defined by the Microsoft CRT to indicate that
-// the x87 FPU control word was properly initialized.
-#pragma comment(linker,"/export:_fltused")
-
-#endif  // _MSC_VER
+}

warp/native/crt.h

@@ -30,15 +30,15 @@
     #define WP_API
 #endif
 
-extern "C" {
+#if !defined(__CUDA_ARCH__)
 
 // Helper for implementing assert() macro
-WP_API void _wp_assert(const char* message, const char* file, unsigned int line);
+extern "C" WP_API void _wp_assert(const char* message, const char* file, unsigned int line);
 
 // Helper for implementing isfinite()
-WP_API int _wp_isfinite(double);
+extern "C" WP_API int _wp_isfinite(double);
 
-}  // extern "C"
+#endif  // !__CUDA_ARCH__
 
 #if !defined(WP_NO_CRT)
 
@@ -52,106 +52,6 @@ WP_API int _wp_isfinite(double);
 
 #else
 
-#if defined(__CUDACC__)
-
-// stdio.h
-extern "C" __device__ int printf(const char* format, ... );
-
-#else
-
-extern "C" {
-
-// stdio.h
-int printf(const char * format, ... );
-
-// stdlib.h
-int abs(int);
-long long llabs(long long);
-
-// math.h
-float fmodf(float, float);
-double fmod(double, double);
-float logf(float);
-double log(double);
-float log2f(float);
-double log2(double);
-float log10f(float);
-double log10(double);
-float expf(float);
-double exp(double);
-float sqrtf(float);
-double sqrt(double);
-float powf(float, float);
-double pow(double, double);
-float floorf(float);
-double floor(double);
-float ceilf(float);
-double ceil(double);
-float fabsf(float);
-double fabs(double);
-float roundf(float);
-double round(double);
-float truncf(float);
-double trunc(double);
-float rintf(float);
-double rint(double);
-float acosf(float);
-double acos(double);
-float asinf(float);
-double asin(double);
-float atanf(float);
-double atan(double);
-float atan2f(float, float);
-double atan2(double, double);
-float cosf(float);
-double cos(double);
-float sinf(float);
-double sin(double);
-float tanf(float);
-double tan(double);
-float sinhf(float);
-double sinh(double);
-float coshf(float);
-double cosh(double);
-float tanhf(float);
-double tanh(double);
-float fmaf(float, float, float);
-
-// stddef.h
-#if defined(_WIN32)
-using size_t = unsigned __int64;
-#else
-using size_t = unsigned long;
-#endif
-
-// string.h
-void* memset(void*, int, size_t);
-void* memcpy(void*, const void*, size_t);
-
-// stdlib.h
-void* malloc(size_t);
-void free(void*);
-
-}  // extern "C"
-
-// cmath
-inline bool isfinite(double x)
-{
-    return _wp_isfinite(x);
-}
-
-// assert.h
-#ifdef NDEBUG
-    #define assert(expression) ((void)0)
-#else
-    #define assert(expression) (void)(                                    \
-            (!!(expression)) ||                                           \
-            (_wp_assert((#expression), (__FILE__), (unsigned)(__LINE__)), 0) \
-        )
-#endif
-
-#endif  // !__CUDACC__
-
 // These definitions are taken from Jitify: https://github.com/NVIDIA/jitify
 
 /// float.h
@@ -221,6 +121,9 @@ enum {
 #define LLONG_MIN  (-LLONG_MAX - 1LL)
 #define ULLONG_MAX 18446744073709551615ULL
 
+#define INFINITY   ((float)(DBL_MAX * DBL_MAX))
+#define HUGE_VAL   ((double)INFINITY)
+#define HUGE_VALF  ((float)INFINITY)
 
 /// stdint.h
 typedef signed char      int8_t;
@@ -325,4 +228,108 @@ typedef unsigned long long uint64_t;
 
 #define M_PI 3.14159265358979323846
 
+#if defined(__CUDACC__)
+
+#if defined(__clang__)
+// When compiling CUDA with barebones Clang we need to define its builtins and runtime functions ourselves.
+#include "cuda_crt.h"
+#endif
+
+#else
+
+extern "C" {
+
+// stdio.h
+int printf(const char * format, ... );
+
+// stdlib.h
+int abs(int);
+long long llabs(long long);
+
+// math.h
+float fmodf(float, float);
+double fmod(double, double);
+float logf(float);
+double log(double);
+float log2f(float);
+double log2(double);
+float log10f(float);
+double log10(double);
+float expf(float);
+double exp(double);
+float sqrtf(float);
+double sqrt(double);
+float cbrtf(float);
+double cbrt(double);
+float powf(float, float);
+double pow(double, double);
+float floorf(float);
+double floor(double);
+float ceilf(float);
+double ceil(double);
+float fabsf(float);
+double fabs(double);
+float roundf(float);
+double round(double);
+float truncf(float);
+double trunc(double);
+float rintf(float);
+double rint(double);
+float acosf(float);
+double acos(double);
+float asinf(float);
+double asin(double);
+float atanf(float);
+double atan(double);
+float atan2f(float, float);
+double atan2(double, double);
+float cosf(float);
+double cos(double);
+float sinf(float);
+double sin(double);
+float tanf(float);
+double tan(double);
+float sinhf(float);
+double sinh(double);
+float coshf(float);
+double cosh(double);
+float tanhf(float);
+double tanh(double);
+float fmaf(float, float, float);
+
+// stddef.h
+#if defined(_WIN32)
+using size_t = unsigned __int64;
+#else
+using size_t = unsigned long;
+#endif
+
+// string.h
+void* memset(void*, int, size_t);
+void* memcpy(void*, const void*, size_t);
+
+// stdlib.h
+void* malloc(size_t);
+void free(void*);
+
+}  // extern "C"
+
+// cmath
+inline bool isfinite(double x)
+{
+    return _wp_isfinite(x);
+}
+
+// assert.h
+#ifdef NDEBUG
+    #define assert(expression) ((void)0)
+#else
+    #define assert(expression) (void)(                                    \
+            (!!(expression)) ||                                           \
+            (_wp_assert((#expression), (__FILE__), (unsigned)(__LINE__)), 0) \
+        )
+#endif
+
+#endif  // !__CUDACC__
+
 #endif // WP_NO_CRT