Redesigning the Seastar Tracer (Frontend)

[JoelMathewC]

The need to redesign the seastar frontend are as follows

1. Monkey-patching is considered a very hacky solution and it cannot be pushed to production
2. The current tracer design makes the definition of helper functions(leaky_relu needs to be defined in self and then passed to nb_compute) is unintuitive.
3. Support for normal backend operations, as shown below. (Dont know how much changing the frontend will help with this)

@self.cm.zoomIn(nspace=[self, th])
        def nb_forward(v):
           coeff = [th.exp((self.leaky_relu(nb.el + v.er) * self.attn_l).sum(dim=-1).unsqueeze(-1)) for nb in v.innbs]
           s = sum(coeff)
           alpha = [c/s for c in coeff]
           feat_src = [nb.feat_src for nb in v.innbs]
           return sum([alpha[i] * feat_src[i] for i in range(len(feat_src))])
        rst = nb_forward(g=graph, n_feats= {'el':el, 'er': er, 'feat_src':feat_src})

fails with error

AttributeError: 'TorchVal' object has no attribute 'unsqueeze'

[JoelMathewC]

A bit early to start looking into this, a better time would be after finalizing on the requirements or we hit a roadblock and this becomes absolutely necessary. But out of curiosity, this issue would be fun to look into time to time. Something interesting I found that could play a part in the redesign is https://devguide.python.org/internals/compiler/

Additionally seems like the ast.parse() method as explained here takes a python program string as input and outputs the AST. But it does seem like there are limitations to this approach. It does seem like Redesigning the frontend is going to be quite an involved process.

To get the source code of a python function we can follow one of the many suggestions here. It maybe noted that these solutions also do seem a bit hacky.

With what we have so far it seems like we know how to convert a function into an AST within the current seastar setup. After getting the AST, a simple middleware should be sufficient to generate the grid. Still uncertain how functions like sum, leaky_relu etc. would be handled though.

[JoelMathewC]

Found some details on possible parsers for python here. However building a custom parser would be arduous and difficult to maintain. Refering to the parsing python in python sub-topic of the above reference maybe using pythons in-built ast and parse mechanisms are the best approach.

[JoelMathewC]

Realized that seastar is performing just in-time compilation (JIT). This is when compilation happens during runtime, instead of ahead of time. Was looking into numba, which is a JIT compiler for python, primarily as a replacement of templating used to make kernel code. Realized that maybe looking into numba might help understand what good practices are in building a frontend. Additionally, noticed that they used LLVM compiler library, might be something to look into at some point.

[JoelMathewC]

Found the numba compiler portion: https://numba.pydata.org/numba-doc/latest/developer/architecture.html It seems like they are converting to bytecode which is quite complicated and best to avoid at the moment.

[JoelMathewC]

Since type inference is one aspect of monkey-patching being used, maybe we can use something like https://mypy.readthedocs.io/en/stable/index.html to help with inference.