search

IntelLabs / numba

NumPy aware dynamic Python compiler using LLVM
http://numba.pydata.org/
BSD 2-Clause "Simplified" License
12 stars 2 forks source link

easier parfor construction/translation #29

Open ninegua opened 7 years ago

ninegua commented 7 years ago

I'm thinking maybe we could build facilities similar to @lower_builtin to help translating various numpy functions to parfor. Here is an idea:

@parfor_call(np.transpose, types.Array)
def array_transpose(sig, args, arg_shapes)
    return array_T(sig.args[0], args[0], arg_shapes[0])

@parfor_getattr(types.Array, 'T')
def array_T(typ, inp, shape)
    out_shape = [shape[1], shape[0]]
    init = allocate_array(typ, out_shape)
    def body_func(out, i, j):
        out[i, j] = inp[j, i]
    loopnest = mk_loopnest(body_func, out_shape)
    return mk_parfor(init, loopnest)

This would translate both np.transpose(a) and a.T, by doing some pattern match on incoming IR (usually an assignment).

I think we can extend it further to incorporate array analysis too. So in the end we build two dictionaries that map function names (or getattr) to its corresponding (1) shape analysis, and (2) parfor translation. Eventually I'd see it merged into @lower_builtin so that all related handling of a single numpy call would sit in one place, but that is perhaps quite a stretch.

What do you think? @ehsantn @DrTodd13

ehsantn commented 7 years ago

The overall approach is very useful. In terms of details, I think we can write actual functions similar to @lower_builtin instead of things like mk_loopnest. prange can help with this.