Extending standard library to include ML-relevant operations

[achowd32]

Hi, I'm a first time open source contributor looking to get involved and add something to XLS. After reading the XLS ideas page and going through issues such as this one, I decided to start by making an extension to the standard library to include more ML-relevant operations (which are becoming increasingly important in the hardware domain). This work has been overseen by @msaligane, who is a visiting faculty at Google.

I wanted to keep the pull requests fairly atomic, so I started with just two operations: ReLU (which is a pure fn) and a weight-stationary systolic array (which is a proc). I have more operations I am working on (e.g. approximated Softmax, output-stationary systolic array, etc) which I would be happy to contribute if reviewers are satisfied with this work, alongside appropriate documentation. I could pivot to a related idea in the standard library if the XLS team would prefer (e.g. a quantized floating point format).

I would like to note that the code for the systolic array here is different to the one provided in examples here. I found that my code got improved compilation times for the same architecture (details below). I am happy to gather PPA stats if it is deemed important.

	IR Generation Time	IR Optimization Time	Code Generation Time	Verilog LoC
My code	3.659 s	137.033 s	30.124 s	207498
Example code	6.945 s	150.122 s	40.654 s	285690

The above times were averaged on 5 runs on a 16x16 systolic array.

I would sincerely appreciate all feedback, as I am trying to learn and would love to make this work.

P. S. I seem to have found a bug in the DSLX formatter wherein comments cannot be placed between multi-line arrays. There doesn't seem to be an open issue for this, so I may file one later.

[allight]

@proppy @richmckeever Maybe one of you can take a look at this.

Given the many latency/throughput/implementation tradeoffs inherent in any matrix multiplier design I'm not sure we want to promote this to stdlib status but this seems like an excellent addition to the xls examples directory as the sort of flexibility DSLX allows a designer. Or maybe adding an xls/contrib/ directory for an ML dslx library might be worthwhile.

[achowd32]

@allight I would be happy to add and contribute to an xls/contrib/ directory if you deem it ok. I could perhaps also focus more on operations with fewer inherent tradeoffs, like activation functions (e.g. approximated sigmoid) or approximated softmax.

If you'd rather not do the above, is there something similar I could contribute to stdlib/ or contrib/? Potentially a quantized FP library?

(We can add the systolic array implementation in the PR to examples/ if you don't think it fits elsewhere).

[dplassgit]

We've discussed this some more internally and feel it would be best to move this code to a new directory in xls/modules. We also acknowledge that we don't have published criteria regarding what belongs in stdlib, so we've opened an internal bug to track that. Thanks for the contribution and your patience!