Commit Graph

26 Commits

Author SHA1 Message Date
Nicolas Vasilache c449e46ceb Introduce AffineExpr::compose(AffineMap)
This CL is the 1st on the path to simplifying AffineMap composition.
This CL uses the now accepted AffineExpr.replaceDimsAndSymbols to
implement `AffineExpr::compose(AffineMap)`.

Arguably, `simplifyAffineExpr` should be part of IR and not Analysis but
this CL does not yet pull the trigger on that.

PiperOrigin-RevId: 228265845
2019-03-29 15:03:36 -07:00
Uday Bondhugula 94c2d969ce Rename getAffineBinaryExpr -> getAffineBinaryOpExpr, getBinaryAffineOpExpr ->
getAffineBinaryOpExpr for consistency (NFC)

- this is consistent with the name of the class and getAffineDimExpr/ConstantExpr, etc.

PiperOrigin-RevId: 228164959
2019-03-29 14:59:52 -07:00
Nicolas Vasilache 28cf580555 Cleanup spurious DenseMap include
PiperOrigin-RevId: 228059305
2019-03-29 14:58:38 -07:00
Nicolas Vasilache 618c6a74c6 [MLIR] Introduce normalized single-result unbounded AffineApplyOp
Supervectorization does not plan on handling multi-result AffineMaps and
non-canonical chains of > 1 AffineApplyOp.
This CL introduces a simpler abstraction and composition of single-result
unbounded AffineApplyOp by using the existing unbound AffineMap composition.

This CL adds a simple API call and relevant tests:

```c++
OpPointer<AffineApplyOp> makeNormalizedAffineApply(
  FuncBuilder *b, Location loc, AffineMap map, ArrayRef<Value*> operands);
```

which creates a single-result unbounded AffineApplyOp.
The operands of AffineApplyOp are not themselves results of AffineApplyOp by
consrtuction.

This represent the simplest possible interface to complement the composition
of (mathematical) AffineMap, for the cases when we are interested in applying
it to Value*.

In this CL the composed AffineMap is not compressed (i.e. there exist operands
that are not part of the result). A followup commit will compress to normal
form.

The single-result unbounded AffineApplyOp abstraction will be used in a
followup CL to support the MaterializeVectors pass.

PiperOrigin-RevId: 227879021
2019-03-29 14:56:37 -07:00
Chris Lattner 7983bbc251 Introduce a simple canonicalization of affine_apply that drops unused dims and
symbols.

Included with this is some other infra:
 - Testcases for other canonicalizations that I will implement next.
 - Some helpers in AffineMap/Expr for doing simple walks without defining whole
   visitor classes.
 - A 'replaceDimsAndSymbols' facility that I'll be using to simplify maps and
   exprs, e.g. to fold one constant into a mapping and to drop/renumber unused dims.
 - Allow index (and everything else) to work in memref's, as we previously
   discussed, to make the testcase easier to write.
 - A "getAffineBinaryExpr" helper to produce a binop when you know the kind as
   an enum.

This line of work will eventually subsume the ComposeAffineApply pass, but it is no where close to that yet :-)

PiperOrigin-RevId: 227852951
2019-03-29 14:56:07 -07:00
Nicolas Vasilache 4adc169bd0 [MLIR] Add AffineMap composition and use it in Materialization
This CL adds the following free functions:
```
/// Returns the AffineExpr e o m.
AffineExpr compose(AffineExpr e, AffineMap m);
/// Returns the AffineExpr f o g.
AffineMap compose(AffineMap f, AffineMap g);
```

This addresses the issue that AffineMap composition is only available at a
distance via AffineValueMap and is thus unusable on Attributes.
This CL thus implements AffineMap composition in a more modular and composable
way.

This CL does not claim that it can be a good replacement for the
implementation in AffineValueMap, in particular it does not support bounded
maps atm.

Standalone tests are added that replicate some of the logic of the AffineMap
composition pass.

Lastly, affine map composition is used properly inside MaterializeVectors and
a standalone test is added that requires permutation_map composition with a
projection map.

PiperOrigin-RevId: 224376870
2019-03-29 14:20:22 -07:00
Nicolas Vasilache 3013dadb7c [MLIR] Basic infrastructure for vectorization test
This CL implements a very simple loop vectorization **test** and the basic
infrastructure to support it.

The test simply consists in:
1. matching the loops in the MLFunction and all the Load/Store operations
nested under the loop;
2. testing whether all the Load/Store are contiguous along the innermost
memory dimension along that particular loop. If any reference is
non-contiguous (i.e. the ForStmt SSAValue appears in the expression), then
the loop is not-vectorizable.

The simple test above can gradually be extended with more interesting
behaviors to account for the fact that a layout permutation may exist that
enables contiguity etc. All these will come in due time but it is worthwhile
noting that the test already supports detection of outer-vetorizable loops.

In implementing this test, I also added a recursive MLFunctionMatcher and some
sugar that can capture patterns
such as `auto gemmLike = Doall(Doall(Red(LoadStore())))` and allows iterating
on the matched IR structures. For now it just uses in order traversal but
post-order DFS will be useful in the future once IR rewrites start occuring.

One may note that the memory management design decision follows a different
pattern from MLIR. After evaluating different designs and how they quickly
increase cognitive overhead, I decided to opt for the simplest solution in my
view: a class-wide (threadsafe) RAII context.

This way, a pass that needs MLFunctionMatcher can just have its own locally
scoped BumpPtrAllocator and everything is cleaned up when the pass is destroyed.
If passes are expected to have a longer lifetime, then the contexts can easily
be scoped inside the runOnMLFunction call and storage lifetime reduced.
Lastly, whatever the scope of threading (module, function, pass), this is
expected to also be future-proof wrt concurrency (but this is a detail atm).

PiperOrigin-RevId: 217622889
2019-03-29 13:32:13 -07:00
Nicolas Vasilache 8ebb6ff171 [MLIR] Sketch AffineExpr value type
This CL sketches what it takes for AffineExpr to fully have by-value semantics
and not be a not-so-smart pointer anymore.

This essentially makes the underyling class a simple storage struct and
implements the operations on the value type directly. Since there is no
forwarding of operations anymore, we can full isolate the storage class and
make a hard visibility barrier by moving detail::AffineExpr into
AffineExprDetail.h.

AffineExprDetail.h is only included where storage-related information is
needed.

PiperOrigin-RevId: 216385459
2019-03-29 13:25:42 -07:00
Nicolas Vasilache 6707c7bea1 [MLIR] AffineExpr final cleanups
This CL:
1. performs the global codemod AffineXExpr->AffineXExprClass and
AffineXExprRef -> AffineXExpr;
2. simplifies function calls by removing the redundant MLIRContext parameter;
3. adds missing binary operator versions of scalar op AffineExpr where it
makes sense.

PiperOrigin-RevId: 216242674
2019-03-29 13:25:14 -07:00
Nicolas Vasilache ce2edea135 [MLIR] Cleanup AffineExpr
This CL introduces a series of cleanups for AffineExpr value types:
1. to make it clear that the value types should be used, the pointer
AffineExpr types are put in the detail namespace. Unfortunately, since the
value type operator-> only forwards to the underlying pointer type, we
still
need to expose this in the include file for now;
2. AffineExprKind is ok to use, it thus comes out of detail and thus of
AffineExpr
3. getAffineDimExpr, getAffineSymbolExpr, getAffineConstantExpr are
similarly
extracted as free functions and their naming is mande consistent across
Builder, MLContext and AffineExpr
4. AffineBinaryOpEx::simplify functions are made into static free
functions.
In particular it is moved away from AffineMap.cpp where it does not belong
5. operator AffineExprType is made explicit
6. uses the binary operators everywhere possible
7. drops the pointer usage everywhere outside of AffineExpr.cpp,
MLIRContext.cpp and AsmPrinter.cpp

PiperOrigin-RevId: 216207212
2019-03-29 13:24:45 -07:00
Nicolas Vasilache 4911978f7e [MLIR] Value types for AffineXXXExpr
This CL makes AffineExprRef into a value type.

Notably:
1. drops llvm isa, cast, dyn_cast on pointer type and uses member functions on
the value type. It may be possible to still use classof  (in a followup CL)
2. AffineBaseExprRef aggressively casts constness away: if we mean the type is
immutable then let's jump in with both feet;
3. Drop implicit casts to the underlying pointer type because that always
results in surprising behavior and is not needed in practice once enough
cleanup has been applied.

The remaining negative I see is that we still need to mix operator. and
operator->. There is an ugly solution that forwards the methods but that ends
up duplicating the class hierarchy which I tried to avoid as much as
possible. But maybe it's not that bad anymore since AffineExpr.h would still
contain a single class hierarchy (the duplication would be impl detail in.cpp)

PiperOrigin-RevId: 216188003
2019-03-29 13:24:31 -07:00
Nicolas Vasilache b55b407601 [RFC][MLIR] Use AffineExprRef in place of AffineExpr* in IR
This CL starts by replacing AffineExpr* with value-type AffineExprRef in a few
places in the IR. By a domino effect that is pretty telling of the
inconsistencies in the codebase, const is removed where it makes sense.

The rationale is that the decision was concisously made that unique'd types
have pointer semantics without const specifier. This is fine but we should be
consistent. In the end, the only logical invariant is that there should never
be such a thing as a const AffineExpr*, const AffineMap* or const IntegerSet*
in our codebase.

This CL takes a number of shortcuts to killing const with fire, in particular
forcing const AffineExprRef to return the underlying non-const
AffineExpr*. This will be removed once AffineExpr* has disappeared in
containers but for now such shortcuts allow a bit of sanity in this long quest
for cleanups.

The **only** places where const AffineExpr*, const AffineMap* or const
IntegerSet* may still appear is by transitive needs from containers,
comparison operators etc.

There is still one major thing remaining here: figure out why cast/dyn_cast
return me a const AffineXXX*, which in turn requires a bunch of ugly
const_casts. I suspect this is due to the classof
taking const AffineXXXExpr*. I wonder whether this is a side effect of 1., if
it is coming from llvm itself (I'd doubt it) or something else (clattner@?)

In light of this, the whole discussion about const makes total sense to me now
and I would systematically apply the rule that in the end, we should never
have any const XXX in our codebase for unique'd types (assuming we can remove
them all in containers and no additional constness constraint is added on us
from the outside world).

PiperOrigin-RevId: 215811554
2019-03-29 13:23:05 -07:00
Nicolas Vasilache 5b8017db18 [MLIR] Templated AffineExprBaseRef
This CL implements AffineExprBaseRef as a templated type to allow LLVM-style
casts to work properly. This also allows making AffineExprBaseRef::expr
private.

To achieve this, it is necessary to use llvm::simplify_type and make
AffineConstExpr derive from both AffineExpr and llvm::simplify<AffineExprRef>.
Note that llvm::simplify_type is just an interface to enable the proper
template resolution of isa/cast/dyn_cast but it otherwise holds no value.

Lastly note that certain dyn_cast operations wanted the const AffineExpr* form
of AffineExprBaseRef so I made the implicit constructor take that by default
and documented the immutable behavior. I think this is consistent with the
decision to make unique'd type immutable by convention and never use const on
them.

PiperOrigin-RevId: 215642247
2019-03-29 13:22:49 -07:00
Nicolas Vasilache 544f5e7a9b [MLIR] Remove uses of AffineExpr* outside of IR
This CL uniformizes the uses of AffineExprWrap outside of IR.
The public API of AffineExpr builder is modified to only use AffineExprWrap.
A few places access AffineExprWrap.expr, this is only while the API is in
transition to easily keep track (i.e. make expr private and let the compiler
track the errors).

Parser.cpp exhibits patterns that are dependent on nullptr values so
converting it is left for another CL.

PiperOrigin-RevId: 215642005
2019-03-29 13:22:35 -07:00
Nicolas Vasilache 9ef87c4b6b [MLIR] AffineExpr lightweight value type for operators
This CL proposes adding MLIRContext* to AffineExpr as discussed previously.
This allows the value class to not require the context in its constructor and
makes it a POD that it makes sense to pass by value everywhere.
A list of other RFC CLs will build on this. The RFC CLs are small incremental
pushes of the API which would be a pretty big change otherwise.

Pushing the thinking a little bit more it seems reasonable to use implicit
cast/constructor to/from AffineExpr*.
As this thing evolves, it looks to me like IR (and
probably Parser, for not so good reasons) want to operate on AffineExpr* and
the rest of the code wants to operate on the value type.

For this reason I think AffineExprImpl*/AffineExpr may also make sense but I
do not have a particular naming preference.
The jury is still out for naming decision between the above and
AffineExprBase*/AffineExpr or AffineExpr*/AffineExprRef.

PiperOrigin-RevId: 215641596
2019-03-29 13:22:21 -07:00
Nicolas Vasilache 4805e629c5 [MLIR] Use chainable ligthweight wrapper for AffineExpr
This CL argues that the builder API for AffineExpr should be used
with a lightweight wrapper that supports operators chaining.
This CL takes the ill-named AffineExprWrap and proposes a simple
set of operators with builtin constant simplifications.

This allows:
1. removing the getAddMulPureAffineExpr function;
2. avoiding concerns about constant vs non-constant simplifications
at **every call site**;
3. writing the mathematical expressions we want to write without unnecessary
obfuscations.

The points above represent pure technical debt that we don't want to carry on.
It is important to realize that this is not a mere convenience or "just sugar"
but reduction in cognitive overhead.

This thinking can be pushed significantly further, I have added some comments
with some basic ideas but we could make AffineMap, AffineApply and other
objects that use map applications more functional and value-based.

I am putting this out to get a first batch of reviews and see what people
think.

I think in my preferred design I would have the Builder directly return such
AffineExprPtr objects by value everywhere and avoid the boilerplate explicit
creations that I am doing by hand at this point.

Yes this AffineExprPtr would implicitly convert to AffineExpr* because that is
what it is.

PiperOrigin-RevId: 215641317
2019-03-29 13:22:07 -07:00
Uday Bondhugula a7611790f8 Add misc builder convenience methods for AffineMap's, for statement's.
Use these methods to simplify existing code. Rename getConstantMap
getConstantAffineMap. Move declarations to group similar ones together.

PiperOrigin-RevId: 212814829
2019-03-29 13:13:15 -07:00
Uday Bondhugula 64812a56c7 Extend getConstantTripCount to deal with a larger subset of loop bounds; make loop
unroll/unroll-and-jam more powerful; add additional affine expr builder methods

- use previously added analysis/simplification to infer multiple of unroll
  factor trip counts, making loop unroll/unroll-and-jam more general.

- for loop unroll, support bounds that are single result affine map's with the
  same set of operands. For unknown loop bounds, loop unroll will now work as
  long as trip count can be determined to be a multiple of unroll factor.

- extend getConstantTripCount to deal with single result affine map's with the
  same operands. move it to mlir/Analysis/LoopAnalysis.cpp

- add additional builder utility methods for affine expr arithmetic
  (difference, mod/floordiv/ceildiv w.r.t postitive constant). simplify code to
  use the utility methods.

- move affine analysis routines to AffineAnalysis.cpp/.h from
  AffineStructures.cpp/.h.

- Rename LoopUnrollJam to LoopUnrollAndJam to match class name.

- add an additional simplification for simplifyFloorDiv, simplifyCeilDiv

- Rename AffineMap::getNumOperands() getNumInputs: an affine map by itself does
  not have operands. Operands are passed to it through affine_apply, from loop
  bounds/if condition's, etc., operands are stored in the latter.

This should be sufficiently powerful for now as far as unroll/unroll-and-jam go for TPU
code generation, and can move to other analyses/transformations.

Loop nests like these are now unrolled without any cleanup loop being generated.

  for %i = 1 to 100 {
    // unroll factor 4: no cleanup loop will be generated.
    for %j = (d0) -> (d0) (%i) to (d0) -> (5*d0 + 3) (%i) {
      %x = "foo"(%j) : (affineint) -> i32
    }
  }

  for %i = 1 to 100 {
    // unroll factor 4: no cleanup loop will be generated.
    for %j = (d0) -> (d0) (%i) to (d0) -> (d0 - d mod 4 - 1) (%i) {
      %y = "foo"(%j) : (affineint) -> i32
    }
  }

  for %i = 1 to 100 {
    for %j = (d0) -> (d0) (%i) to (d0) -> (d0 + 128) (%i) {
      %x = "foo"() : () -> i32
    }
  }

TODO(bondhugula): extend this to LoopUnrollAndJam as well in the next CL (with minor
changes).

PiperOrigin-RevId: 212661212
2019-03-29 13:13:00 -07:00
Uday Bondhugula 0122a99cbb Affine expression analysis and simplification.
Outside of IR/
- simplify a MutableAffineMap by flattening the affine expressions
- add a simplify affine expression pass that uses this analysis
- update the FlatAffineConstraints API (to be used in the next CL)

In IR:
- add isMultipleOf and getKnownGCD for AffineExpr, and make the in-IR
  simplication of simplifyMod simpler and more powerful.
- rename the AffineExpr visitor methods to distinguish b/w visiting and
  walking, and to simplify API names based on context.

The next CL will use some of these for the loop unrolling/unroll-jam to make
the detection for the need of cleanup loop powerful/non-trivial.

A future CL will finally move this simplification to FlatAffineConstraints to
make it more powerful. For eg., currently, even if a mod expr appearing in a
part of the expression tree can't be simplified, the whole thing won't be
simplified.

PiperOrigin-RevId: 211012256
2019-03-29 13:07:44 -07:00
Uday Bondhugula 6d242fcf4b Simplify affine binary op expression class hierarchy
- Drop sub-classing of affine binary op expressions.
- Drop affine expr op kind sub. Represent it as multiply by -1 and add. This
  will also be in line with the math form when we'll need to represent a system of
  linear equalities/inequalities: the negative number goes into the coefficient
  of an affine form. (For eg. x_1 + (-1)*x_2 + 3*x_3 + (-2) >= 0). The folding
  simplification will transparently deal with multiplying the -1 with any other
  constants. This also means we won't need to simplify a multiply expression
  like in x_1 + (-2)*x_2 to a subtract expression (x_1 - 2*x_2) for
  canonicalization/uniquing.
- When we print the IR, we will still pretty print to a subtract when possible.

PiperOrigin-RevId: 205298958
2019-03-29 12:36:46 -07:00
Uday Bondhugula 8bbdd04365 Rename isSymbolic to isSymbolicOrConstant to avoid confusion.
PiperOrigin-RevId: 205288794
2019-03-29 12:36:33 -07:00
Uday Bondhugula b488a035aa Implement some simple affine expr canonicalization/simplification.
- fold constants when possible.
- for a mul expression, canonicalize to always keep the LHS as the
  constant/symbolic term, and similarly, the RHS for an add expression to keep
  it closer to the mathematical form. (Eg: f(x) = 3*x + 5)); other similar simplifications;
- verify binary op expressions at creation time.

TODO: we can completely drop AffineSubExpr, and instead use add and mul by -1.
This way something like x - 4 and -4 + x get canonicalized to x + -1 * 4
instead of being x - 4 and x + -4. (The other alternative if wanted to retain
AffineSubExpr would be to simplify x + -1*y to x - y and x + <neg number> to x
- <pos number>).
PiperOrigin-RevId: 204240258
2019-03-29 12:32:45 -07:00
Chris Lattner d6c4c5dbb8 Add attributes and affine expr/map to the Builder, switch the parser over to
use it.

This also removes "operand" from the affine expr classes: it is unnecessary
verbosity and "operand" will mean something very specific for SSA stuff (we
will have an Operand type).

PiperOrigin-RevId: 203976504
2019-03-29 12:32:08 -07:00
Uday Bondhugula fc46bcf51d Complete affine expr parsing support
- check for non-affine expressions
- handle negative numbers and negation of id's, expressions
- functions to check if a map is pure affine or semi-affine
- simplify/clean up affine map parsing code
- report more errors messages, more accurate error messages

PiperOrigin-RevId: 203773633
2019-03-29 12:31:03 -07:00
Uday Bondhugula fdf7bc4e25 [WIP] Sketching IR and parsing support for affine maps, affine expressions
Run test case:

$ mlir-opt test/IR/parser-affine-map.mlir
test/IR/parser-affine-map.mlir:3:30: error: expect '(' at start of map range
#hello_world2 (i, j) [s0] -> i+s0, j)
                             ^

PiperOrigin-RevId: 202736856
2019-03-29 12:27:20 -07:00
MLIR Team 8901448f14 Add some scaffolding for parsing affine maps:
- parsing affine map identifiers
- place-holder classes for AffineMap
- module contains a list of affine maps (defined at the top level).

PiperOrigin-RevId: 202336919
2019-03-29 12:26:28 -07:00