warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/native/clang/clang.cpp

@@ -16,6 +16,7 @@
  */
 
 #include "../native/crt.h"
+#include "../version.h"
 
 #include <clang/Frontend/CompilerInstance.h>
 #include <clang/Basic/DiagnosticOptions.h>
@@ -58,27 +59,14 @@
 #if defined(_WIN64)
     extern "C" void __chkstk();
 #elif defined(__APPLE__)
-
-#if defined(__MACH__) && defined(__aarch64__)
     extern "C" void _bzero(void *s, size_t n) {
         memset(s, 0, n);
     }
     extern "C" void __bzero(void *s, size_t n) {
         memset(s, 0, n);
     }
-
-    extern "C" void _memset_pattern16(void *s, const void *pattern, size_t n);
-    extern "C" void __memset_pattern16(void *s, const void *pattern, size_t n);
-
-#else
-    // // Intel Mac's define bzero in libSystem.dylib
-    extern "C" void __bzero(void *s, size_t n);
-
     extern "C" void _memset_pattern16(void *s, const void *pattern, size_t n);
     extern "C" void __memset_pattern16(void *s, const void *pattern, size_t n);
-
-#endif
-
     extern "C" __double2 __sincos_stret(double);
     extern "C" __float2 __sincosf_stret(float);
 #endif // defined(__APPLE__)
@@ -114,7 +102,7 @@ static void initialize_llvm()
     llvm::InitializeAllAsmPrinters();
 }
 
-static std::unique_ptr<llvm::Module> cpp_to_llvm(const std::string& input_file, const char* cpp_src, const char* include_dir, bool debug, bool verify_fp, llvm::LLVMContext& context)
+static std::unique_ptr<llvm::Module> source_to_llvm(bool is_cuda, const std::string& input_file, const char* cpp_src, const char* include_dir, bool debug, bool verify_fp, llvm::LLVMContext& context, bool tiles_in_stack_memory)
 {
     // Compilation arguments
     std::vector<const char*> args;
@@ -125,84 +113,50 @@ static std::unique_ptr<llvm::Module> cpp_to_llvm(const std::string& input_file,
 
     args.push_back(debug ? "-O0" : "-O2");
 
-    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);
-    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)
+    if(is_cuda)
     {
-        #if LLVM_VERSION_MAJOR >= 18
-        compiler_invocation.getCodeGenOpts().setDebugInfo(llvm::codegenoptions::FullDebugInfo);
-        #else
-        compiler_invocation.getCodeGenOpts().setDebugInfo(clang::codegenoptions::FullDebugInfo);
-        #endif
-    }
+        args.push_back("-triple");
+        args.push_back("nvptx64-nvidia-cuda");
 
-    // 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());
-
-    if(!debug)
-    {
-        compiler_instance.getPreprocessorOpts().addMacroDef("NDEBUG");
+        args.push_back("-target-cpu");
+        args.push_back("sm_70");
     }
-
-    if(verify_fp)
+    else
     {
-        compiler_instance.getPreprocessorOpts().addMacroDef("WP_VERIFY_FP");
-    }
-
-    compiler_instance.getLangOpts().MicrosoftExt = 1;  // __forceinline / __int64
-    compiler_instance.getLangOpts().DeclSpecKeyword = 1;  // __declspec
-
-    compiler_instance.createDiagnostics(text_diagnostic_printer.get(), false);
+        args.push_back("-triple");
+        args.push_back(target_triple);
 
-    clang::EmitLLVMOnlyAction emit_llvm_only_action(&context);
-    bool success = compiler_instance.ExecuteAction(emit_llvm_only_action);
-    (void)buffer.release();
-
-    return success ? std::move(emit_llvm_only_action.takeModule()) : nullptr;
-}
-
-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)
-{
-    // 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");
+        #if defined(__x86_64__) || defined(_M_X64)
+            args.push_back("-target-feature");
+            args.push_back("+f16c");  // Enables support for _Float16
+        #endif
 
-    args.push_back("-target-cpu");
-    args.push_back("sm_70");
+        #if defined(__aarch64__)
+        if(tiles_in_stack_memory)
+        {
+            // Static memory support is broken on AArch64 CPUs. As a workaround we reserve some stack memory on kernel entry,
+            // and point the callee-saved x28 register to it so we can access it anywhere. See tile_shared_storage_t in tile.h.
+            args.push_back("-target-feature");
+            args.push_back("+reserve-x28");
+        }
+        #endif
+    }
 
+    #if LLVM_VERSION_MAJOR >= 21
+    clang::DiagnosticOptions diagnostic_options;
+    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());
+    #else
     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());
+    #endif
 
     clang::CompilerInstance compiler_instance;
 
@@ -222,21 +176,43 @@ static std::unique_ptr<llvm::Module> cuda_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());
 
-    // 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");
     }
+    
+    if(is_cuda)
+    {
+        // 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__");
+
+        compiler_instance.getLangOpts().CUDA = 1;
+        compiler_instance.getLangOpts().CUDAIsDevice = 1;
+        compiler_instance.getLangOpts().CUDAAllowVariadicFunctions = 1;
+    }
+    else
+    {
+        if(verify_fp)
+        {
+            compiler_instance.getPreprocessorOpts().addMacroDef("WP_VERIFY_FP");
+        }
 
-    compiler_instance.getLangOpts().CUDA = 1;
-    compiler_instance.getLangOpts().CUDAIsDevice = 1;
-    compiler_instance.getLangOpts().CUDAAllowVariadicFunctions = 1;
+        if(tiles_in_stack_memory)
+        {
+            compiler_instance.getPreprocessorOpts().addMacroDef("WP_ENABLE_TILES_IN_STACK_MEMORY");
+        }
+
+        compiler_instance.getLangOpts().MicrosoftExt = 1;  // __forceinline / __int64
+        compiler_instance.getLangOpts().DeclSpecKeyword = 1;  // __declspec
+    }
 
+    #if LLVM_VERSION_MAJOR >= 21
+    compiler_instance.createDiagnostics(*llvm::vfs::getRealFileSystem(), text_diagnostic_printer.get(), false);
+    #else
     compiler_instance.createDiagnostics(text_diagnostic_printer.get(), false);
+    #endif
 
     clang::EmitLLVMOnlyAction emit_llvm_only_action(&context);
     bool success = compiler_instance.ExecuteAction(emit_llvm_only_action);
@@ -247,12 +223,12 @@ static std::unique_ptr<llvm::Module> cuda_to_llvm(const std::string& input_file,
 
 extern "C" {
 
-WP_API int wp_compile_cpp(const char* cpp_src, const char *input_file, const char* include_dir, const char* output_file, bool debug, bool verify_fp, bool fuse_fp)
+WP_API int wp_compile_cpp(const char* cpp_src, const char *input_file, const char* include_dir, const char* output_file, bool debug, bool verify_fp, bool fuse_fp, bool tiles_in_stack_memory)
 {
     initialize_llvm();
 
     llvm::LLVMContext context;
-    std::unique_ptr<llvm::Module> module = cpp_to_llvm(input_file, cpp_src, include_dir, debug, verify_fp, context);
+    std::unique_ptr<llvm::Module> module = source_to_llvm(false, input_file, cpp_src, include_dir, debug, verify_fp, context, tiles_in_stack_memory);
 
     if(!module)
     {
@@ -260,7 +236,11 @@ WP_API int wp_compile_cpp(const char* cpp_src, const char *input_file, const cha
     }
 
     std::string error;
+     #if LLVM_VERSION_MAJOR >= 22
+    const llvm::Target* target = llvm::TargetRegistry::lookupTarget(llvm::Triple(target_triple), error);
+    #else
     const llvm::Target* target = llvm::TargetRegistry::lookupTarget(target_triple, error);
+    #endif
 
     const char* CPU = "generic";
     const char* features = "";
@@ -271,7 +251,11 @@ WP_API int wp_compile_cpp(const char* cpp_src, const char *input_file, const cha
         target_options.AllowFPOpFusion = llvm::FPOpFusion::Strict;
     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
+    #if LLVM_VERSION_MAJOR >= 20
+    llvm::TargetMachine* target_machine = target->createTargetMachine(llvm::Triple(target_triple), CPU, features, target_options, relocation_model, code_model);
+    #else
     llvm::TargetMachine* target_machine = target->createTargetMachine(target_triple, CPU, features, target_options, relocation_model, code_model);
+    #endif
 
     module->setDataLayout(target_machine->createDataLayout());
 
@@ -299,7 +283,7 @@ WP_API int wp_compile_cuda(const char* cpp_src, const char *input_file, const ch
     initialize_llvm();
 
     llvm::LLVMContext context;
-    std::unique_ptr<llvm::Module> module = cuda_to_llvm(input_file, cpp_src, include_dir, debug, context);
+    std::unique_ptr<llvm::Module> module = source_to_llvm(true, input_file, cpp_src, include_dir, debug, false, context, false);
 
     if(!module)
     {
@@ -307,13 +291,22 @@ WP_API int wp_compile_cuda(const char* cpp_src, const char *input_file, const ch
     }
 
     std::string error;
+
+    #if LLVM_VERSION_MAJOR >= 22
+    const llvm::Target* target = llvm::TargetRegistry::lookupTarget(llvm::Triple("nvptx64-nvidia-cuda"), error);
+    #else
     const llvm::Target* target = llvm::TargetRegistry::lookupTarget("nvptx64-nvidia-cuda", error);
+    #endif
 
     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_;
+    #if LLVM_VERSION_MAJOR >= 20
+    llvm::TargetMachine* target_machine = target->createTargetMachine(llvm::Triple("nvptx64-nvidia-cuda"), CPU, features, target_options, relocation_model);
+    #else
     llvm::TargetMachine* target_machine = target->createTargetMachine("nvptx64-nvidia-cuda", CPU, features, target_options, relocation_model);
+    #endif
 
     module->setDataLayout(target_machine->createDataLayout());
 
@@ -363,8 +356,16 @@ WP_API int wp_load_obj(const char* object_file, const char* module_name)
 
         auto jit_expected = llvm::orc::LLJITBuilder()
             .setObjectLinkingLayerCreator(
+#if LLVM_VERSION_MAJOR >= 21
+                [&](llvm::orc::ExecutionSession &session) {
+#else
                 [&](llvm::orc::ExecutionSession &session, const llvm::Triple &triple) {
+#endif
+#if LLVM_VERSION_MAJOR >= 21
+                    auto get_memory_manager = [](const llvm::MemoryBuffer &) {
+#else
                     auto get_memory_manager = []() {
+#endif
                         return std::make_unique<llvm::SectionMemoryManager>();
                     };
                     auto obj_linking_layer = std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(session, std::move(get_memory_manager));
@@ -443,6 +444,10 @@ WP_API int wp_load_obj(const char* object_file, const char* module_name)
             SYMBOL(coshf), SYMBOL_T(cosh, double(*)(double)),
             SYMBOL(tanhf), SYMBOL_T(tanh, double(*)(double)),
             SYMBOL(fmaf), SYMBOL_T(fma, double(*)(double, double, double)),
+            SYMBOL(erff), SYMBOL_T(erf, double(*)(double)),
+            SYMBOL(erfcf), SYMBOL_T(erfc, double(*)(double)),
+            SYMBOL(erfinvf), SYMBOL_T(erfinv, double(*)(double)),
+            SYMBOL(erfcinvf), SYMBOL_T(erfcinv, double(*)(double)),
             SYMBOL(memcpy), SYMBOL(memset), SYMBOL(memmove),
             SYMBOL(_wp_assert),
             SYMBOL(_wp_isfinite),
@@ -454,13 +459,8 @@ WP_API int wp_load_obj(const char* object_file, const char* module_name)
             // triggering the stack overflow guards.
             SYMBOL(__chkstk),
         #elif defined(__APPLE__)
-            #if defined(__MACH__) && defined(__aarch64__)
-                SYMBOL(bzero),
-                SYMBOL(_bzero),
-            #else
-                // Intel Mac
-                SYMBOL(__bzero),
-            #endif
+            SYMBOL(bzero),
+            SYMBOL(_bzero),
             SYMBOL(memset_pattern16),
             SYMBOL(__sincos_stret), SYMBOL(__sincosf_stret),
         #else
@@ -531,6 +531,11 @@ WP_API uint64_t wp_lookup(const char* dll_name, const char* function_name)
     return func->getValue();
 }
 
+WP_API const char* wp_warp_clang_version()
+{
+    return WP_VERSION_STRING;
+}
+
 }  // extern "C"
 
 }  // namespace wp

warp/native/coloring.cpp

@@ -346,7 +346,6 @@ public:
         return node_weights[node_idx];
     }
 
-
     void add_node(int weight, int node_idx)
     {
         if (weight >= weight_buckets.size()) 

warp/native/crt.h

@@ -311,6 +311,14 @@ float tanhf(float);
 double tanh(double);
 float fmaf(float, float, float);
 double fma(double, double, double);
+double erf(double);
+float erff(float);
+double erfc(double);
+float erfcf(float);
+double erfinv(double);
+float erfinvf(float);
+double erfcinv(double);
+float erfcinvf(float);
 
 // stddef.h
 #if defined(_WIN32)
@@ -358,3 +366,203 @@ inline bool isinf(double x)
 #endif  // !__CUDACC__
 
 #endif // WP_NO_CRT
+
+#if !defined(__CUDACC__)
+
+/*
+ * From Cephes Library polevl.c
+ * Original source: https://www.netlib.org/cephes/
+ * Copyright (c) 1984 by Stephen L. Moshier.
+ * All rights reserved.
+*/
+// evaluate polynomial using Horner's method
+static inline double polevl(double x, const double* coefs, int N)
+{
+    double ans = coefs[0];
+    for (int i = 1; i <= N; i++)
+    {
+        ans = ans * x + coefs[i];
+    }
+    return ans;
+}
+
+/*
+ * From Cephes Library polevl.c
+ * Original source: https://www.netlib.org/cephes/
+ * Copyright (c) 1984 by Stephen L. Moshier.
+ * All rights reserved.
+*/
+// evaluate polynomial assuming leading coef = 1, using Horner's method
+static inline double p1evl(double x, const double* coefs, int N)
+{
+    double ans = x + coefs[0];
+    for (int i = 1; i < N; i++)
+    {
+        ans = ans * x + coefs[i];
+    }
+    return ans;
+}
+
+/*
+ * From Cephes Library ndtri.c
+ * Original source: https://www.netlib.org/cephes/
+ * Copyright (c) 1984 by Stephen L. Moshier.
+ * All rights reserved.
+*/
+// inverse normal distribution function (ndtri)
+static inline double ndtri(double y)
+{
+    // domain check
+    if (y <= 0.0 || y >= 1.0)
+    {
+        return (y <= 0.0) ? -HUGE_VAL : HUGE_VAL;
+    }
+    
+    // constants from Cephes
+    const double s2pi = 2.50662827463100050242E0;  // sqrt(2*pi)
+    const double exp_neg2 = 0.13533528323661269189;  // exp(-2)
+    
+    // approximation for 0 <= abs(z - 0.5) <= 3/8
+    static const double P0[5] = {
+        -5.99633501014107895267e1,
+        9.80010754185999661536e1,
+        -5.66762857469070293439e1,
+        1.39312609387279679503e1,
+        -1.23916583867381258016e0
+    };
+    
+    static const double Q0[8] = {
+        1.95448858338141759834e0,
+        4.67627912898881538453e0,
+        8.63602421390890590575e1,
+        -2.25462687854119370527e2,
+        2.00260212380060660359e2,
+        -8.20372256168333339912e1,
+        1.59056225126211695515e1,
+        -1.18331621121330003142e0
+    };
+    
+    // approximation for interval z = sqrt(-2 log y) between 2 and 8
+    static const double P1[9] = {
+        4.05544892305962419923e0,
+        3.15251094599893866154e1,
+        5.71628192246421288162e1,
+        4.40805073893200834700e1,
+        1.46849561928858024014e1,
+        2.18663306850790267539e0,
+        -1.40256079171354495875e-1,
+        -3.50424626827848203418e-2,
+        -8.57456785154685413611e-4
+    };
+    
+    static const double Q1[8] = {
+        1.57799883256466749731e1,
+        4.53907635128879210584e1,
+        4.13172038254672030440e1,
+        1.50425385692907503408e1,
+        2.50464946208309415979e0,
+        -1.42182922854787788574e-1,
+        -3.80806407691578277194e-2,
+        -9.33259480895457427372e-4
+    };
+    
+    // approximation for interval z = sqrt(-2 log y) between 8 and 64
+    static const double P2[9] = {
+        3.23774891776946035970e0,
+        6.91522889068984211695e0,
+        3.93881025292474443415e0,
+        1.33303460815807542389e0,
+        2.01485389549179081538e-1,
+        1.23716634817820021358e-2,
+        3.01581553508235416007e-4,
+        2.65806974686737550832e-6,
+        6.23974539184983293730e-9
+    };
+    
+    static const double Q2[8] = {
+        6.02427039364742014255e0,
+        3.67983563856160859403e0,
+        1.37702099489081330271e0,
+        2.16236993594496635890e-1,
+        1.34204006088543189037e-2,
+        3.28014464682127739104e-4,
+        2.89247864745380683936e-6,
+        6.79019408009981274425e-9
+    };
+    
+    int code = 1;
+    double y_work = y;
+    
+    if (y_work > (1.0 - exp_neg2))
+    {
+        y_work = 1.0 - y_work;
+        code = 0;
+    }
+    
+    // middle region: 0 <= |y - 0.5| <= 3/8
+    if (y_work > exp_neg2)
+    {
+        y_work -= 0.5;
+        double y2 = y_work * y_work;
+        double x = y_work + y_work * (y2 * polevl(y2, P0, 4) / p1evl(y2, Q0, 8));
+        x = x * s2pi;
+        return x;
+    }
+    
+    double x = ::sqrt(-2.0 * ::log(y_work));
+    double x0 = x - ::log(x) / x;
+    
+    double z = 1.0 / x;
+    double x1;
+    if (x < 8.0)
+    {
+        x1 = z * polevl(z, P1, 8) / p1evl(z, Q1, 8);
+    }
+    else
+    {
+        x1 = z * polevl(z, P2, 8) / p1evl(z, Q2, 8);
+    }
+    
+    x = x0 - x1;
+    if (code != 0)
+    {
+        x = -x;
+    }
+    
+    return x;
+}
+
+// inverse error function (not in standard C library)
+// only compiled for non-CUDA builds - CUDA provides these in its math headers
+inline double erfinv(double z)
+{
+    // handle special cases
+    if (z == 0.0)
+        return 0.0;
+    if (z == 1.0)
+        return HUGE_VAL;  // infinity
+    if (z == -1.0)
+        return -HUGE_VAL;  // -infinity
+    if (z < -1.0 || z > 1.0)
+        return NAN;  // outside valid range
+    
+    // erfinv(z) = ndtri((z + 1) / 2) / sqrt(2)
+    return ndtri((z + 1.0) / 2.0) / ::sqrt(2.0);
+}
+
+inline float erfinvf(float x)
+{
+    return (float)erfinv((double)x);
+}
+
+inline double erfcinv(double x)
+{
+    return erfinv(1.0 - x);
+}
+
+inline float erfcinvf(float x)
+{
+    return (float)erfcinv((double)x);
+}
+
+#endif  // !defined(__CUDACC__)