[MLIR] New Frontend with HeteroCL AST

[zzzDavid]

Summary

This PR revamps the Python front to resolve the issues with generating MLIR IR during tracing. Generating IR on the fly can insert MLIR operations at the wrong scope for imperative programs. This is caused by the Python expression evaluation order and MLIR being SSA.

We propose a solution that builds an "intermediate" layer of abstraction in Python and takes a second pass to build IR in MLIR. This layer of abstraction is HeteroCL AST. There are more benefits of introducing HeteroCL AST:

1. Simplifies front-end implementation. HeteroCL AST has a simple and regular structure that is easy to build and transform.
2. Removes context variables and global data structures such as If/For stack, making debugging easier.
3. Modular infrastructure. Steps such as building stages, dataflow graph, and building IR from AST can be separated into passes on the AST.
4. Local reasoning. MLIR operation insertion point can be resolved locally.

Major Changes

• Added HCL AST, a program representation layer in Python above MLIR IR.
• Added TypeInfer class, a type inference engine that outputs result data type given an operation and its operands.
• Added type rules. Type rules are defined as lambda functions to be extensible and decoupled from the type inference engine.
• Added IRBuilder, a builder that takes in HCL AST and generates MLIR.
• Added passes to check and transform HCL AST.

HCL AST Model

• HCL AST is an intermediate representation consisting of AST nodes. An AST node is either an operation or an expression. An operation cannot be assigned, an expression can be assigned.
  Operation: FuncOp, ComputeOp, IfOp, ElifOp, ElseOp, ForOp, StoreOp, ReuseAtOp …
  Expr: LoadOp, UnaryOp, BinaryOp, TernaryOp, …
• An AST operation can have a body region which is a list of operations, therefore, operations can be nested. In the figure above, each box is an operation.
• The AST is non-SSA, each AST node has a one-to-one mapping with HeteroCL API, making AST building straightforward.

New Frontend Workflow

1. All HeteroCL APIs, including customization APIs, generate operations in AST.
2. A series of passes are applied to transform the AST from "close-to-frontend-API" to "close-to-IR", including:
   • NestElseIf: “elif” to nested “else”
   • Promote FuncOps: promotes local functions to global scope
   • Separate host-xcel: separates host and xcel program
3. HCL AST is passed to IR Builder, which generates an MLIR module containing the IR represented in MLIR operations.

API Changes

1. Added API to print AST: schedule.ast
2. Deprecated device-agnostic MLIR module schedule.device_module, changing to schedule.module
3. Creating schedule without function only requires input tensors, the output tensors are not required. This follows the main branch syntax.

a = hcl.placeholder((10, 20))
b = hcl.placeholder((10, 20))
c = hcl.compute(a.shape, lambda i, j: a[i, j] + b[i, j])
s = hcl.create_schedule([a, b])

4. Since customization operations are built in AST as well, the MLIR IR is generated during hcl.lower instead of during hcl.create_schedule.

Tests

The new front-end implementation passes the same tests as the hcl-mlir branch.
Additionally, it passes tutorial tests. Tutorial tests are also enabled by this PR.