[ONNX] onnx.Loop

[renxida]

Will need to extend importer for blocks and regions.

example test case here
https://gist.github.com/renxida/6e859dbfab286916dd8b99542c0a2332

Todo list for now:

• find the torch aten / prim ops and types for if op, loop op, and graph regions
• find todo items for extending the importer

Torch ops to translate to:

PrimLoopOp

PrimIfOp

    This op (together with prim.If.yield) define a conditional control flow
    construct. It is analogous to `scf.if` for MLIR folks that are familiar
    with that. The main differences from that op are:

    - `!torch.bool` condition value.
    - The "else" region is always present. This is reflective of invariants of
      the TorchScript IR.
    - No special prettiness for the "no yielded values" case. These are
      interesting for modeling mostly-non-SSA programs, but TorchScript IR
      is already in SSA form.

PrimIfYieldOp

Methods / types of interest:

mlir::Region

PrimIfOp::getThenRegion
PrimIfOp::getElseRegion
PrimIfOp::getSuccessorRegions

[renxida]

Rob's Guide for Implementing If and Loop Constructs with Regions in ONNX Lowering

TLDR:

• Regions are lists of related blocks. They don't provide methods or impose constraints, but other methods are built on top of the regions type to e.g. enforce single entry, single exit, or return types / return and yield ops.
• Our onnx import doesn't provide classes, so we directly interface with onnx ops that are all represented as the same op with an attribute telling us which op it is.

Some concepts

Block: A basic block with a single entry and exit point. Fundamental unit within a region.
Region:

• What Regions: A collection of connected basic blocks. Analogous to a function body, it can contain any number of blocks and represents a procedural logic unit within an operation.
• Why Regions: The introduction of regions allows for a more structured and efficient approach to control flow constructs like If and Loop. Regions encapsulate a collection of connected blocks, facilitating common optimizations such as common subexpression elimination (CSE) and dead code elimination (DCE). They don't provide methods or impose constraints, but other methods are built on top of the regions type to e.g. enforce single entry, single exit, or return types / return and yield ops.

How Linalg and SCF use Regions

• Linalg Operations: In the context of Linalg, regions function similarly to map-reduce components. A Linalg operation is a generic operation that can run both map-like and reduce-like operations using a region
• Structured Control Flow (SCF) Dialect: The introduction of regions is pivotal in optimizing code that uses the SCF dialect. Regions facilitate reasoning about conditions and enable Dead Code Elimination (DCE) and Common Subexpression Elimination (CSE).

[renxida]

Rob gave this reference for manipulating regions:

https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp

Inline if case

Rewrite terminators: the region's parent defines the rules about what terminators the region has. E.g. "must have 1 terminator block that ends in a particular op. e.g. functions must end with return"

e.g. in SCF, 2 blocks for If / else. Has to "scf yield" at the end. What's getting returned is given in the yield.

e.g. in while, one of the blocks has to have a conditional, and the other block has to have a yield block.
the conditional block and a yield / body block.

before / after block for for style loops

do while: before block and after block.

scf while: yield op is both the conditional and the before block

[renxida]

onnx binder: binder.op

binder.op->getRegions() for a regionrange

RegionRange.size()

[renxida]

Running this to test if import:

wget -O onnx_conditional_example.py https://gist.githubusercontent.com/renxida/6e859dbfab286916dd8b99542c0a2332/raw/e38d1558a2bdc9553a14338ca84c8f1ae4bb40b5/onnx_conditional_example.py
python onnx_conditional_example.py
PYTHONPATH=~/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir python -m torch_mlir.tools.import_onnx ./conditional_example.onnx -o conditional_example.mlir

[renxida]

@dosu-bot