Npcomp notes - Ufunc modeling Part 1 - April 29, 2020.md

I've been fiddling with the modeling of the numpy ufunc abstractions, because if you get that right, you have a large swath of standard operations.

Mapped to MLIR, they represent several ops: a definition (which defines a module level symbol) and various operations on the ufunc (such as ufunc_call, ufunc_reduce, ufunc_accumulate, ufunc_reduceat, ufunc_outer, ufunc_at). (The at variants perform in place updates)

The ufunc definition itself defines some metadata but is primarily an overloaded set of FunctionType signatures that operate on scalars. I've modeled this as an op that combines an array of FunctionType and variadic of regions:

def Numpy_GenericUfuncOp : Numpy_Op<"generic_ufunc", [
    IsolatedFromAbove,
    Symbol]> {
  let summary = "Defines a ufunc in terms of overloaded element-wise functions";
  let description = [{
  }];

  let arguments = (ins
    TypeArrayAttr:$overload_types);

  let regions = (region 
    VariadicRegion<AnyRegion>:$overloads);
}

Example (with some printer/parser fiddling):

  module @example_generic_ufunc {
    numpy.generic_ufunc @numpy.add(overload(%arg0: i32, %arg1: i32) -> i32 {
      %0 = addi %arg0, %arg1 : i32
      numpy.ufunc_return %0 : i32
    }, overload(%arg0: f32, %arg1: f32) -> f32 {
      %0 = addf %arg0, %arg1 : f32
      numpy.ufunc_return %0 : f32
    })
  }

We can auto-generate a module of such definitions covering the 60+ some primitives that come with numpy by generating overloads that map to the ufunc and the tuple of (MLIR scalar op, Type). A lot of them are simple, being implementable with a single std op. Others are more complex (some will require simple branching, etc). If we load that bootstrap module, then when tracing/compiling a python ast, when we encounter a ufunc that we know about, we just splat out the right call/reduce/accumulate/reduceat/outer/at op referencing its symbol. If tracing, there is a hook to get called back on this... it is just a few lines of code at the python level. Keeping it isomorphic with numpy in this way should give us everything we need to generically lower most of these to the next level down, do dtype, shape inference, etc. Note that in this formulation, things like "add" are not special ops like they are in XLA, TF, etc. They are more in line with what we expect from MLIR-born tooling like LinAlg, being mostly about the structure.