Commit Graph

179 Commits

Author SHA1 Message Date
Feng Liu 5f69643cbf Support TF Variant type in the tf/mlir roundtrip pass.
PiperOrigin-RevId: 213748573
2019-03-29 13:16:18 -07:00
Feng Liu 4bc5dc9602 Handle the TF resource data type in the TF/XLA roundtrip pass.
PiperOrigin-RevId: 213650346
2019-03-29 13:16:03 -07:00
Uday Bondhugula ab4797229c Extend loop unroll/unroll-and-jam to affine bounds + refactor related code.
- extend loop unroll-jam similar to loop unroll for affine bounds
- extend both loop unroll/unroll-jam to deal with cleanup loop for non multiple
  of unroll factor.
- extend promotion of single iteration loops to work with affine bounds
- fix typo bugs in loop unroll
- refactor common code b/w loop unroll and loop unroll-jam
- move prototypes of non-pass transforms to LoopUtils.h
- add additional builder methods.
- introduce loopUnrollUpTo(factor) to unroll by either factor or trip count,
  whichever is less.
- remove Statement::isInnermost (not used for now - will come back at the right
  place/in right form later)

PiperOrigin-RevId: 213471227
2019-03-29 13:15:06 -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
Tatiana Shpeisman 1a56ee7093 Implement operands for the 'if' statement.
This CL also includes two other minor changes:
- change the implemented syntax from 'if (cond)' to 'if cond', as specified by MLIR spec.
- a minor fix to the implementation of the ForStmt.

PiperOrigin-RevId: 210618122
2019-03-29 13:07:16 -07:00
Chris Lattner adf48e1bd2 Introduce a new Location abstraction to represent location data in a structured
(and more useful) way rather than hacking up a pile of attributes for it.  In
the future this will grow to represent inlined locations, fusion cases etc, but
for now we start with simple Unknown and File/Line/Col locations.  NFC.

PiperOrigin-RevId: 210485775
2019-03-29 13:06:49 -07:00
Tatiana Shpeisman d32a28c520 Implement operands for the lower and upper bounds of the for statement.
This revamps implementation of the loop bounds in the ForStmt, using general representation that supports operands. The frequent case of constant bounds is supported
via special access methods.

This also includes:
- Operand iterators for the Statement class.
- OpPointer::is() method to query the class of the Operation.
- Support for the bound shorthand notation parsing and printing.
- Validity checks for the bound operands used as dim ids and symbols

I didn't mean this CL to be so large. It just happened this way, as one thing led to another.

PiperOrigin-RevId: 210204858
2019-03-29 13:05:16 -07:00
Chris Lattner acd5bd98d1 First steps towards TF/XLA control flow lowering: functional if lowering.
- Implement support for the TensorFlow 'If' op, the first TF op definition.
 - Fill in some missing basic infra, including the ability to split a basic block, the ability to create a branch with operands, etc.
 - Implement basic lowering for some simple forms of If, where the condition is a zero-D bool tensor and when all the types line up.  Future patches will generalize this.

There is still much to be done here.  I'd like to get some example graphs coming from the converter to play with to direct this work.

PiperOrigin-RevId: 210198760
2019-03-29 13:05:01 -07:00
Chris Lattner 956e0f7e21 Push location information more tightly into the IR, providing space for every
operation and statement to have a location, and make it so a location is
required to be specified whenever you make one (though a null location is still
allowed).  This is to encourage compiler authors to propagate loc info
properly, allowing our failability story to work well.

This is still a WIP - it isn't clear if we want to continue abusing Attribute
for location information, or whether we should introduce a new class heirarchy
to do so.  This is good step along the way, and unblocks some of the tf/xla
work that builds upon it.

PiperOrigin-RevId: 210001406
2019-03-29 13:04:33 -07:00
Chris Lattner 84259c7def Implement call and call_indirect ops.
This also fixes an infinite recursion in VariadicOperands that this turned up.

PiperOrigin-RevId: 209692932
2019-03-29 13:03:51 -07:00
Uday Bondhugula 00bed4bd99 Extend loop unrolling to unroll by a given factor; add builder for affine
apply op.

- add builder for AffineApplyOp (first one for an operation that has
  non-zero operands)
- add support for loop unrolling by a given factor; uses the affine apply op
  builder.

While on this, change 'step' of ForStmt to be 'unsigned' instead of
AffineConstantExpr *. Add setters for ForStmt lb, ub, step.

Sample Input:

// CHECK-LABEL: mlfunc @loop_nest_unroll_cleanup() {
mlfunc @loop_nest_unroll_cleanup() {
  for %i = 1 to 100 {
    for %j = 0 to 17 {
      %x = "addi32"(%j, %j) : (affineint, affineint) -> i32
      %y = "addi32"(%x, %x) : (i32, i32) -> i32
    }
  }
  return
}

Output:

$ mlir-opt -loop-unroll -unroll-factor=4 /tmp/single2.mlir
#map0 = (d0) -> (d0 + 1)
#map1 = (d0) -> (d0 + 2)
#map2 = (d0) -> (d0 + 3)
mlfunc @loop_nest_unroll_cleanup() {
  for %i0 = 1 to 100 {
    for %i1 = 0 to 17 step 4 {
      %0 = "addi32"(%i1, %i1) : (affineint, affineint) -> i32
      %1 = "addi32"(%0, %0) : (i32, i32) -> i32
      %2 = affine_apply #map0(%i1)
      %3 = "addi32"(%2, %2) : (affineint, affineint) -> i32
      %4 = affine_apply #map1(%i1)
      %5 = "addi32"(%4, %4) : (affineint, affineint) -> i32
      %6 = affine_apply #map2(%i1)
      %7 = "addi32"(%6, %6) : (affineint, affineint) -> i32
    }
    for %i2 = 16 to 17 {
      %8 = "addi32"(%i2, %i2) : (affineint, affineint) -> i32
      %9 = "addi32"(%8, %8) : (i32, i32) -> i32
    }
  }
  return
}

PiperOrigin-RevId: 209676220
2019-03-29 13:03:38 -07:00
Chris Lattner 9265197c4e Implement initial support for function attributes, including parser, printer,
resolver support.

Still TODO are verifier support (to make sure you don't use an attribute for a
function in another module) and the TODO in ModuleParser::finalizeModule that I
will handle in the next patch.

PiperOrigin-RevId: 209361648
2019-03-29 13:02:44 -07:00
Jacques Pienaar 067d70f20d Add convenience builder for MemRefType.
PiperOrigin-RevId: 208245701
2019-03-29 12:58:32 -07:00
Uday Bondhugula 8a663870e8 Support for affine integer sets
- introduce affine integer sets into the IR
- parse and print affine integer sets (both inline or outlined) similar to
  affine maps
- use integer set for IfStmt's conditional, and implement parsing of IfStmt's
  conditional

- fixed an affine expr paren omission bug while one this.

TODO: parse/represent/print MLValue operands to affine integer set references.
PiperOrigin-RevId: 207779408
2019-03-29 12:56:58 -07:00
Chris Lattner 17ef97bf7e Refactor the asmparser hook to work with a new OperationState type that fully
encapsulates an operation that is yet to be created.  This is a patch towards
custom ops providing create methods that don't need to be templated, allowing
them to move out of line in the future.

PiperOrigin-RevId: 207725557
2019-03-29 12:56:30 -07:00
James Molloy 72645b31b8 [mlir] Add a TypeAttr class, allow type attributes
PiperOrigin-RevId: 207235956
2019-03-29 12:54:11 -07:00
James Molloy 1e793eb8dc [mlir] Add a string type
PiperOrigin-RevId: 206977161
2019-03-29 12:52:35 -07:00
Tatiana Shpeisman 8189a12bce Clean up and extend MLFuncBuilder to allow creating statements in the middle of a statement block. Rename Statement::getFunction() and StmtBlock()::getFunction() to findFunction() to make it clear that this is not a constant time getter.
Fix b/112039912 - we were recording 'i' instead of '%i' for loop induction variables causing "use of undefined SSA value" error.

PiperOrigin-RevId: 206884644
2019-03-29 12:51:38 -07:00
Tatiana Shpeisman c8b0273f19 Implement induction variables. Pretty print induction variable operands as %i<ssa value number>. Add support for future pretty printing of ML function arguments as %arg<ssa value number>.
Induction variables are implemented by inheriting ForStmt from MLValue. ForStmt provides APIs that make this design decision invisible to the ForStmt users.

This CL in combination with cl/206253643 resolves  http://b/111769060.

PiperOrigin-RevId: 206655937
2019-03-29 12:49:36 -07:00
Chris Lattner c77f39f55c Eliminate "primitive" types from being a thing, splitting them into FloatType
and OtherType.  Other type is now the thing that holds AffineInt, Control,
eventually Resource, Variant, String, etc.  FloatType holds the floating point
types, and allows convenient query of isa<FloatType>().

This fixes issues where we allowed control to be the element type of tensor,
memref, vector.  At the same time, ban AffineInt from being an element of a
vector/memref/tensor as well since we don't need it.

I updated the spec to match this as well.

PiperOrigin-RevId: 206361942
2019-03-29 12:46:57 -07:00
Jacques Pienaar 6a93e146c0 Add tf_control type and allow $ in bare-id.
* Add tf_control as primitive type;
* Allow $ in bare-id to allow attributes with $ (to make it trivially to mangle a TF attribute);

PiperOrigin-RevId: 206342642
2019-03-29 12:46:30 -07:00
Tatiana Shpeisman 1b24c48b91 Scaffolding for convertToCFG pass that replaces all instances of ML functions with equivalent CFG functions. Traverses module MLIR, generates CFG functions (empty for now) and removes ML functions. Adds Transforms library and tests.
PiperOrigin-RevId: 205848367
2019-03-29 12:41:15 -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
Tatiana Shpeisman 6ada91db02 Parse ML function arguments, return statement operands, and for statement loop header.
Loop bounds and presumed to be constants for now and are stored in ForStmt as affine constant expressions.  ML function arguments, return statement operands and loop variable name are dropped for now.

PiperOrigin-RevId: 205256208
2019-03-29 12:36:20 -07:00
MLIR Team f1e039617b Support for AffineMapAttr.
PiperOrigin-RevId: 205157390
2019-03-29 12:35:40 -07:00
Uday Bondhugula 8fbaf79afb Parse affine map range sizes.
PiperOrigin-RevId: 204240947
2019-03-29 12:32:59 -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
Chris Lattner a5a6c77e91 Introduce the start of IR builder APIs, which makes it easier and less error
prone to create things.

PiperOrigin-RevId: 203703229
2019-03-29 12:30:49 -07:00