text: # MLIR: The case for a simplified polyhedral form
  
  MLIR embraces polyhedral compiler techniques for their many advantages
  representing and transforming dense numerical kernels, but it uses a form that
  differs significantly from other polyhedral frameworks.
  
  **Disclaimer / Warning**
  
  This document is a very early design proposal (which has since been accepted)
  that explored the tradeoffs of using this simplified form vs the traditional
  polyhedral schedule list form. At some point, this document could be dusted off
  and written as a proper academic paper, but until now, it is better to included
  it in this crafty form than not to. Beware that this document uses archaic
  syntax and should not be considered a canonical reference to modern MLIR.
  
  ## Introduction
  
  This document discusses general goals of the project, introduces context and the
  two alternatives, then talks about the tradeoffs of these designs. Written by
  Chris Lattner.
  
  ## General goals of an IR, and goals of mlfunc's specifically
  
  Our currently planned representation for MLIR consists of two kinds of
  functions: an LLVM-like "CFG Function" and an "ML Function": a function
  represented in multidimensional loop form. The idea is that a CFG function is
  capable of full generality for expressing arbitrary computation, but is awkward
  for loop transformations. In contrast, mlfunc's are limited (e.g. to control
  flow involving loop nests over affine spaces) but these limitations make it much
  easier to transform and analyze, particularly for the set of computations in a
  machine learning kernel.
  
  The design of an intermediate representations is an optimization problem, which
  makes intentional tradeoffs that aim to make certain kinds of compiler
  transformations simple. After all, it is "possible" to do almost any
  transformation on any IR: we could theoretically do loop transformations on
  assembly language. OTOH, such transformations would take too long to write,
  would be fragile due to irrelevant changes, would be difficult to maintain, and
  difficult to make target independent. Performing transformations on the "right
  level" of IR makes it much easier to do analysis and transformation of code, and
  can make them faster by reducing the size of the IR, and eliminating
  possibilities that would have otherwise have to be considered.
  
  This is the reason we're interested in adding polyhedral techniques to an IR in
  the first place: though our base "CFG function" representation is fully capable
  of expressing any computation, it is "too" expressive. The limitations imposed
  by polyhedral techniques (e.g. on affine loop bounds and array subscripts)
  define a closed algebra that can represent an interesting range of
  transformations and their compositions, and because of their simplicity, we can
  perform (e.g.) dependence analysis more efficiently and more reliably.
  
  This raises an important question that this document examines: given we are
  introducing a redundant and limited way to express code and transformations,
  exactly what form is best to perform the analyses and transformations we want?
  
  We explore two different design points that are capable of expressing the same
  class of affine loop computations, but which use different representational
  forms. These forms trade off verbosity, ease of transformation, and ease of
  analysis in interesting ways.
  
  ## Context: Traditional Polyhedral Form
  
  We started by discussing a representation that uses the traditional polyhedral
  schedule set + domain representation, e.g. consider C-like code like:
  
  ```c
    void simple_example(...) {
      for (int i = 0; i < N; ++i) {
        for (int j = 0; j < N; ++j) {
           float tmp = X[i,j]    // S1
           A[i,j] = tmp + 1      // S2
           B[i,j] = tmp * 42     // S3
         }
      }
    }
  ```
  
  The polyhedral representation doesn't care about the actual computation, so we
  will abstract them into S1/S2/S3 in the discussion below. Originally, we planned
  to represent this with a classical form like (syntax details are not important
  and probably slightly incorrect below):
  
  ```mlir
    mlfunc @simple_example(... %N) {
      %tmp = call @S1(%X, %i, %j)
        domain: (0 <= %i < %N), (0 <= %j < %N)
        schedule: (i, j, 0)
  
      call @S2(%tmp, %A, %i, %j)
        domain: (0 <= %i < %N), (0 <= %j < %N)
        schedule: (i, j, 1)
  
      call @S3(%tmp, %B, %i, %j)
        domain: (0 <= %i < %N), (0 <= %j < %N)
        schedule: (i, j, 2)
    }
  ```
  
  In this design, an mlfunc is an unordered bag of instructions whose execution
  order is fully controlled by their schedule.
  
  However, we recently agreed that a more explicit schedule tree representation is
  a better fit for our needs, because it exposes important structure that will
  make analyses and optimizations more efficient, and also makes the scoping of
  SSA values more explicit. This leads us to a representation along the lines of:
  
  ```mlir
    mlfunc @simple_example(... %N) {
      d0/d1 = mlspace
      for S1(d0), S2(d0), S3(d0) {
        for S1(d1), S2(d1), S3(d1) {
  
          %tmp = call @S1(%X, d0, d1)      ;; S1
            domain: (0 <= d0 < %N), (0 <= d1 < %N)
  
          call @S2(%tmp, %A, d0, d1)      ;; S2
            domain: (0 <= d0 < %N), (0 <= d1 < %N)
  
          call @S3(%tmp, %B, d0, d1)      ;; S3
            domain: (0 <= d0 < %N), (0 <= d1 < %N)
        }
      }
    }
  ```
  
  This change makes the nesting structure of the loops an explicit part of the
  representation, and makes lexical ordering within a loop significant
  (eliminating the constant 0/1/2 of schedules).
  
  It isn't obvious in the example above, but the representation allows for some
  interesting features, including the ability for instructions within a loop nest
  to have non-equal domains, like this - the second instruction ignores the outer
  10 points inside the loop:
  
  ```mlir
    mlfunc @reduced_domain_example(... %N) {
      d0/d1 = mlspace
      for S1(d0), S2(d0) {
        for S1(d1), S2(d1) {
          %tmp = call @S1(%X, d0, d1)    ;; S1
            domain: (0 <= d0 < %N), (0 <= d1 < %N)
  
          call @S2(%tmp, %A, d0, d1)      ;; S2
            domain: (10 <= d0 < %N-10), (10 <= d1 < %N-10)
        }
      }
    }
  ```
  
  It also allows schedule remapping within the instruction, like this example that
  introduces a diagonal skew through a simple change to the schedules of the two
  instructions:
  
  ```mlir
    mlfunc @skewed_domain_example(... %N) {
      d0/d1 = mlspace
      for S1(d0), S2(d0+d1) {
        for S1(d0+d1), S2(d1) {
          %tmp = call @S1(%X, d0, d1)    ;; S1
            domain: (0 <= d0 < %N), (0 <= d1 < %N)
  
          call @S2(%tmp, %A, d0, d1)      ;; S2
            domain: (0 <= d0 < %N), (0 <= d1 < %N)
        }
      }
    }
  ```
  
  This form has great power, and the polyhedral code generator (which lowers from
  an mlfunc to a cfgfunc representation) handles this power so things that
  introduce loop transformations don't have to explicitly manipulate the looping
  structure.
  
  ## Proposal: Simplified Polyhedral Form
  
  This document proposes and explores the idea of going one step further, moving
  all of the domain and schedule information into the "schedule tree". In this
  form, we would have a representation where all instructions inside of a given
  for-loop are known to have the same domain, which is maintained by the loop. In
  the simplified form, we also have an "if" instruction that takes an affine
  condition.
  
  Our simple example above would be represented as:
  
  ```mlir
    mlfunc @simple_example(... %N) {
      affine.for %i = 0 ... %N step 1 {
        affine.for %j = 0 ... %N step 1 {
          // identity noop in this case, but can exist in general.
          %0,%1 = affine.apply #57(%i, %j)
  
          %tmp = call @S1(%X, %0, %1)
  
          call @S2(%tmp, %A, %0, %1)
  
          call @S3(%tmp, %B, %0, %1)
        }
      }
    }
  ```
  
  The example with the reduced domain would be represented with an if instruction:
  
  ```mlir
    mlfunc @reduced_domain_example(... %N) {
      affine.for %i = 0 ... %N step 1 {
        affine.for %j = 0 ... %N step 1 {
          // identity noop in this case, but can exist in general.
          %0,%1 = affinecall #57(%i, %j)
  
          %tmp = call @S1(%X, %0, %1)
  
          if (10 <= %i < %N-10), (10 <= %j < %N-10) {
  
            %2,%3 = affine.apply(%i, %j)    // identity noop in this case
  
            call @S2(%tmp, %A, %2, %3)
          }
        }
      }
    }
  ```
  
  These IRs represent exactly the same information, and use a similar information
  density. The 'traditional' form introduces an extra level of abstraction
  (schedules and domains) that make it easy to transform instructions at the
  expense of making it difficult to reason about how those instructions will come
  out after code generation. With the simplified form, transformations have to do
  parts of code generation inline with their transformation: instead of simply
  changing a schedule to **(i+j, j)** to get skewing, you'd have to generate this
  code explicitly (potentially implemented by making polyhedral codegen a library
  that transformations call into):
  
  ```mlir
  mlfunc @skewed_domain_example(... %N) {
    affine.for %t1 = 0 ... 2*N-2 step 1 {
      affine.for %t2 = max(0, t1-N+1) ... min(N, t1) step 1 {
        (%i, %j) = (%t1-%t2, %t2)
        ...
      }
    }
  }
  ```
  
  ## Evaluation
  
  Both of these forms are capable of expressing the same class of computation:
  multidimensional loop nests with affine loop bounds and affine memory
  references. That said, they pose very different tradeoffs in other ways.
  
  ### Commonality: can express same computation
  
  Both of these can express the same sorts of computation, e.g. kernels written in
  one form are representable in the other form in all cases.
  
  ### Commonality: dependence analysis
  
  These representations both use affine functions for data layout mapping and
  access subscripts, and dependence analysis works the same way.
  
  ### Commonality: difficulty of determining optimal transformation series
  
  One major challenge in performance of optimization of this sort of code is
  choosing the ordering and behavior of various loop transformations that get
  applied. There are non-local effects of every decision, and neither
  representation helps solve this inherently hard problem.
  
  ### Commonality: compactness of IR
  
  In the cases that are most relevant to us (hyper rectangular spaces) these forms
  are directly equivalent: a traditional instruction with a limited domain (e.g.
  the "reduced_domain_example" above) ends up having one level of ML 'if' inside
  its loops. The simplified form pays for this by eliminating schedules and
  domains from the IR. Both forms allow code duplication to reduce dynamic
  branches in the IR: the traditional approach allows instruction splitting, the
  simplified form supports instruction duplication.
  
  It is important to point out that the traditional form wins on compactness in
  the extreme cases: e.g. the loop skewing case. These cases will be rare in
  practice for our workloads, and are exactly the cases that downstream
  transformations want to be explicit about what they are doing.
  
  ### Simplicity of code generation
  
  A key final stage of an mlfunc is its conversion to a CFG function, which is
  required as part of lowering to the target machine. The simplified form has a
  clear advantage here: the IR has a direct correspondence to the structure of the
  generated code.
  
  In contrast, the traditional form has significant complexity in the lowering
  process to a CFG function, because the verbosity not imbued in the IR needs to
  come out during code generation. Code generation from ISL shows that it is
  possible to do this, but it is a non-trivial transformation.
  
  ### Ease of transformation
  
  An advantage for the traditional form is that it is easier to perform certain
  transformations on it: skewing and tiling are just transformations on the
  schedule of the instructions in question, it doesn't require changing the loop
  structure.
  
  In practice, the simplified form requires moving the complexity of code
  generation into the transformations themselves - this is sometimes trivial,
  sometimes involved. The author believes that this should be possible by making
  the code generation algorithms themselves be library functions that
  transformations call into, instead of an opaque block that happens at the end of
  the mlfunc processing.
  
  Also, the sorts of transformations performed today by XLA (including tiling,
  padding, unrolling, and other rectangular transformations) should be easy enough
  to implement on either representation. The only cases that are a challenge are
  more advanced cases like skewing, e.g. for DMA data movement generation.
  
  ### Ease of analysis: Cost models
  
  The simplified form is much easier for analyses and transformations to build
  cost models for (e.g. answering the question of "how much code bloat will be
  caused by unrolling a loop at this level?"), because it is easier to predict
  what target code will be generated. With the traditional form, these analyses
  will have to anticipate what polyhedral codegen will do to a set of instructions
  under consideration: something that is non-trivial in the interesting cases in
  question (see "Cost of code generation").
  
  ### Cost of code generation
  
  State of the art polyhedral code generation is
  [expensive and complicated](https://lirias.kuleuven.be/bitstream/123456789/497238/1/toplas-astgen.pdf),
  sometimes exponential time complexity. We expect that most machine learning
  workloads will be hyper-rectangular, and thus it should be easy to specialize in
  important cases. That said, the traditional polyhedral representation makes it
  very easy to introduce complicated and expensive schedules, and provides no way
  to understand and project a cost model for using them. All downstream clients of
  the IR need to be prepared to handle the full generality of IR that may come to
  them.
  
  The simplified form defines this away: the concepts in the IR remain simple, and
  the code much more directly reflects the cost model for lowering to CFG
  functions and machine code. This is expected to be very important in the late
  stages of a code generator for an accelerator.
  
  ### SSA in ML Functions
  
  We agree already that values defined in an mlfunc can include scalar values and
  they are defined based on traditional dominance. In the simplified form, this is
  very simple: arguments and induction variables defined in for-loops are live
  inside their lexical body, and linear series of instructions have the same "top
  down" dominance relation that a basic block does.
  
  In the traditional form though, this is not the case: it seems that a lot of
  knowledge about how codegen will emit the code is necessary to determine if SSA
  form is correct or not. For example, this is invalid code:
  
  ```mlir
    %tmp = call @S1(%X, %0, %1)
      domain: (10 <= %i < %N), (0 <= %j < %N)
      schedule: (i, j)
  
    call @S2(%tmp, %A, %0, %1)
      domain: (0 <= %i < %N), (0 <= %j < %N)
      schedule: (i, j)
  ```
  
  Because `%tmp` isn't defined on some iterations of the %i loop.
  
  This matters because it makes the verifier more complicated, but more
  significantly, it means that load promotion and other optimizations that will
  produce SSA form will need to be aware of this and be able to model what codegen
  does.
  
  An emergent property of this that we discussed recently is that PHI nodes in
  mlfunc's (if we support them) will also have to have domains.
  
  ### Lack of redundancy in IR
  
  The traditional form has multiple encodings for the same sorts of behavior: you
  end up having bits on `affine.for` loops to specify whether codegen should use
  "atomic/separate" policies, unroll loops, etc. Instructions can be split or can
  generate multiple copies of their instruction because of overlapping domains,
  etc.
  
  This is a problem for analyses and cost models, because they each have to reason
  about these additional forms in the IR.
  
  ### Suitability to purpose: lowering to machine code
  
  One of the main drivers for this work is lowering to low-level accelerator code,
  including two-dimensional vectorization, insertion of DMAs, and other
  utilization of the matrix accelerator units. In the author's opinion, the extra
  compactness of the traditional form is a negative for this purpose: reasoning
  about the generated machine code will require understanding the mapping from
  mlfunc to lowered code, which means that it must understand what code generation
  will do.
  
  In the simplified form, the effect of "code generation" is always obvious from
  the IR itself, which should make it easier to perform vectorization to target
  instructions and other analyses we need to perform.
  
  ## Third Alternative: two different levels of mlfunc
  
  One hybrid alternative is to support both the traditional and simplified forms
  of mlfunc in our IR.
  
  The stages could look like this, for example:
  
  1.  Early performance transformations could be done on the traditional form.
  1.  Partial code generation lowers to the simplified form
  1.  Target specific lowering phases for tiling, and vectorization and other 2D
      transforms that don't benefit much from the traditional form could be run.
  1.  Final codegen to a cfg func can be done when all of the instructions are
      replaced with ones valid on the target.
  
  While this is possible, it isn't clear what would justify the complexity of this
  approach. Unless there is a super compelling reason for this, it would be nice
  to not do this. **Update:** we discussed this as a design team and agreed that
  this wouldn't be a good way to go.
  