This CL adds support for named custom instructions in declarative builders.
To allow this, it introduces a templated `CustomInstruction` class.
This CL also splits ValueHandle which can capture only the **value** in single-valued instructions from InstructionHandle which can capture any instruction but provide no typing and sugaring to extract the potential Value*.
PiperOrigin-RevId: 237543222
This CL adds the same helper classes that exist in the AST form of EDSCs to support a basic indexing notation and emit the proper load and store operations and capture MemRefViews as function arguments.
This CL also adds a wrapper class LoopNestBuilder to allow generic rank-agnostic loops over indices.
PiperOrigin-RevId: 237113755
When building unstructured control-flow there is a need to construct mlir::Block* before being able to fill them. This invites goto-style programming.
This CL introduces an alternative eager API for BR and COND_BR in which blocks are created eagerly and captured on the fly.
This allows reducing the number of calls to `BlockBuilder` from 4 to 2 in the `builder_blocks_eager` test and from 3 to 2 in the `builder_cond_branch_eager` test.
PiperOrigin-RevId: 237046114
This CL adds support for BranchHandle and BranchBuilder that require a slightly different
abstraction since an mlir::Block is not an mlir::Value.
This CL also adds support for the BR and COND_BR instructions and the relevant tests.
PiperOrigin-RevId: 237034312
This CL reworks the design of EDSCs from first principles.
It introduces a ValueHandle which can hold either:
1. an eagerly typed, delayed Value*
2. a precomputed Value*
A ValueHandle can be manipulated with intrinsic operations a nested within a NestedBuilder. These NestedBuilder are a more idiomatic nested builder abstraction that should feel intuitive to program in C++.
Notably, this abstraction does not require an AST to stage the construction of MLIR snippets from C++. Instead, the abstraction makes use of orderings between definition and declaration of ValueHandles and provides a NestedBuilder and a LoopBuilder helper classes to handle those orderings.
All instruction creations are meant to use the templated ValueHandle::create<> which directly calls mlir::Builder.create<>.
For now the EDSC AST and the builders live side-by-side until the C API is ported.
PiperOrigin-RevId: 237030945
This CL changes dialect op source files (.h, .cpp, .td) to follow the following
convention:
<full-dialect-name>/<dialect-namespace>Ops.{h|cpp|td}
Builtin and standard dialects are specially treated, though. Both of them do
not have dialect namespace; the former is still named as BuiltinOps.* and the
latter is named as Ops.*.
Purely mechanical. NFC.
PiperOrigin-RevId: 236371358
EDSC provide APIs for constructing and modifying the IR. These APIs are
currently tested by a "test" module pass that reads the dummy IR (empty
functions), recognizes certain function names and injects the IR into those
functions based on their name. This situation is unsatisfactory because the
expected outcome of the test lives in a different file than the input to the
test, i.e. the API calls.
Create a new binary for tests that constructs the IR from scratch using EDSC
APIs and prints it. Put FileCheck comments next to the printing. This removes
the need to have a file with dummy inputs and assert on its contents in the
test driver. The test source includes a simplistic test harness that runs all
functions marked as TEST_FUNC but intentionally does not include any
value-testing functionality.
PiperOrigin-RevId: 235886629
This CL adds a primitive to perform stripmining of a loop by a given factor and
sinking it under multiple target loops.
In turn this is used to implement imperfectly nested loop tiling (with interchange) by repeatedly calling the stripmineSink primitive.
The API returns the point loops and allows repeated invocations of tiling to achieve declarative, multi-level, imperfectly-nested tiling.
Note that this CL is only concerned with the mechanical aspects and does not worry about analysis and legality.
The API is demonstrated in an example which creates an EDSC block, emits the corresponding MLIR and applies imperfectly-nested tiling:
```cpp
auto block = edsc::block({
For(ArrayRef<edsc::Expr>{i, j}, {zero, zero}, {M, N}, {one, one}, {
For(k1, zero, O, one, {
C({i, j, k1}) = A({i, j, k1}) + B({i, j, k1})
}),
For(k2, zero, O, one, {
C({i, j, k2}) = A({i, j, k2}) + B({i, j, k2})
}),
}),
});
// clang-format on
emitter.emitStmts(block.getBody());
auto l_i = emitter.getAffineForOp(i), l_j = emitter.getAffineForOp(j),
l_k1 = emitter.getAffineForOp(k1), l_k2 = emitter.getAffineForOp(k2);
auto indicesL1 = mlir::tile({l_i, l_j}, {512, 1024}, {l_k1, l_k2});
auto l_ii1 = indicesL1[0][0], l_jj1 = indicesL1[1][0];
mlir::tile({l_jj1, l_ii1}, {32, 16}, l_jj1);
```
The edsc::Expr for the induction variables (i, j, k_1, k_2) provide the programmatic hooks from which tiling can be applied declaratively.
PiperOrigin-RevId: 235548228
Leverage the recently introduced support for multiple argument groups and
multiple destination blocks in EDSC Expressions to implement conditional
branches in EDSC. Conditional branches have two successors and three argument
groups. The first group contains a single expression of i1 type that
corresponds to the condition of the branch. The two following groups contain
arguments of the two successors of the conditional branch instruction, in the
same order as the successors. Expose this instruction to the C API and Python
bindings.
PiperOrigin-RevId: 235542768
The new implementation of blocks was designed to support blocks with arguments.
More specifically, StmtBlock can be constructed with a list of Bindables that
will be bound to block aguments upon construction. Leverage this functionality
to implement branch instructions with arguments.
This additionally requires the statement storage to have a list of successors,
similarly to core IR operations.
Becauase successor chains can form loops, we need a possibility to decouple
block declaration, after which it becomes usable by branch instructions, from
block body definition. This is achieved by creating an empty block and by
resetting its body with a new list of instructions. Note that assigning a
block from another block will not affect any instructions that may have
designated this block as their successor (this behavior is necessary to make
value-type semantics of EDSC types consistent). Combined, one can now write
generators like
EDSCContext context;
Type indexType = ...;
Bindable i(indexType), ii(indexType), zero(indexType), one(indexType);
StmtBlock loopBlock({i}, {});
loopBlock.set({ii = i + one,
Branch(loopBlock, {ii})});
MLIREmitter(&builder)
.bindConstant<ConstantIndexOp>(zero, 0)
.bindConstant<ConstantIndexOp>(one, 1)
.emitStmt(Branch(loopBlock, {zero}));
where the emitter will emit the statement and its successors, if present.
PiperOrigin-RevId: 235541892
Introduce support for binding MLIR functions as constant expressions. Standard
constant operation supports functions as possible constant values.
Provide C APIs to look up existing named functions in an MLIR module and expose
them to the Python bindings. Provide Python bindings to declare a function in
an MLIR module without defining it and to add a definition given a function
declaration. These declarations are useful when attempting to link MLIR
modules with, e.g., the standard library.
Introduce EDSC support for direct and indirect calls to other MLIR functions.
Internally, an indirect call is always emitted to leverage existing support for
delayed construction of MLIR Values using EDSC Exprs. If the expression is
bound to a constant function (looked up or declared beforehand), MLIR constant
folding will be able to replace an indirect call by a direct call. Currently,
only zero- and one-result functions are supported since we don't have support
for multi-valued expressions in EDSC.
Expose function calling interface to Python bindings on expressions by defining
a `__call__` function accepting a variable number of arguments.
PiperOrigin-RevId: 234959444
The recent rework of MLIREmitter switched to using the generic call to
`builder.createOperation` from OperationState instead of individual customized
calls to `builder.create<>`. As a result, regular non-composed affine apply
operations where emitted. Introduce a special case in Expr::build to always
create composed affine maps instead, as it used to be the case before the
rework.
Such special-casing goes against the idea of EDSC generality and extensibility.
Instead, we should consider declaring the composed form canonical for
affine.apply operations and using the builder support for creating operations
and canonicalizing them immediately (ongoing effort).
PiperOrigin-RevId: 234790129
MLIR supports 'for' loops with lower(upper) bound defined by taking a
maximum(minimum) of a list of expressions, but does not have first-class affine
constructs for the maximum(minimum). All these expressions must have affine
provenance, similarly to a single-expression bound. Add support for
constructing such loops using EDSC. The expression factory function is called
`edsc::MaxMinFor` to (1) highlight that the maximum(minimum) operation is
applied to the lower(upper) bound expressions and (2) differentiate it from a
`edsc::For` that creates multiple perfectly nested loops (and should arguably
be called `edsc::ForNest`).
PiperOrigin-RevId: 234785996
Introduce a functionality to create EDSC expressions from typed constants.
This complements the current functionality that uses "unbound" expressions and
binds them to a specific constant before emission. It comes in handy in cases
where we want to check if something is a constant early during construciton
rather than late during emission, for example multiplications and divisions in
affine expressions. This is also consistent with MLIR vision of constants
being defined by an operation (rather than being special kinds of values in the
IR) by exposing this operation as EDSC expression.
PiperOrigin-RevId: 234758020
This CL fixes 2 recent issues with EDSCs:
1. the type of the LHS in Stmt::operator=(Expr rhs) should be the same as the (asserted unique) return type;
2. symbols coming from DimOp should be admissible as lower / upper bounds in For
The relevant tests are added.
PiperOrigin-RevId: 234750249
EDSC currently implement a block as a statement that is itself a list of
statements. This suffers from two modeling problems: (1) these blocks are not
addressable, i.e. one cannot create an instruction where thus constructed block
is a successor; (2) they support block nesting, which is not supported by MLIR
blocks. Furthermore, emitting such "compound statement" (misleadingly named
`Block` in Python bindings) does not actually produce a new Block in the IR.
Implement support for creating actual IR Blocks in EDSC. In particular, define
a new StmtBlock EDSC class that is neither an Expr nor a Stmt but contains a
list of Stmts. Additionally, StmtBlock may have (early-) typed arguments.
These arguments are Bindable expressions that can be used inside the block.
Provide two calls in the MLIREmitter, `emitBlock` that actually emits a new
block and `emitBlockBody` that only emits the instructions contained in the
block without creating a new block. In the latter case, the instructions must
not use block arguments.
Update Python bindings to make it clear when instruction emission happens
without creating a new block.
PiperOrigin-RevId: 234556474
The existing implementation in EDSC of 'for' loops in MLIREmitter is
unnecessarily restricted to constant bounds. The underlying AffineForOp can be
constructed from (a list of) Values and AffineMaps instead of constants. Its
verifier will check that the "affine provenance" conditions, i.e. that the
values used in the loop conditions are defined in such a way that they can be
analyzed by affine passes, are respected. One can use non-constant values in
affine loop bounds in conjunction with a single-dimensional identity affine
map. Implement this in MLIREmitter while maintaining the special case for
constant bounds that leads to significantly simpler generated IR when
applicable.
Test this change using the EDSC lowering test pass to inject code emitted from
EDSC into functions with predefined names.
PiperOrigin-RevId: 233578220
This CL applies the following simplifications to EDSCs:
1. Rename Block to StmtList because an MLIR Block is a different, not yet
supported, notion;
2. Rework Bindable to drop specific storage and just use it as a simple wrapper
around Expr. The only value of Bindable is to force a static cast when used by
the user to bind into the emitter. For all intended purposes, Bindable is just
a lightweight check that an Expr is Unbound. This simplifies usage and reduces
the API footprint. After playing with it for some time, it wasn't worth the API
cognition overhead;
3. Replace makeExprs and makeBindables by makeNewExprs and copyExprs which is
more explicit and less easy to misuse;
4. Add generally useful functionality to MLIREmitter:
a. expose zero and one for the ubiquitous common lower bounds and step;
b. add support to create already bound Exprs for all function arguments as
well as shapes and views for Exprs bound to memrefs.
5. Delete Stmt::operator= and replace by a `Stmt::set` method which is more
explicit.
6. Make Stmt::operator Expr() explicit.
7. Indexed.indices assertions are removed to pave the way for expressing slices
and views as well as to work with 0-D memrefs.
The CL plugs those simplifications with TableGen and allows emitting a full MLIR function for
pointwise add.
This "x.add" op is both type and rank-agnostic (by allowing ArrayRef of Expr
passed to For loops) and opens the door to spinning up a composable library of
existing and custom ops that should automate a lot of the tedious work in
TF/XLA -> MLIR.
Testing needs to be significantly improved but can be done in a separate CL.
PiperOrigin-RevId: 231982325
This CL also makes ScopedEDSCContexts to reset the Bindable numbering when
creating a new context.
This is useful to write minimal tests that don't use FileCheck pattern
captures for now.
PiperOrigin-RevId: 230079997
This is mostly plumbing to start allowing testing EDSC lowering. Prototype specifying reference implementation using verbose format without any generation/binding support. Add test pass that dumps the constructed EDSC (of which there can only be one). The idea is to enable iterating from multiple sides, this is wrong on many dimensions at the moment.
PiperOrigin-RevId: 229570535
Nicolas Vasilache