Code Generation with DxT:
Improved Prototyping and
Development
Bryan Marker and Dillon Huff
The 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
interface
graph1
primitive
graph2
BLIS2014-7
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
interface
ok
better
BLIS2014-8
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
Let’s Look at Code
BLIS2014-12
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)
Martin Schatz. “Anatomy of Parallel Computation with Tensors.”
PhD Proposal. UT CS TR-3-21. 2013.
BLIS2014-15
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
Marker, van de Geijn, Batory. “Understanding Performance Stairs:
Elucidating Heuristics.” ASE 2014.
BLIS2014-19
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!
Square gemv
Square gemm
20
20
15
15
10
10
5
5
0
4
8
16
32
64
128
256
512
0
1024
4
8
16
32
Long n, Short m, varied p gemm
20
15
dxter GFLOPS
10
mkl GFLOPS
handwritten GFLOPS
5
0
128
256
384
512
640
786
896
1024
BLIS2014-30
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
Questions?
Comments?
bamarker@cs.utexas.edu
www.cs.utexas.edu/~bamarker
Thanks to Martin Schatz, Field Van Zee, Tyler Smith, and TACC!
BLIS2014-32