Skip to content

Commit

Permalink
Tile-level partitioning in jr/ir loops (ex-trsm). (#695)
Browse files Browse the repository at this point in the history
Details:
- Reimplemented parallelization of the JR loop in gemmt (which is
  recycled for herk, her2k, syrk, and syr2k). Previously, the
  rectangular region of the current MC x NC panel of C would be
  parallelized separately from from the diagonal region of that same
  submatrix, with the rectangular portion being assigned to threads via
  slab or round-robin (rr) partitioning (as determined at configure-
  time) and the diagonal region being assigned via round-robin. This
  approach did not work well when extracting lots of parallelism from
  the JR loop and was often suboptimal even for smaller degrees of
  parallelism. This commit implements tile-level load balancing (tlb) in
  which the IR loop is effectively subjugated in service of more
  equitably dividing work in the JR loop. This approach is especially
  potent for certain situations where the diagonal region of the MC x NR
  panel of C are significant relative to the entire region. However, it
  also seems to benefit many problem sizes of other level-3 operations
  (excluding trsm, which has an inherent algorithmic dependency in the
  IR loop that prevents the application of tlb). For now, tlb is
  implemented as _var2b.c macrokernels for gemm (which forms the basis
  for gemm, hemm, and symm), gemmt (which forms the basis of herk,
  her2k, syrk, and syr2k), and trmm (which forms the basis of trmm and
  trmm3). Which function pointers (_var2() or _var2b()) are embedded in
  the control tree will depend on whether the BLIS_ENABLE_JRIR_TLB cpp
  macro is defined, which is controlled by the value passed to the
  existing --thread-part-jrir=METHOD (or -r METHOD) configure option.
  This script adds 'tlb' as a valid option alongside the previously
  supported values of 'slab' and 'rr'. ('slab' is still the default.)
  Thanks to Leick Robinson for abstractly inspiring this work, and to
  Minh Quan Ho for inquiring (in PR #562, and before that in Issue #437)
  about the possibility of improved load balance in macrokernel loops,
  and even prototyping what it might look like, long before I fully
  understood the problem.
- In bli_thread_range_weighted_sub(), tweaked the the way we compute the
  area of the current MC x NC trapezoidal panel of C by better taking
  into account the microtile structure along the diagonal. Previously,
  it was an underestimate, as it assumed MR = NR = 1 (that is, it
  assumed that the microtile column of C that overlapped with microtiles
  exactly coincided with the diagonal). Now, we only assume MR = NR.
  This is still a slight underestimate when MR != NR, so the additional
  area is scaled by 1.5 in a hackish attempt to compensate for this, as
  well as other additional effects that are difficult to model (such as
  the increased cost of writing to temporary tiles before finally
  updating C). The net effect of this better estimation of the
  trapezoidal area should be (on average) slightly larger regions
  assigned to threads that have little or no overlap with the diagonal
  region (and correspondingly slightly smaller regions in the diagonal
  region), which we expect will lead to slightly better load balancing
  in most situations.
- Spun off the contents of bli_thread.[ch] that relate to computing
  thread ranges into one of three source/header file pairs:
  - bli_thread_range.[ch], which define functions that are not specific
    to the jr/ir loops;
  - bli_thread_range_slab_rr.[ch], which define functions that implement
    slab or round-robin partitioning for the jr/ir loops;
  - bli_thread_range_tlb.[ch], which define functions that implement
    tlb for the jr/ir loops.
- Fixed the computation of a_next in the last iteration of the IR loop
  in bli_gemmt_l_ker_var2(). Previously, it always "wrapped" back around
  to the first micropanel of the current MC x KC packed block of A.
  However, this is almost never actually the micropanel that is used
  next. A new macro, bli_gemmt_l_wrap_a_upanel(), computes a_next
  correctly, with a similarly named bli_gemmt_u_wrap_a_upanel() for use
  in the upper-stored case (which *does* actually always choose the
  first micropanel of A as its a_next at the end of the IR loop).
- Removed adjustments for a_next/b_next (a2/b2) for the diagonal-
  intersecting case of gemmt_l_ker_var2() and the above-diagonal case
  of gemmt_u_ker_var2() since these cases will only coincide with the
  last iteration of the IR loop in very small problems.
- Defined bli_is_last_iter_l() and bli_is_last_iter_u(), the latter of
  which explicitly considers whether the current microtile is the last
  tile that intersects the diagonal. (The former does the same, but the
  computation coincides with the original bli_is_last_iter().) These
  functions are now used in gemmt to test when a_next (or a2) should
  "wrap" (as discussed above). Also defined bli_is_last_iter_tlb_l()
  and bli_is_last_iter_tlb_u(), which are similar to the aforementioned
  functions but are used when employing tlb in gemmt.
- Redefined macros in bli_packm_thrinfo.h, which test whether an
  iteration of work is assigned to a thread, as static inline functions
  in bli_param_macro_defs.h (and then deleted bli_packm_thrinfo.h).
  In the process of redefining these macros, I also renamed them from
  bli_packm_my_iter_rr/sl() to bli_is_my_iter_rr/sl().
- Renamed
    bli_thread_range_jrir_rr() -> bli_thread_range_rr()
    bli_thread_range_jrir_sl() -> bli_thread_range_sl()
    bli_thread_range_jrir()    -> bli_thread_range_slrr()
- Renamed
    bli_is_last_iter() -> bli_is_last_iter_slrr()
- Defined
    bli_info_get_thread_jrir_tlb()
  and renamed:
  - bli_info_get_thread_part_jrir_slab() ->
    bli_info_get_thread_jrir_slab()
  - bli_info_get_thread_part_jrir_rr() ->
    bli_info_get_thread_jrir_rr()
- Modified bli_rntm_set_ways_for_op() to redirect IR loop parallelism
  into the JR loop when tlb is enabled for non-trsm level-3 operations.
- Added a sanity check to prevent bli_prune_unref_mparts() from being
  used on packed objects. This prohibition is necessary because the
  current implementation does not take into account the atomicity of
  packed micropanel widths relative to the diagonal of structured
  matrices. That is, the function prunes greedily without regard to
  whether doing so would prune off part of a micropanel *which has
  already been packed* and assigned to a thread for inclusion in the
  computation.
- Further restricted early returns in bli_prune_unref_mparts() to
  situations where the primary matrix is not only of general structure
  but also dense (in terms of its uplo_t value). The addition of the
  matrix's dense-ness to the conditional is required because gemmt is
  somewhat unusual in that its C matrix has general structure but is
  marked as lower- or upper-stored via its uplo_t. By only checking
  for general structure, attempts to prune gemmt C matrices would
  incorrectly result in early returns, even though that operation
  effectively treats the matrix as symmetric (and stored in only one
  triangle).
- Fixed a latent bug in bli_thread_range_rr() wherein incorrect ranges
  were computed when 1 < bf. Thankfully, this bug was not yet
  manifesting since all current invocations used bf == 1.
- Fixed a latent bug in some unexercised code in bli_?gemmt_l_ker_var2()
  that would perform incorrect pruning of unreferenced regions above
  where the diagonal of a lower-stored matrix intersects the right edge.
  Thankfully, the bug was not harming anything since those unreferenced
  regions were being pruned prior to the macrokernel.
- Rewrote slab/rr-based gemmt macrokernels so that they no longer carved
  C into rectangular and diagonal regions prior to parallelizing each
  separately. The new macrokernels use a unified loop structure where
  quadratic (slab) partitioning is used.
- Updated all level-3 macrokernels to have a more uniform coding style,
  such as wrt combining variable declarations with initializations as
  well as the use of const.
- Updated bls_l3_packm_var[123].c to use bli_thrinfo_n_way() and
  bli_thrinfo_work_id() instead of bli_thrinfo_num_threads() and
  bli_thrinfo_thread_id(), respectively. This change probably should
  have been included in aeb5f0c.
- Removed old prototypes in bli_gemmt_var.h and bli_trmm_var.h that
  corresponded to functions that were removed in aeb5f0c.
- Other very minor cleanups.
- Comment updates.
  • Loading branch information
fgvanzee authored Jan 11, 2023
1 parent b6735ca commit 2e1ba9d
Show file tree
Hide file tree
Showing 67 changed files with 10,597 additions and 2,022 deletions.
4 changes: 4 additions & 0 deletions build/bli_config.h.in
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,10 @@
#define BLIS_ENABLE_JRIR_RR
#endif

#if @enable_jrir_tlb@
#define BLIS_ENABLE_JRIR_TLB
#endif

#if @enable_pba_pools@
#define BLIS_ENABLE_PBA_POOLS
#else
Expand Down
61 changes: 43 additions & 18 deletions configure
Original file line number Diff line number Diff line change
Expand Up @@ -340,16 +340,36 @@ print_usage()
echo " "
echo " -r METHOD, --thread-part-jrir=METHOD"
echo " "
echo " Request a method of assigning micropanels to threads in"
echo " the JR and IR loops. Valid values for METHOD are 'slab'"
echo " and 'rr'. Using 'slab' assigns (as much as possible)"
echo " contiguous regions of micropanels to each thread while"
echo " using 'rr' assigns micropanels to threads in a round-"
echo " robin fashion. The chosen method also applies during"
echo " the packing of A and B. The default method is 'slab'."
echo " NOTE: Specifying this option constitutes a request,"
echo " which may be ignored in select situations if the"
echo " implementation has a good reason to do so."
echo " Select a strategy for partitioning computation in JR and"
echo " IR loops and assigning that computation to threads. Valid"
echo " values for METHOD are 'rr', 'slab', and 'tlb':"
echo " 'rr': Assign the computation associated with whole"
echo " columns of microtiles to threads in a round-"
echo " robin fashion. When selected, round-robin"
echo " assignment is also employed during packing."
echo " 'slab': Partition the computation into N contiguous"
echo " regions, where each region contains a whole"
echo " number of microtile columns, and assign one"
echo " region to each thread. For some operations, the"
echo " number of microtile columns contained within a"
echo " given region may differ from that of other"
echo " regions, depending on how much work is implied"
echo " by each region. When selected, slab assignment"
echo " is also employed during packing."
echo " 'tlb': Tile-level load balancing is similar to slab,"
echo " except that regions will be divided at a more"
echo " granular level (individual microtiles instead"
echo " of whole columns of microtiles) to ensure more"
echo " equitable assignment of work to threads. When"
echo " selected, tlb will only be employed for level-3"
echo " operations except trsm; due to practical and"
echo " algorithmic limitations, slab partitioning will"
echo " be used instead during packing and for trsm."
echo " The default strategy is 'slab'. NOTE: Specifying this"
echo " option constitutes a request, which may be ignored in"
echo " select situations if implementation has a good reason to"
echo " do so. (See description of 'tlb' above for an example of"
echo " this.)"
echo " "
echo " --disable-trsm-preinversion, --enable-trsm-preinversion"
echo " "
Expand Down Expand Up @@ -3731,16 +3751,20 @@ main()

# Check the method of assigning micropanels to threads in the JR and IR
# loops.
enable_jrir_slab_01=0
enable_jrir_rr_01=0
if [ "x${thread_part_jrir}" = "xslab" ]; then
echo "${script_name}: requesting slab threading in jr and ir loops."
enable_jrir_slab_01=1
elif [ "x${thread_part_jrir}" = "xrr" ]; then
echo "${script_name}: requesting round-robin threading in jr and ir loops."
enable_jrir_slab_01=0
enable_jrir_tlb_01=0
if [ "x${thread_part_jrir}" = "xrr" ]; then
echo "${script_name}: requesting round-robin (rr) work partitioning in jr and/or ir loops."
enable_jrir_rr_01=1
elif [ "x${thread_part_jrir}" = "xslab" ]; then
echo "${script_name}: requesting slab work partitioning in jr and/or ir loops."
enable_jrir_slab_01=1
elif [ "x${thread_part_jrir}" = "xtlb" ]; then
echo "${script_name}: requesting tile-level load balancing (tlb) in unified jr+ir loop."
enable_jrir_tlb_01=1
else
echo "${script_name}: *** Unsupported method of thread partitioning in jr and ir loops: ${thread_part_jrir}."
echo "${script_name}: *** Unsupported method of work partitioning in jr/ir loops: ${thread_part_jrir}."
exit 1
fi

Expand Down Expand Up @@ -4177,8 +4201,9 @@ main()
| sed -e "s/@enable_pthreads_as_def@/${enable_pthreads_as_def_01}/g" \
| sed -e "s/@enable_hpx@/${enable_hpx_01}/g" \
| sed -e "s/@enable_hpx_as_def@/${enable_hpx_as_def_01}/g" \
| sed -e "s/@enable_jrir_slab@/${enable_jrir_slab_01}/g" \
| sed -e "s/@enable_jrir_rr@/${enable_jrir_rr_01}/g" \
| sed -e "s/@enable_jrir_slab@/${enable_jrir_slab_01}/g" \
| sed -e "s/@enable_jrir_tlb@/${enable_jrir_tlb_01}/g" \
| sed -e "s/@enable_pba_pools@/${enable_pba_pools_01}/g" \
| sed -e "s/@enable_sba_pools@/${enable_sba_pools_01}/g" \
| sed -e "s/@enable_mem_tracing@/${enable_mem_tracing_01}/g" \
Expand Down
1 change: 0 additions & 1 deletion frame/1m/packm/bli_packm.h
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@
#include "bli_packm_init.h"
#include "bli_packm_int.h"
#include "bli_packm_scalar.h"
#include "bli_packm_thrinfo.h"

#include "bli_packm_part.h"

Expand Down
16 changes: 8 additions & 8 deletions frame/1m/packm/bli_packm_blk_var1.c
Original file line number Diff line number Diff line change
Expand Up @@ -170,11 +170,11 @@ void bli_packm_blk_var1
const dim_t tid = bli_thrinfo_work_id( thread );

// Determine the thread range and increment using the current thread's
// packm thrinfo_t node. NOTE: The definition of bli_thread_range_jrir()
// packm thrinfo_t node. NOTE: The definition of bli_thread_range_slrr()
// will depend on whether slab or round-robin partitioning was requested
// at configure-time.
dim_t it_start, it_end, it_inc;
bli_thread_range_jrir( thread, n_iter, 1, FALSE, &it_start, &it_end, &it_inc );
bli_thread_range_slrr( thread, n_iter, 1, FALSE, &it_start, &it_end, &it_inc );

char* p_begin = p_cast;

Expand All @@ -195,10 +195,10 @@ void bli_packm_blk_var1

char* c_begin = c_cast + (ic )*incc*dt_c_size;

// Hermitian/symmetric and general packing may use slab or
// round-robin (bli_packm_my_iter()), depending on which was
// selected at configure-time.
if ( bli_packm_my_iter( it, it_start, it_end, tid, nt ) )
// Hermitian/symmetric and general packing may use slab or round-
// robin (bli_is_my_iter()), depending on which was selected at
// configure-time.
if ( bli_is_my_iter( it, it_start, it_end, tid, nt ) )
{
packm_ker_cast( bli_is_triangular( strucc ) ? BLIS_GENERAL : strucc,
diagc,
Expand Down Expand Up @@ -286,9 +286,9 @@ void bli_packm_blk_var1
// We nudge the imaginary stride up by one if it is odd.
is_p_use += ( bli_is_odd( is_p_use ) ? 1 : 0 );

// NOTE: We MUST use round-robin work allocation (bli_packm_my_iter_rr())
// NOTE: We MUST use round-robin work allocation (bli_is_my_iter_rr())
// when packing micropanels of a triangular matrix.
if ( bli_packm_my_iter_rr( it, it_start, it_end, tid, nt ) )
if ( bli_is_my_iter_rr( it, tid, nt ) )
{
packm_ker_cast( strucc,
diagc,
Expand Down
16 changes: 8 additions & 8 deletions frame/3/bli_l3_sup_packm_var.c
Original file line number Diff line number Diff line change
Expand Up @@ -155,10 +155,10 @@ void PASTEMAC(ch,varname) \
dim_t it_start, it_end, it_inc; \
\
/* Determine the thread range and increment using the current thread's
packm thrinfo_t node. NOTE: The definition of bli_thread_range_jrir()
packm thrinfo_t node. NOTE: The definition of bli_thread_range_slrr()
will depend on whether slab or round-robin partitioning was requested
at configure-time. */ \
bli_thread_range_jrir( thread, n_iter, 1, FALSE, &it_start, &it_end, &it_inc ); \
bli_thread_range_slrr( thread, n_iter, 1, FALSE, &it_start, &it_end, &it_inc ); \
\
/* Iterate over every logical micropanel in the source matrix. */ \
for ( ic = ic0, it = 0; it < n_iter; \
Expand All @@ -175,9 +175,9 @@ void PASTEMAC(ch,varname) \
panel_len_i = panel_len_full; \
panel_len_max_i = panel_len_max; \
\
/* The definition of bli_packm_my_iter() will depend on whether slab
/* The definition of bli_is_my_iter() will depend on whether slab
or round-robin partitioning was requested at configure-time. */ \
if ( bli_packm_my_iter( it, it_start, it_end, tid, nt ) ) \
if ( bli_is_my_iter( it, it_start, it_end, tid, nt ) ) \
{ \
f \
( \
Expand Down Expand Up @@ -398,10 +398,10 @@ void PASTEMAC(ch,varname) \
dim_t it_start, it_end, it_inc; \
\
/* Determine the thread range and increment using the current thread's
packm thrinfo_t node. NOTE: The definition of bli_thread_range_jrir()
packm thrinfo_t node. NOTE: The definition of bli_thread_range_slrr()
will depend on whether slab or round-robin partitioning was requested
at configure-time. */ \
bli_thread_range_jrir( thread, n_iter, 1, FALSE, &it_start, &it_end, &it_inc ); \
bli_thread_range_slrr( thread, n_iter, 1, FALSE, &it_start, &it_end, &it_inc ); \
\
/* Iterate over every logical micropanel in the source matrix. */ \
for ( it = 0; it < n_iter; it += 1 ) \
Expand All @@ -412,9 +412,9 @@ void PASTEMAC(ch,varname) \
ctype* p_use = p_begin; \
\
{ \
/* The definition of bli_packm_my_iter() will depend on whether slab
/* The definition of bli_is_my_iter() will depend on whether slab
or round-robin partitioning was requested at configure-time. */ \
if ( bli_packm_my_iter( it, it_start, it_end, tid, nt ) ) \
if ( bli_is_my_iter( it, it_start, it_end, tid, nt ) ) \
{ \
PASTEMAC2(ch,scal2v,BLIS_TAPI_EX_SUF) \
( \
Expand Down
4 changes: 2 additions & 2 deletions frame/3/bli_l3_sup_var12.c
Original file line number Diff line number Diff line change
Expand Up @@ -357,11 +357,11 @@ void PASTEMAC(ch,varname) \
object. */ \
/*
ctype* a2 = bli_gemm_get_next_a_upanel( a_ir, irstep_a, ir_inc ); \
if ( bli_is_last_iter( i, ir_iter, 0, 1 ) ) \
if ( bli_is_last_iter_slrr( i, ir_iter, 0, 1 ) ) \
{ \
a2 = a_00; \
b2 = bli_gemm_get_next_b_upanel( b_jr, jrstep_b, jr_inc ); \
if ( bli_is_last_iter( j, jr_iter, 0, 1 ) ) \
if ( bli_is_last_iter_slrr( j, jr_iter, 0, 1 ) ) \
b2 = b_00; \
} \
\
Expand Down
12 changes: 6 additions & 6 deletions frame/3/bli_l3_thrinfo.h
Original file line number Diff line number Diff line change
Expand Up @@ -39,22 +39,22 @@

// gemm

// NOTE: The definition of bli_gemm_get_next_?_upanel() does not need to
// change depending on BLIS_ENABLE_JRIR_SLAB / BLIS_ENABLE_JRIR_RR.
#define bli_gemm_get_next_a_upanel( a1, step, inc ) ( a1 + step * inc )
#define bli_gemm_get_next_b_upanel( b1, step, inc ) ( b1 + step * inc )

// gemmt

// NOTE: The definition of bli_gemmt_get_next_?_upanel() does not need to
// change depending on BLIS_ENABLE_JRIR_SLAB / BLIS_ENABLE_JRIR_RR.
#define bli_gemmt_get_next_a_upanel( a1, step, inc ) ( a1 + step * inc )
#define bli_gemmt_get_next_b_upanel( b1, step, inc ) ( b1 + step * inc )

// NOTE: Here, we assume NO parallelism in the IR loop.
#define bli_gemmt_l_wrap_a_upanel( a0, step, doff_j, mr, nr ) \
( a0 + ( (-doff_j + 1*nr) / mr ) * step )
#define bli_gemmt_u_wrap_a_upanel( a0, step, doff_j, mr, nr ) \
( a0 )

// trmm

// NOTE: The definition of bli_trmm_get_next_?_upanel() does not need to
// change depending on BLIS_ENABLE_JRIR_SLAB / BLIS_ENABLE_JRIR_RR.
#define bli_trmm_get_next_a_upanel( a1, step, inc ) ( a1 + step * inc )
#define bli_trmm_get_next_b_upanel( b1, step, inc ) ( b1 + step * inc )

Expand Down
23 changes: 19 additions & 4 deletions frame/3/gemm/bli_gemm_cntl.c
Original file line number Diff line number Diff line change
Expand Up @@ -61,10 +61,25 @@ cntl_t* bli_gemmbp_cntl_create
void_fp macro_kernel_fp;

// Choose the default macrokernel based on the operation family...
if ( family == BLIS_GEMM ) macro_kernel_fp = bli_gemm_ker_var2;
else if ( family == BLIS_GEMMT ) macro_kernel_fp = bli_gemmt_x_ker_var2;
else if ( family == BLIS_TRMM ) macro_kernel_fp = bli_trmm_xx_ker_var2;
else /* should never execute */ macro_kernel_fp = NULL;
if ( family == BLIS_GEMM ) macro_kernel_fp =
#ifdef BLIS_ENABLE_JRIR_TLB
bli_gemm_ker_var2b;
#else // ifdef ( _SLAB || _RR )
bli_gemm_ker_var2;
#endif
else if ( family == BLIS_GEMMT ) macro_kernel_fp =
#ifdef BLIS_ENABLE_JRIR_TLB
bli_gemmt_x_ker_var2b;
#else // ifdef ( _SLAB || _RR )
bli_gemmt_x_ker_var2;
#endif
else if ( family == BLIS_TRMM ) macro_kernel_fp =
#ifdef BLIS_ENABLE_JRIR_TLB
bli_trmm_xx_ker_var2b;
#else // ifdef ( _SLAB || _RR )
bli_trmm_xx_ker_var2;
#endif
else /* should never execute */ macro_kernel_fp = NULL;

// ...unless a non-NULL kernel function pointer is passed in, in which
// case we use that instead.
Expand Down
Loading

0 comments on commit 2e1ba9d

Please sign in to comment.