Changes

Author: DrTodd13

numba/npyufunc/csa.py

@@ -10,6 +10,7 @@
 
 import pdb
 import copy
+import sys
 
 class CSATypingContext(typing.BaseContext):
     def load_additional_registries(self):
@@ -51,6 +52,7 @@ class CSATargetDesc(TargetDescriptor):
 def compile_csa(func_ir, return_type, args, inflags):
     if config.DEBUG_CSA:
         print("compile_csa", func_ir, return_type, args)
+        sys.stdout.flush()
 
     cput = registry.dispatcher_registry['cpu'].targetdescr 
     typingctx = cput.typing_context
@@ -85,11 +87,15 @@ def compile_csa(func_ir, return_type, args, inflags):
                                flags,
                                locals={})
     library = cres.library
+    if config.DEBUG_CSA:
+        print("library", library, type(library))
+        sys.stdout.flush()
     library.finalize()
 
     if config.DEBUG_CSA:
         print("compile_csa cres", cres, type(cres))
         print("LLVM")
+        sys.stdout.flush()
 
         llvm_str = cres.library.get_llvm_str()
         llvm_out = "compile_csa" + ".ll"
@@ -191,6 +197,7 @@ def compile_csa_kernel(func_ir, args, flags, link, fastmath=False):
         print("lib", lib, type(lib))
         print("kernel", kernel, type(kernel))
         print("wrapfnty", wrapfnty, type(wrapfnty))
+        sys.stdout.flush()
     csakern = CSAKernel(llvm_module=lib._final_module,
                         library=wrapper_library,
                         wrapper_module=wrapper_library._final_module,
@@ -228,6 +235,9 @@ def compile(self, sig):
         if kernel is None:
             if 'link' not in self.targetoptions:
                 self.targetoptions['link'] = ()
+            if config.DEBUG_CSA:
+                print("Before compile_csa_kernel.")
+                sys.stdout.flush()
             kernel = compile_csa_kernel(self.func_ir, argtypes, self.flags,
                                     **self.targetoptions)
             self.definitions[argtypes] = kernel

numba/npyufunc/parfor.py

@@ -218,7 +218,7 @@ def lower(self, lowerer):
         builder = lowerer.builder
         library = lowerer.library
 
-        assert self.start_spm.prs_var != None
+        assert self.start_spmd.prs_var != None
 
         llvm_int32_t = lc.Type.int(32)
         fnty = lir.FunctionType(lc.Type.void(), [llvm_int32_t])
@@ -1123,9 +1123,9 @@ def _create_gufunc_for_parfor_body(
         if config.DEBUG_CSA:
             print("eachdim", eachdim)
         if targetctx.auto_parallel.csa:
-            for indent in range(eachdim + 1):
-                gufunc_txt += "    "
-            gufunc_txt += "numba.npyufunc.parfor.par_start_spmd(" + str(eachdim) + ", 4)\n"
+#            for indent in range(eachdim + 1):
+#                gufunc_txt += "    "
+#            gufunc_txt += "numba.npyufunc.parfor.par_start_spmd(" + str(eachdim) + ", 4)\n"
             for indent in range(eachdim + 1):
                 gufunc_txt += "    "
             gufunc_txt += "numba.npyufunc.parfor.par_start_region(" + str(eachdim) + ")\n"
@@ -1179,9 +1179,9 @@ def _create_gufunc_for_parfor_body(
             for indent in range(eachdim + 1):
                 gufunc_txt += "    "
             gufunc_txt += "numba.npyufunc.parfor.par_end_region(" + str(eachdim) + ")\n"
-            for indent in range(eachdim + 1):
-                gufunc_txt += "    "
-            gufunc_txt += "numba.npyufunc.parfor.par_end_spmd(" + str(eachdim) + ")\n"
+#            for indent in range(eachdim + 1):
+#                gufunc_txt += "    "
+#            gufunc_txt += "numba.npyufunc.parfor.par_end_spmd(" + str(eachdim) + ")\n"
 
 
     # Add assignments of reduction variables (for returning the value)
@@ -1285,7 +1285,7 @@ def _create_gufunc_for_parfor_body(
                             elif callname[0] == 'par_start_spmd':
                                 pss_loc = gufunc_ir._definitions[rhs.args[0].name][0].value
                                 pnum_threads = gufunc_ir._definitions[rhs.args[1].name][0].value
-                                pspmd_dict[pss_loc] = parallel_spmd_start(loc, pnum_threads)
+                                pspmd_dict[pss_loc] = parallel_spmd_start(pnum_threads, loc)
                                 new_block.append(pspmd_dict[pss_loc])
                                 continue
                             elif callname[0] == 'par_end_spmd':
@@ -1428,6 +1428,10 @@ def _create_gufunc_for_parfor_body(
     gufunc_ir.blocks = rename_labels(gufunc_ir.blocks)
     remove_dels(gufunc_ir.blocks)
 
+    if config.DEBUG_ARRAY_OPT:
+        print("flush")
+        sys.stdout.flush()
+
     if config.DEBUG_ARRAY_OPT:
         print("gufunc_ir last dump")
         gufunc_ir.dump()
@@ -1450,6 +1454,8 @@ def _create_gufunc_for_parfor_body(
 
     if config.DEBUG_CSA:
         print("Before compile gufunc.")
+    if config.DEBUG_ARRAY_OPT:
+        sys.stdout.flush()
     if targetctx.auto_parallel.csa:
         ajck = csa.AutoJitCSAKernel(gufunc_ir, True, flags, {})
         # Returns CSAKernel
@@ -1458,6 +1464,7 @@ def _create_gufunc_for_parfor_body(
         if config.DEBUG_CSA:
             print("After compile gufunc.")
             print("kernel_func", kernel_func, type(kernel_func))
+            sys.stdout.flush()
         kernel_func = compiler.compile_result(typing_context=typingctx,
                                        target_context=targetctx,
                                        entry_point=kernel_func.kernel,

numba/targets/codegen.py

@@ -17,6 +17,7 @@
 from numba.compiler_lock import require_global_compiler_lock
 
 import pdb
+import sys
 
 _x86arch = frozenset(['x86', 'i386', 'i486', 'i586', 'i686', 'i786',
                       'i886', 'i986'])
@@ -1075,6 +1076,7 @@ def get_asm_str(self, filename):
         if config.DEBUG_CSA:
             print("CSACodeLibrary::get_asm_str", filename)
             print(self._final_module)
+            sys.stdout.flush()
         return str(self._codegen._tm.emit_assembly_file(self._final_module, filename))
 
 #class AOTCSACodegen(BaseCSACodegen):

numba/targets/csa.py

@@ -1,6 +1,7 @@
 from __future__ import print_function, absolute_import
 
 import sys
+import os
 
 import llvmlite.llvmpy.core as lc
 import llvmlite as ll
@@ -16,6 +17,7 @@
 from . import fastmathpass
 from llvmlite.llvmpy.core import (Type, Builder, LINKAGE_INTERNAL, Constant, ICMP_EQ)
 from llvmlite import ir
+from llvmlite import binding
 from ctypes import *
 
 
@@ -50,6 +52,23 @@ def init(self):
         self.is32bit = (utils.MACHINE_BITS == 32)
         self._internal_codegen = codegen.JITCSACodegen("numba.csa.exec")
 
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+
+        #with open(dir_path + "/libcsac.bc", 'rb') as f:
+        #    bccode = f.read()
+        #    self._csamod = binding.parse_bitcode(bccode)
+
+        #with open(dir_path + "/libcsamath.bc", 'rb') as f:
+        #    bccode = f.read()
+        #    self._csamath = binding.parse_bitcode(bccode)
+
+        with open(dir_path + "/libcsamath_fixed.bc", 'rb') as f:
+            bccode = f.read()
+            self._csamath = binding.parse_bitcode(bccode)
+
+        self.sqrt32extern = "sp_sqrt_RN"
+        self.sqrt64extern = "dp_sqrt_RN"
+
         # Map external C functions.
         #externals.c_math_functions.install(self)
 
@@ -152,6 +171,9 @@ def post_lowering(self, mod, library):
             intrinsics.fix_divmod(mod)
 
         library.add_linking_library(rtsys.library)
+        print("CSA post-lowering adding csa library.", library, type(library))
+        library._final_module.link_in(self._csamod)
+        library._final_module.link_in(self._csamath)
 
     def create_cpython_wrapper(self, library, fndesc, env, call_helper,
                                release_gil=False):
@@ -216,11 +238,22 @@ def prepare_csa_kernel(self, codelib, fname, argtypes):
             print("fname", fname, type(fname))
             print("argtypes", argtypes, type(argtypes))
             print("csa_asm_name", csa_asm_name, type(csa_asm_name))
+            print("codelib", codelib, type(codelib))
+            sys.stdout.flush()
 
         codelib.get_asm_str(csa_asm_name)
+        if config.DEBUG_CSA:
+            print("After get_asm_str")
+            sys.stdout.flush()
         library = self.codegen().create_library('')
+        if config.DEBUG_CSA:
+            print("After create_library")
+            sys.stdout.flush()
         #library.add_linking_library(codelib)
         wrapper, wrapfnty, wrapper_library = self.generate_kernel_wrapper(library, fname, argtypes, csa_asm_name)
+        if config.DEBUG_CSA:
+            print("After generate_kernel_wrapper")
+            sys.stdout.flush()
         return library, wrapper, wrapfnty, wrapper_library
 
     def generate_kernel_wrapper(self, library, fname, argtypes, csa_asm_name):
@@ -229,7 +262,13 @@ def generate_kernel_wrapper(self, library, fname, argtypes, csa_asm_name):
         The function being wrapped have the name ``fname`` and argument types
         ``argtypes``.  The wrapper function is returned.
         """
+        if config.DEBUG_CSA:
+            print("Before finalize in generate_kernel_wrapper")
+            sys.stdout.flush()
         library.finalize()
+        if config.DEBUG_CSA:
+            print("After finalize in generate_kernel_wrapper")
+            sys.stdout.flush()
         cput = registry.dispatcher_registry['cpu'].targetdescr 
         context = cput.target_context
 
@@ -260,6 +299,7 @@ def generate_kernel_wrapper(self, library, fname, argtypes, csa_asm_name):
             print("prefixed", prefixed)
             print("wrapfn", wrapfn)
             print("builder", builder)
+            sys.stdout.flush()
 
         ll.binding.load_library_permanently("/home/taanders/numba/numba_csa2/numba/numba/targets/libgeneric.so")
 

numba/targets/mathimpl.py

@@ -168,6 +168,45 @@ def float_impl(context, builder, sig, args):
 
     return float_impl
 
+def unary_sqrt_extern(fn, int_restype=False):
+    """
+    Register implementations of Python function *fn* using the
+    external function named *f32extern* and *f64extern* (for float32
+    and float64 inputs, respectively).
+    If *int_restype* is true, then the function's return value should be
+    integral, otherwise floating-point.
+    """
+    f_restype = types.int64 if int_restype else None
+
+    def float_impl(context, builder, sig, args):
+        """
+        Implement *fn* for a types.Float input.
+        """
+        [val] = args
+        mod = builder.module
+        input_type = sig.args[0]
+        lty = context.get_value_type(input_type)
+        sqrt32extern = getattr(context, 'sqrt32extern', "sqrtf")
+        sqrt64extern = getattr(context, 'sqrt64extern', "sqrt")
+
+        func_name = {
+            types.float32: sqrt32extern,
+            types.float64: sqrt64extern,
+            }[input_type]
+        fnty = Type.function(lty, [lty])
+        fn = cgutils.insert_pure_function(builder.module, fnty, name=func_name)
+        res = builder.call(fn, (val,))
+        res = context.cast(builder, res, input_type, sig.return_type)
+        return impl_ret_untracked(context, builder, sig.return_type, res)
+
+    lower(fn, types.Float)(float_impl)
+
+    # Implement wrapper for integer inputs
+    unary_math_int_impl(fn, float_impl)
+
+    return float_impl
+
+
 
 unary_math_intr(math.fabs, lc.INTR_FABS)
 #unary_math_intr(math.sqrt, lc.INTR_SQRT)
@@ -206,7 +245,8 @@ def float_impl(context, builder, sig, args):
     ceil_impl = unary_math_extern(math.ceil, "ceilf", "ceil")
     floor_impl = unary_math_extern(math.floor, "floorf", "floor")
 gamma_impl = unary_math_extern(math.gamma, "numba_gammaf", "numba_gamma") # work-around
-sqrt_impl = unary_math_extern(math.sqrt, "sqrtf", "sqrt")
+sqrt_impl = unary_sqrt_extern(math.sqrt)
+#sqrt_impl = unary_math_extern(math.sqrt, "sqrtf", "sqrt")
 trunc_impl = unary_math_extern(math.trunc, "truncf", "trunc", True)
 lgamma_impl = unary_math_extern(math.lgamma, "lgammaf", "lgamma")