GPU codegen fixes (#746)

* GPU codegen fixes

* Fix parseCode() API

* Fix parseCode

* Import gpu on @Par(gpu)

* Mark GPU functions as nothrow

---------

Co-authored-by: Ibrahim Numanagić <ibrahimpasa@gmail.com>

Author: arshajii

codon/cir/llvm/gpu.cpp

@@ -187,9 +187,12 @@ void linkLibdevice(llvm::Module *M, const std::string &path) {
   seqassertn(!fail, "linking libdevice failed");
 }
 
-llvm::Function *copyPrototype(llvm::Function *F, const std::string &name) {
+llvm::Function *copyPrototype(llvm::Function *F, const std::string &name,
+                              bool external = false) {
   auto *M = F->getParent();
-  return llvm::Function::Create(F->getFunctionType(), llvm::GlobalValue::PrivateLinkage,
+  return llvm::Function::Create(F->getFunctionType(),
+                                external ? llvm::GlobalValue::ExternalLinkage
+                                         : llvm::GlobalValue::PrivateLinkage,
                                 name.empty() ? F->getName() : name, *M);
 }
 
@@ -651,7 +654,7 @@ void remapFunctions(llvm::Module *M) {
       } else {
         G = M->getFunction(pair.second);
         if (!G)
-          G = copyPrototype(F, pair.second);
+          G = copyPrototype(F, pair.second, /*external=*/true);
       }
 
       G->setWillReturn();
@@ -712,6 +715,12 @@ std::string moduleToPTX(llvm::Module *M, std::vector<llvm::GlobalValue *> &kerne
       llvm::codegen::getExplicitRelocModel(), llvm::codegen::getExplicitCodeModel(),
       llvm::CodeGenOptLevel::Aggressive));
 
+  // Remove personality functions
+  for (auto &F : *M) {
+    F.setDoesNotThrow();
+    F.setPersonalityFn(nullptr);
+  }
+
   M->setDataLayout(machine->createDataLayout());
   auto keep = getRequiredGVs(kernels);
 

codon/cir/transform/parallel/openmp.cpp

@@ -1560,7 +1560,7 @@ void OpenMPPass::handle(ImperativeForFlow *v) {
 
   if (sched->gpu) {
     std::unordered_set<id_t> kernels;
-    const std::string gpuAttr = ast::getMangledFunc("std.internal.gpu", "kernel");
+    const std::string gpuAttr = ast::getMangledFunc(gpuModule, "kernel");
     for (auto *var : *M) {
       if (auto *func = cast<BodiedFunc>(var)) {
         if (util::hasAttribute(func, gpuAttr)) {

codon/parser/cache.cpp

@@ -208,14 +208,18 @@ std::vector<size_t> Cache::getChildRealizationIds(types::ClassType *type) {
   return childIds;
 }
 
-void Cache::parseCode(const std::string &code) {
+std::vector<ir::SeriesFlow *> Cache::parseCode(const std::string &code) {
   auto nodeOrErr = ast::parseCode(this, "<internal>", code, /*startLine=*/0);
-  if (nodeOrErr)
+  if (!nodeOrErr)
     throw exc::ParserException(nodeOrErr.takeError());
   auto sctx = imports[MAIN_IMPORT].ctx;
   auto node = ast::TypecheckVisitor::apply(sctx, *nodeOrErr);
+  auto old = codegenCtx->series;
+  codegenCtx->series.clear();
   ast::TranslateVisitor(codegenCtx).initializeGlobals();
   ast::TranslateVisitor(codegenCtx).translateStmts(node);
+  std::swap(old, codegenCtx->series);
+  return old;
 }
 
 std::vector<std::shared_ptr<types::ClassType>>

codon/parser/cache.h

@@ -331,7 +331,7 @@ struct Cache {
   std::vector<size_t> getBaseRealizationIds(types::ClassType *type);
   std::vector<size_t> getChildRealizationIds(types::ClassType *type);
 
-  void parseCode(const std::string &code);
+  std::vector<ir::SeriesFlow *> parseCode(const std::string &code);
 
   static std::vector<std::shared_ptr<types::ClassType>>
   mergeC3(std::vector<std::vector<types::TypePtr>> &);

codon/parser/visitors/typecheck/loops.cpp

@@ -102,6 +102,20 @@ void TypecheckVisitor::visit(WhileStmt *stmt) {
 void TypecheckVisitor::visit(ForStmt *stmt) {
   stmt->decorator = transformForDecorator(stmt->getDecorator());
 
+  if (auto fc = cast<CallExpr>(stmt->getDecorator())) {
+    if (auto fi = cast<IdExpr>(fc->getExpr());
+        fi && fi->getType()->getFunc() &&
+        fi->getType()->getFunc()->getFuncName() ==
+            getMangledFunc("std.openmp", "for_par")) {
+      if (auto n = extractFuncGeneric(fi->getType(), 3)->getBoolStatic();
+          n && n->value) {
+        prependStmts->push_back(
+            transform(N<ImportStmt>(N<IdExpr>("gpu"), nullptr, std::vector<Param>{},
+                                    nullptr, getTemporaryVar("_"))));
+      }
+    }
+  }
+
   std::string breakVar;
   // Needs in-advance transformation to prevent name clashes with the iterator variable
   stmt->getIter()->setAttribute(

stdlib/internal/gpu.codon

@@ -138,8 +138,10 @@ def nvptx_load_module():
     module = CUmodule()
     # NOTE: 2nd argument to cuModuleLoadData() will
     #       be replaced by Codon's GPU pass
-    cuda_check(cuModuleLoadData(__ptr__(module), __codon_ptx__()))
-    modules.append(module)
+    ptx = __codon_ptx__()
+    if ptx:
+        cuda_check(cuModuleLoadData(__ptr__(module), ptx))
+        modules.append(module)
 
 
 def cuda_init(debug: bool = False):

stdlib/internal/types/str.codon

@@ -144,16 +144,17 @@ class str:
                 n += s.len
             return str(p, n)
         else:
-            total = 0
-            for i in args:
-                if not isinstance(i, str):
+            for i in static.range(static.len(args)):
+                s = args[i]
+                if not isinstance(s, str):
                     compile_error("not a string")
-                total += i.len
+                total += s.len
             p = cobj(total)
             n = 0
-            for i in args:
-                str.memcpy(p + n, i.ptr, i.len)
-                n += i.len
+            for i in static.range(static.len(args)):
+                s = args[i]
+                str.memcpy(p + n, s.ptr, s.len)
+                n += s.len
             return str(p, total)
 
     def __prefix_b__(s: str, N: Literal[int]):

stdlib/numpy/ndarray.codon

@@ -1080,7 +1080,7 @@ class ndarray[dtype, ndim: Literal[int]]:
 
     def _check_order(order: str):
         if order not in ('C', 'F', 'A', 'K'):
-            raise ValueError(f"order must be one of 'C', 'F', 'A', or 'K' (got {repr(order)})")
+            raise ValueError(f"order must be one of 'C', 'F', 'A', or 'K' (got '{order}')")
 
     def astype(self, dtype: type, order: str = 'K', copy: bool = True):
         ndarray._check_order(order)