Code Generation with DxT:Improved Prototyping and

Development

Bryan Marker and Dillon HuffThe University of Texas at Austin

BLIS2014-1

Recap

• Last year I presented how code generation helps– to improve or (performance) validate existing code– to explore and develop implementation options for new

code

• Generated distributed memory dense linear algebra (DLA) code using the Elemental API

• Generated sequential and parallel BLAS3 code for large matrices using BLIS packing and macrokernel calls as an API– Quickly explored parallelization options via code

generation– Performance rivaled or beat MKL

BLIS2014-2

Code Generation

• Now, we are generating code for distributed-memory tensor contractions– Using Martin Schatz’s notation and API for output code

• Dillon will discuss complimentary work to generate code at a much lower level– Calling vector intrinsics

BLIS2014-3

Code Generation

• In all cases, a system is searching many software design options as a person would– Parallelization options– Blocking options– Loop transformations

• Sometimes the combination of options is not massive, but people still make mistakes – Correctness bugs– Performance bug

• Sometimes the combination is massive and a system must be used

BLIS2014-4

DESIGN BY TRANSFORMATION

BLIS2014-5

Graphs

• Data-flow, directed acyclic graphs

• A box or node represents an operation– An interface without implementation details– OR a primitive operation that maps to given code

• A starting algorithm specification is represented as a graph without implementation details

• We want to transform it into a graph with complete implementation details

BLIS2014-6

Transform with Implementations

• Refinements replace a box without implementation details– Chooses a specific way to implement the box’s functionality– E.g., choose how to parallelize a contraction or how to load

data into CPU registers

BLIS2014-7

interface

graph1

graph2

primitive

Transform to Optimize

• Optimizations replace a subgraph with another subgraph– Same functionality– A different way of implementing it– E.g., fuse loops or remove redundant communication

BLIS2014-8

interface

ok

better

Examples from Past Work

• Refinements choose – Parallelization of BLAS operations like Gemm and Trsm– Dimensions over which to partition data (maybe from

FLAME-derived algorithms)

• Optimizations – Explore alternate collective communication patterns to

redistribute data– Remove redundant communication or computation

• Just rewrite rules to explore design options

BLIS2014-9

Tensors!

BLIS2014-10

Tensors!

• Martin Schatz will talk about his notation for parallelizing contractions– Generalization of Jack Poulson’s Elemental notation

• Tensors have an arbitrary number of modes (aka dimensions)– Not known until algorithm is specified

• Therefore, Martin cannot implement (an efficient) distributed contraction as Jack has done for Gemm

• Infinite options for contractions a user might need– Infinite options for data distributions within the

contractions

BLIS2014-11

BLIS2014-12

Let’s Look at Code

AND IT WAS A MOVING TARGET (KIND OF)

BLIS2014-13

Code Generation!

• We encode software design knowledge– The knowledge one would use to implement a given

contraction (once it is known)

• Apply design knowledge automatically for a user-specified contraction – Explore a search space of options– Use cost estimates to rank order them – Output the fastest

• Notice that even implementing a given contraction is HARD– Could have tens of billions of ways to distribute data and

parallelize for a relatively simple contraction– Why should a person explore the options manually?

BLIS2014-14

Knowledge Encoding

• What are the valid ways to parallelize a contraction and how must the data be distributed to enable that?

• Three options– Keep tensor A, B, or C stationary – or don’t communicate it

from the default distribution– Encoded as refinements in DxT

• Martin’s proposal shows how to formalize these ideas– Martin will talk about that later– We encode his ideas

• Two of the inputs must be redistributed• The output might need to be redistributed (and summed across

processes)

BLIS2014-15

Martin Schatz. “Anatomy of Parallel Computation with Tensors.”PhD Proposal. UT CS TR-3-21. 2013.

Knowledge Encoding

• How can relatively simple communication patterns from Martin’s API be combined to form arbitrarily complex communication?

• Martin’s API provides collective communication for an arbitrary number of modes– AllToAll, AllGather, ReduceScatter, etc.– Not combinations of multiple collectives

• How can we combine these collectives?

BLIS2014-16

Knowledge Encoding

• How can we use complicated communication to redistribute data as needed?

• Allows for communication in an arbitrary number of modes and with arbitrary complexity– Implementations are generated when the required

communication pattern is known

• Refinements encode how to breakdown arbitrary communication into pieces of implemented communication patterns

BLIS2014-17

Knowledge Encoding

• Parallelizing a given contraction is not sufficient– The same tensor might be redistributed for multiple

contractions

• How can combinations of communication be optimized?

• How can we reduce memory operation to copy / pack data?

• Optimization transformations explore these options

BLIS2014-18

Implementation Search Space

• This leads to many options– Combinatorial search space– Billions of implementations

• We are studying how to limit the search space by cutting out clearly bad options– A developer does not consider every single design

option– He uses intuition / heuristics to limit consideration

• Search optimization

BLIS2014-19

Marker, van de Geijn, Batory. “Understanding Performance Stairs: Elucidating Heuristics.” ASE 2014.

Lessons Learned

• By formally defining and encoding design options we learn a lot

• Different experience to implement code directly than to encode the knowledge to implement all similar code– You have to justify your design choices

BLIS2014-20

Lessons Learned

• We find misconceptions in what we need from the API– New communication patterns– New computation variants

BLIS2014-21

Lessons Learned

• We identify new optimizations– We are spending less time encoding knowledge than we

would to implement the code directly– We are free to identify (and effect) optimizations

• Optimizations can be turned on and off easily, so we can see the effects

• Can add optimizations to the API and to DxTer hand-in-hand– Quickly re-generate all code with new optimizations– Trust generated code for correctness

BLIS2014-22

AND NOW FOR SOMETHING COMPLETELY DIFFERENT

BLIS2014-23

Small DLA

• Think of the software stack for a DLA library– Layers of loops reduce a large problem to a problem that

is small in some (or all) dimensions– Small, fixed problem sizes are not usually an optimized

case in libraries

• Such operations also show up as part of – Optimization algorithms– Real-time control– Tensor contractions– Small dense blocks in sparse problems

BLIS2014-24

Low-Level DLA

• Intel’s compilers are really good, but they’re not as good as expert developers– That’s why you pay for MKL

• DLA experts are very good– But not available to optimize every problem size in every

application for every hardware architecture

• Once again, why not encode the experts’ design knowledge to enable a system to generate code?

BLIS2014-25

Existing Work

• The BTO BLAS does this for level-1 and level-2 – Genetic representation and search

• SPIRAL and LGen do this with a mathematical operator notation and rewrite rules

• LGen generates code calling functions that are specialized to a particular machine using vector intrinsics– Problem sizes small enough to fit in the L1 cache

BLIS2014-26

DxT’s Flavor

• Bryan wanted to see if the DxT approach would work at this low level

• I was interested in DLA code generation

• We aimed to generate small BLAS like operations using a DxT representation (data fit in L1)– Where does DxT fail?– How much encoded design knowledge can be reused

from previous work?

BLIS2014-27

Knowledge Encoding

• DxTer needs a way to do loop tiling and unrolling

• Loop fusion is already included in DxTer

• Encode how to implement (really) small operations in terms of loading into registers, computing with registers, and saving back to main memory

BLIS2014-28

DxTer’s 3 stage approach

• Generate a variety of loop structures with different tiling, blocking, and loop fusion patterns

• Decide which loops to unroll and the factors to unroll them by

• Convert the updates specified in each implementation into primitives that map to C vector intrinsics

BLIS2014-29

Charts!

BLIS2014-30

4 8 16 32 640

5

10

15

20

Square gemm

4 8 16 32 64 128 256 512 10240

5

10

15

20

Square gemv

128 256 384 512 640 786 896 10240

5

10

15

20

Long n, Short m, varied p gemm

dxter GFLOPSmkl GFLOPShandwritten GFLOPS

Lessons Learned

• DxTer’s partial unrolling is important– Full unrolling is not good

• Functions implemented independently via vector operations do not perform as well as generating code using vector operations directly– And the overhead in the representation is small– So we generate code that calls vector intrinsics

• Empirical + cost modeling is effective and possible in DxTer– Weed out really bad with cost models and test the rest

• Good compilers are important (when the code is amenable)– icc is much better than gcc and clang

BLIS2014-31

So Far

BLIS2014-32

Questions?Comments?

bamarker@cs.utexas.edu

www.cs.utexas.edu/~bamarker

Thanks to Martin Schatz, Field Van Zee, Tyler Smith, and TACC!