v4.1 release
This commit is contained in:
Junkai-Wu
@@ -325,21 +325,6 @@ struct TiledCopy : Copy_Atom
|
||||
return tile2thrfrg(ref_S, right_inverse(AtomLayoutRef{}).compose(AtomLayoutSrc{}))(_,_,Int<0>{});
|
||||
}
|
||||
|
||||
CUTE_HOST_DEVICE constexpr static
|
||||
auto
|
||||
get_layoutS_MN()
|
||||
{
|
||||
// (thr_idx,val_idx) -> (M,N)
|
||||
auto layoutS_TV = get_layoutS_TV();
|
||||
// (M,K) -> (thr_idx,val_idx)
|
||||
auto layoutS_MK = right_inverse(layoutS_TV).with_shape(shape(Tiler_MN{}));
|
||||
|
||||
// athrid = (v,m,k) -> thr_idx
|
||||
auto thrID_S = make_layout(size<0>(TiledLayout_TV{}));
|
||||
|
||||
return cute::make_tuple(layoutS_MK, thrID_S);
|
||||
}
|
||||
|
||||
CUTE_HOST_DEVICE constexpr static
|
||||
auto
|
||||
get_layoutD_TV()
|
||||
@@ -350,21 +335,6 @@ struct TiledCopy : Copy_Atom
|
||||
return tile2thrfrg(ref_D, right_inverse(AtomLayoutRef{}).compose(AtomLayoutDst{}))(_,_,Int<0>{});
|
||||
}
|
||||
|
||||
CUTE_HOST_DEVICE constexpr static
|
||||
auto
|
||||
get_layoutD_MN()
|
||||
{
|
||||
// (thr_idx,val_idx) -> (M,N)
|
||||
auto layoutD_TV = get_layoutD_TV();
|
||||
// (M,K) -> (thr_idx,val_idx)
|
||||
auto layoutD_MK = right_inverse(layoutD_TV).with_shape(shape(Tiler_MN{}));
|
||||
|
||||
// athrid = (v,m,k) -> thr_idx
|
||||
auto thrID_D = make_layout(size<0>(TiledLayout_TV{}));
|
||||
|
||||
return cute::make_tuple(layoutD_MK, thrID_D);
|
||||
}
|
||||
|
||||
template <class ThrIdx,
|
||||
__CUTE_REQUIRES(is_integral<ThrIdx>::value)>
|
||||
CUTE_HOST_DEVICE static
|
||||
@@ -680,101 +650,6 @@ print(ThrCopy<TiledCopy, ThrIdx> const& thr_copy)
|
||||
print(TiledCopy{});
|
||||
}
|
||||
|
||||
// TiledCopy to LaTeX TikZ
|
||||
template <class... Args, class TikzColorFn = TikzColor_TV>
|
||||
CUTE_HOST_DEVICE
|
||||
auto
|
||||
print_latex(TiledCopy<Args...> const& copy,
|
||||
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
|
||||
{
|
||||
auto [layoutS_MN, thrID_S] = copy.get_layoutS_MN();
|
||||
auto [layoutD_MN, thrID_D] = copy.get_layoutD_MN();
|
||||
|
||||
print_latex_copy(layoutS_MN, thrID_S,
|
||||
layoutD_MN, thrID_D);
|
||||
}
|
||||
|
||||
// MNK Copy Layout to LaTeX TikZ
|
||||
template <class LayoutS, class ThrIDS,
|
||||
class LayoutD, class ThrIDD,
|
||||
class TikzColorFn = TikzColor_TV>
|
||||
CUTE_HOST_DEVICE
|
||||
void
|
||||
print_latex_copy(LayoutS const& S, ThrIDS const& TS, // (m,n) -> (tid,vid) and tid -> thr_idx
|
||||
LayoutD const& D, ThrIDD const& TD, // (m,n) -> (tid,vid) and tid -> thr_idx
|
||||
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
|
||||
{
|
||||
CUTE_STATIC_ASSERT_V(rank(S) == Int<2>{});
|
||||
CUTE_STATIC_ASSERT_V(rank(D) == Int<2>{});
|
||||
|
||||
assert(size<0>(S) == size<0>(D));
|
||||
assert(size<1>(S) == size<1>(D));
|
||||
|
||||
// Commented prints
|
||||
printf("%% LayoutS: "); print(S); printf("\n");
|
||||
printf("%% ThrIDS : "); print(TS); printf("\n");
|
||||
printf("%% LayoutD: "); print(D); printf("\n");
|
||||
printf("%% ThrIDD : "); print(TD); printf("\n\n");
|
||||
|
||||
// Header
|
||||
printf("\\documentclass[convert]{standalone}\n"
|
||||
"\\usepackage{tikz}\n\n"
|
||||
"\\begin{document}\n"
|
||||
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n\n");
|
||||
|
||||
// S starting at 0,0
|
||||
for (int i = 0; i < size<0>(S); ++i) {
|
||||
for (int j = 0; j < size<1>(S); ++j) {
|
||||
int thrid = S(i,j) % size(TS);
|
||||
int val_idx = S(i,j) / size(TS);
|
||||
int thr_idx = TS(thrid);
|
||||
|
||||
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
|
||||
color(thr_idx, val_idx),
|
||||
i, j,
|
||||
thr_idx, val_idx);
|
||||
}
|
||||
}
|
||||
// Grid
|
||||
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
|
||||
0, 0, int(size<0>(S)), int(size<1>(S)));
|
||||
// S Labels
|
||||
for (int i = 0, j = -1; i < size<0>(S); ++i) {
|
||||
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i);
|
||||
}
|
||||
for (int i = -1, j = 0; j < size<1>(S); ++j) {
|
||||
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j);
|
||||
}
|
||||
|
||||
// D starting at 0,size<1>(S)+3
|
||||
for (int i = 0; i < size<0>(D); ++i) {
|
||||
for (int j = 0; j < size<1>(D); ++j) {
|
||||
int thrid = D(i,j) % size(TD);
|
||||
int val_idx = D(i,j) / size(TD);
|
||||
int thr_idx = TD(thrid);
|
||||
|
||||
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
|
||||
color(thr_idx, val_idx),
|
||||
i, j + size<1>(S) + 3,
|
||||
thr_idx, val_idx);
|
||||
}
|
||||
}
|
||||
// Grid
|
||||
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
|
||||
0, int(size<1>(S)+3), int(size<0>(D)), int(size<1>(D)+size<1>(S)+3));
|
||||
// D Labels
|
||||
for (int i = 0, j = size<1>(D); i < size<0>(D); ++i) {
|
||||
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j + size<1>(S) + 3, i);
|
||||
}
|
||||
for (int i = -1, j = 0; j < size<1>(D); ++j) {
|
||||
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j + size<1>(S) + 3, j);
|
||||
}
|
||||
|
||||
// Footer
|
||||
printf("\\end{tikzpicture}\n"
|
||||
"\\end{document}\n");
|
||||
}
|
||||
|
||||
} // end namespace cute
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
+13
-429
@@ -180,10 +180,10 @@ struct MMA_Atom<MMA_Traits<MMAOperation, Args...>>
|
||||
if constexpr (has_dereference<FrgTypeB>::value) {
|
||||
// If the intended FrgTypeB is a view (of the current tensor), forward the whole
|
||||
static_assert(is_same<ValTypeB, typename remove_cvref_t<BTensor>::value_type>::value
|
||||
|
||||
|
||||
|| (sizeof_bits_v<typename remove_cvref_t<BTensor>::value_type> == 8 &&
|
||||
(sizeof_bits_v<ValTypeB> == 8 || sizeof_bits_v<ValTypeB> == 6 || sizeof_bits_v<ValTypeB> == 4))
|
||||
|
||||
|
||||
, "Expecting ValTypeB type");
|
||||
return make_tensor<FrgTypeB>(static_cast<BTensor&&>(btensor));
|
||||
} else {
|
||||
@@ -394,55 +394,22 @@ struct TiledMMA : MMA_Atom
|
||||
return size(permutation_mnk<I>());
|
||||
}
|
||||
|
||||
CUTE_HOST_DEVICE constexpr
|
||||
auto
|
||||
get_layoutC_MN() const
|
||||
{
|
||||
// (M,N) -> (M,N)
|
||||
auto ref_C = make_layout(make_shape(tile_size_mnk<0>(), tile_size_mnk<1>()));
|
||||
// (cthrid,val) -> (M,N)
|
||||
auto layoutC_TV = thrfrg_C(ref_C);
|
||||
// (M,N) -> (cthrid,frg)
|
||||
auto layoutC_MN = right_inverse(layoutC_TV).with_shape(shape(ref_C));
|
||||
|
||||
// cthrid = (v,m,n) -> thr_idx
|
||||
auto thrID_C = thr_layout_vmnk_(_,_,_,Int<0>{});
|
||||
|
||||
return cute::make_tuple(layoutC_MN, thrID_C);
|
||||
}
|
||||
|
||||
CUTE_HOST_DEVICE constexpr
|
||||
auto
|
||||
get_layoutC_TV() const
|
||||
{
|
||||
// (M,N) -> (M,N)
|
||||
auto ref_C = make_layout(make_shape(tile_size_mnk<0>(), tile_size_mnk<1>()));
|
||||
// (cthrid,val) -> (M,N)
|
||||
auto layoutC_TV = thrfrg_C(ref_C);
|
||||
|
||||
// thr_idx -> (ThrV,ThrM,ThrN,ThrK)
|
||||
auto thridx_2_thrid = right_inverse(thr_layout_vmnk_);
|
||||
auto thridx_2_thrid = composition(make_layout(make_shape (size(thr_layout_vmnk_), Int<1>{}),
|
||||
make_stride(Int<1>{}, Int<0>{})),
|
||||
right_inverse(make_layout(thr_layout_vmnk_, complement(thr_layout_vmnk_))));
|
||||
|
||||
// (thr_idx,val) -> (M,N)
|
||||
return layoutC_TV.compose(thridx_2_thrid, _);
|
||||
return thrfrg_C(ref_C).compose(thridx_2_thrid, _);
|
||||
}
|
||||
|
||||
CUTE_HOST_DEVICE constexpr
|
||||
auto
|
||||
get_layoutA_MK() const
|
||||
{
|
||||
// (M,K) -> (M,K)
|
||||
auto ref_A = make_layout(make_shape(tile_size_mnk<0>(), tile_size_mnk<2>()));
|
||||
// (athrid,val) -> (M,K)
|
||||
auto layoutA_TV = thrfrg_A(ref_A);
|
||||
// (M,K) -> (athrid,frg)
|
||||
auto layoutA_MK = right_inverse(layoutA_TV).with_shape(shape(ref_A));
|
||||
|
||||
// athrid = (v,m,k) -> thr_idx
|
||||
auto thrID_A = thr_layout_vmnk_(_,_,Int<0>{},_);
|
||||
|
||||
return cute::make_tuple(layoutA_MK, thrID_A);
|
||||
}
|
||||
|
||||
CUTE_HOST_DEVICE constexpr
|
||||
auto
|
||||
@@ -458,29 +425,14 @@ struct TiledMMA : MMA_Atom
|
||||
_));
|
||||
|
||||
// thr_idx -> (ThrV,ThrM,ThrN,ThrK)
|
||||
auto thridx_2_thrid = right_inverse(thr_layout_vmnk_);
|
||||
auto thridx_2_thrid = composition(make_layout(make_shape (size(thr_layout_vmnk_), Int<1>{}),
|
||||
make_stride(Int<1>{}, Int<0>{})),
|
||||
right_inverse(make_layout(thr_layout_vmnk_, complement(thr_layout_vmnk_))));
|
||||
|
||||
// (thr_idx,val) -> (M,K)
|
||||
return thrfrg_A(ref_A).compose(atile, _).compose(thridx_2_thrid, _);
|
||||
}
|
||||
|
||||
CUTE_HOST_DEVICE constexpr
|
||||
auto
|
||||
get_layoutB_NK() const
|
||||
{
|
||||
// (N,K) -> (N,K)
|
||||
auto ref_B = make_layout(make_shape(tile_size_mnk<1>(), tile_size_mnk<2>()));
|
||||
// (bthrid,val) -> (N,K)
|
||||
auto layoutB_TV = thrfrg_B(ref_B);
|
||||
// (N,K) -> (bthrid,frg)
|
||||
auto layoutB_NK = right_inverse(layoutB_TV).with_shape(shape(ref_B));
|
||||
|
||||
// bthrid = (v,n,k) -> thr_idx
|
||||
auto thrID_B = thr_layout_vmnk_(_,Int<0>{},_,_);
|
||||
|
||||
return cute::make_tuple(layoutB_NK, thrID_B);
|
||||
}
|
||||
|
||||
CUTE_HOST_DEVICE constexpr
|
||||
auto
|
||||
get_layoutB_TV() const
|
||||
@@ -495,7 +447,9 @@ struct TiledMMA : MMA_Atom
|
||||
_));
|
||||
|
||||
// thr_idx -> (ThrV,ThrM,ThrN,ThrK)
|
||||
auto thridx_2_thrid = right_inverse(thr_layout_vmnk_);
|
||||
auto thridx_2_thrid = composition(make_layout(make_shape (size(thr_layout_vmnk_), Int<1>{}),
|
||||
make_stride(Int<1>{}, Int<0>{})),
|
||||
right_inverse(make_layout(thr_layout_vmnk_, complement(thr_layout_vmnk_))));
|
||||
|
||||
// (thr_idx,val) -> (N,K)
|
||||
return thrfrg_B(ref_B).compose(btile, _).compose(thridx_2_thrid, _);
|
||||
@@ -733,376 +687,6 @@ print(ThrMMA<TiledMMA, ThrVMNK> const& thr_mma)
|
||||
print(static_cast<TiledMMA>(thr_mma));
|
||||
}
|
||||
|
||||
// MMA Atom to LaTeX TikZ
|
||||
template <class... Args, class TikzColorFn = TikzColor_TV>
|
||||
CUTE_HOST_DEVICE
|
||||
void
|
||||
print_latex(MMA_Atom<Args...> const& mma_atom,
|
||||
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
|
||||
{
|
||||
print_latex(make_tiled_mma(mma_atom));
|
||||
}
|
||||
|
||||
// TiledMMA to LaTeX TikZ
|
||||
template <class... Args, class TikzColorFn = TikzColor_TV>
|
||||
CUTE_HOST_DEVICE
|
||||
void
|
||||
print_latex(TiledMMA<Args...> const& mma,
|
||||
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
|
||||
{
|
||||
auto layout_and_thrid_C = mma.get_layoutC_MN();
|
||||
auto layoutC_MN = get<0>(layout_and_thrid_C);
|
||||
auto thrID_C = get<1>(layout_and_thrid_C);
|
||||
|
||||
auto layout_and_thrid_A = mma.get_layoutA_MK();
|
||||
auto layoutA_MK = get<0>(layout_and_thrid_A);
|
||||
auto thrID_A = get<1>(layout_and_thrid_A);
|
||||
|
||||
auto layout_and_thrid_B = mma.get_layoutB_NK();
|
||||
auto layoutB_NK = get<0>(layout_and_thrid_B);
|
||||
auto thrID_B = get<1>(layout_and_thrid_B);
|
||||
|
||||
print_latex_mma(layoutC_MN, thrID_C,
|
||||
layoutA_MK, thrID_A,
|
||||
layoutB_NK, thrID_B);
|
||||
}
|
||||
|
||||
// MNK MMA Layout to LaTeX TikZ
|
||||
template <class LayoutC, class ThrIDC,
|
||||
class LayoutA, class ThrIDA,
|
||||
class LayoutB, class ThrIDB,
|
||||
class TikzColorFn = TikzColor_TV>
|
||||
CUTE_HOST_DEVICE
|
||||
void
|
||||
print_latex_mma(LayoutC const& C, ThrIDC const& TC, // (m,n) -> (tid,vid) and tid -> thr_idx
|
||||
LayoutA const& A, ThrIDA const& TA, // (m,k) -> (tid,vid) and tid -> thr_idx
|
||||
LayoutB const& B, ThrIDB const& TB, // (n,k) -> (tid,vid) and tid -> thr_idx
|
||||
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
|
||||
{
|
||||
CUTE_STATIC_ASSERT_V(rank(C) == Int<2>{});
|
||||
CUTE_STATIC_ASSERT_V(rank(A) == Int<2>{});
|
||||
CUTE_STATIC_ASSERT_V(rank(B) == Int<2>{});
|
||||
|
||||
assert(size<0>(A) == size<0>(C));
|
||||
assert(size<0>(B) == size<1>(C));
|
||||
assert(size<1>(A) == size<1>(B));
|
||||
|
||||
// Commented prints
|
||||
printf("%% LayoutC: "); print(C); printf("\n");
|
||||
printf("%% ThrIDC : "); print(TC); printf("\n");
|
||||
printf("%% LayoutA: "); print(A); printf("\n");
|
||||
printf("%% ThrIDA : "); print(TA); printf("\n");
|
||||
printf("%% LayoutB: "); print(B); printf("\n");
|
||||
printf("%% ThrIDB : "); print(TB); printf("\n\n");
|
||||
// Header
|
||||
printf("\\documentclass[convert]{standalone}\n"
|
||||
"\\usepackage{tikz}\n\n"
|
||||
"\\begin{document}\n"
|
||||
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n\n");
|
||||
|
||||
// C starting at 0,0
|
||||
for (int m = 0; m < size<0>(C); ++m) {
|
||||
for (int n = 0; n < size<1>(C); ++n) {
|
||||
int thrid = C(m,n) % size(TC);
|
||||
int val_idx = C(m,n) / size(TC);
|
||||
int thr_idx = TC(thrid);
|
||||
|
||||
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
|
||||
color(thr_idx, val_idx),
|
||||
m, n,
|
||||
thr_idx, val_idx);
|
||||
}
|
||||
}
|
||||
// Grid
|
||||
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
|
||||
0, 0, int(size<0>(C)), int(size<1>(C)));
|
||||
|
||||
// A starting at 0,-size<1>(A)-1
|
||||
for (int m = 0; m < size<0>(A); ++m) {
|
||||
for (int k = 0; k < size<1>(A); ++k) {
|
||||
int thrid = A(m,k) % size(TA);
|
||||
int val_idx = A(m,k) / size(TA);
|
||||
int thr_idx = TA(thrid);
|
||||
|
||||
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
|
||||
color(thr_idx, val_idx),
|
||||
m, k-1-size<1>(A),
|
||||
thr_idx, val_idx);
|
||||
}
|
||||
}
|
||||
// Grid
|
||||
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
|
||||
0, int(-size<1>(A)-1), int(size<0>(A)), -1);
|
||||
// A labels
|
||||
for (int m = 0, k = -1; m < size<0>(A); ++m) {
|
||||
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), m);
|
||||
}
|
||||
for (int m = -1, k = 0; k < size<1>(A); ++k) {
|
||||
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), k);
|
||||
}
|
||||
|
||||
// B starting at -size<1>(B)-1,0
|
||||
for (int n = 0; n < size<0>(B); ++n) {
|
||||
for (int k = 0; k < size<1>(B); ++k) {
|
||||
int thrid = B(n,k) % size(TB);
|
||||
int val_idx = B(n,k) / size(TB);
|
||||
int thr_idx = TB(thrid);
|
||||
|
||||
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
|
||||
color(thr_idx, val_idx),
|
||||
k-1-size<1>(B), n,
|
||||
thr_idx, val_idx);
|
||||
}
|
||||
}
|
||||
// Grid
|
||||
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
|
||||
int(-size<1>(B)-1), 0, -1, int(size<0>(B)));
|
||||
// B labels
|
||||
for (int n = 0, k = -1; n < size<0>(B); ++n) {
|
||||
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, n);
|
||||
}
|
||||
for (int n = -1, k = 0; k < size<1>(B); ++k) {
|
||||
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, k);
|
||||
}
|
||||
|
||||
// Footer
|
||||
printf("\\end{tikzpicture}\n"
|
||||
"\\end{document}\n");
|
||||
}
|
||||
|
||||
// MNK MMA Layout to console printer
|
||||
template <class LayoutC, class ThrIDC,
|
||||
class LayoutA, class ThrIDA,
|
||||
class LayoutB, class ThrIDB>
|
||||
CUTE_HOST_DEVICE
|
||||
void
|
||||
print_layout_mma(LayoutC const& C, ThrIDC const& TC, // (m,n) -> (tid,vid) and tid -> thr_idx
|
||||
LayoutA const& A, ThrIDA const& TA, // (m,k) -> (tid,vid) and tid -> thr_idx
|
||||
LayoutB const& B, ThrIDB const& TB) // (n,k) -> (tid,vid) and tid -> thr_idx
|
||||
{
|
||||
CUTE_STATIC_ASSERT_V(rank(C) == Int<2>{});
|
||||
CUTE_STATIC_ASSERT_V(rank(A) == Int<2>{});
|
||||
CUTE_STATIC_ASSERT_V(rank(B) == Int<2>{});
|
||||
|
||||
assert(size<0>(A) == size<0>(C));
|
||||
assert(size<0>(B) == size<1>(C));
|
||||
assert(size<1>(A) == size<1>(B));
|
||||
|
||||
int a_width = size<1>(A) * 6 + 4;
|
||||
|
||||
// Print out B (white-shifted) k-by-n
|
||||
for (int k = 0; k < size<1>(B); ++k) {
|
||||
// Header
|
||||
printf("%*s", a_width, "");
|
||||
for (int n = 0; n < size<0>(B); ++n) printf("+-----");
|
||||
printf("+\n");
|
||||
// Values
|
||||
printf("%*s", a_width, "");
|
||||
for (int n = 0; n < size<0>(B); ++n) printf("|T%02dV%1d", int(TB(B(n,k) % size(TB))), int(B(n,k) / size(TB)));
|
||||
printf("|\n");
|
||||
}
|
||||
// Footer
|
||||
printf("%*s", a_width, "");
|
||||
for (int n = 0; n < size<0>(B); ++n) printf("+-----");
|
||||
printf("+\n\n");
|
||||
|
||||
// Print out A m-by-k and C m-by-n
|
||||
for (int m = 0; m < size<0>(A); ++m) {
|
||||
// Header
|
||||
for (int k = 0; k < size<1>(A); ++k) printf("+-----");
|
||||
printf("+ ");
|
||||
for (int n = 0; n < size<1>(C); ++n) printf("+-----");
|
||||
printf("+\n");
|
||||
// Values
|
||||
for (int k = 0; k < size<1>(A); ++k) printf("|T%02dV%1d", int(TA(A(m,k) % size(TA))), int(A(m,k) / size(TA)));
|
||||
printf("| ");
|
||||
for (int n = 0; n < size<1>(C); ++n) printf("|T%02dV%1d", int(TC(C(m,n) % size(TC))), int(C(m,n) / size(TC)));
|
||||
printf("|\n");
|
||||
}
|
||||
// Footer
|
||||
for (int k = 0; k < size<1>(A); ++k) printf("+-----");
|
||||
printf("+ ");
|
||||
for (int n = 0; n < size<1>(C); ++n) printf("+-----");
|
||||
printf("+\n");
|
||||
}
|
||||
|
||||
// MNK MMA Layout to SVG -- 8-value color coded by thread
|
||||
template <class LayoutC, class ThrIDC,
|
||||
class LayoutA, class ThrIDA,
|
||||
class LayoutB, class ThrIDB>
|
||||
CUTE_HOST_DEVICE
|
||||
void
|
||||
print_svg_mma(LayoutC const& C, ThrIDC const& TC, // (m,n) -> (tid,vid) and tid -> thr_idx
|
||||
LayoutA const& A, ThrIDA const& TA, // (m,k) -> (tid,vid) and tid -> thr_idx
|
||||
LayoutB const& B, ThrIDB const& TB) // (n,k) -> (tid,vid) and tid -> thr_idx
|
||||
{
|
||||
char const *color_map[8] = {"175,175,255", "175,255,175", "255,255,175",
|
||||
"255,175,175", "210,210,255", "210,255,210",
|
||||
"255,255,210", "255,210,210"};
|
||||
|
||||
const int cell_width = 20;
|
||||
const int cell_height = 20;
|
||||
|
||||
const int page_width = (size<1>(A) + size<0>(B) + 2) * cell_width;
|
||||
const int page_height = (size<1>(B) + size<0>(A) + 2) * cell_height;
|
||||
|
||||
// header
|
||||
printf("<svg width=\"100%%\" height=\"100%%\" viewBox=\"0 0 %d %d\" "
|
||||
"preserveAspectRatio=\"xMidYMid meet\" "
|
||||
"xmlns=\"http://www.w3.org/2000/svg\">\n",
|
||||
page_width, page_height);
|
||||
|
||||
// C
|
||||
int c_base_x = (size<1>(A) + 2) * cell_width;
|
||||
int c_base_y = (size<1>(B) + 2) * cell_height;
|
||||
for (int m = 0; m < cute::size<0>(C); ++m) {
|
||||
for (int n = 0; n < cute::size<1>(C); ++n) {
|
||||
|
||||
int thrid = C(m, n) % size(TC);
|
||||
int val_idx = C(m, n) / size(TC);
|
||||
int thr_idx = TC(thrid);
|
||||
|
||||
int x = n * cell_width + c_base_x;
|
||||
int y = m * cell_height + c_base_y;
|
||||
|
||||
int thr_x = x + cell_width / 2;
|
||||
int thr_y = y + cell_height / 4;
|
||||
int val_x = x + cell_width / 2;
|
||||
int val_y = y + cell_height * 3 / 4;
|
||||
|
||||
printf("<rect x=\"%d\" y=\"%d\" width=\"%d\" height=\"%d\" "
|
||||
"fill=\"rgb(%s)\" stroke=\"black\"/>\n",
|
||||
x, y, cell_width, cell_height, color_map[thr_idx % 8]);
|
||||
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"8\">T%d</text>\n",
|
||||
thr_x, thr_y, thr_idx);
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"8\">V%d</text>\n",
|
||||
val_x, val_y, val_idx);
|
||||
}
|
||||
}
|
||||
|
||||
// A
|
||||
int a_base_x = cell_width;
|
||||
int a_base_y = (size<1>(B) + 2) * cell_height;
|
||||
for (int m = 0; m < size<0>(A); ++m) {
|
||||
for (int k = 0; k < size<1>(A); ++k) {
|
||||
int thrid = A(m, k) % size(TA);
|
||||
int val_idx = A(m, k) / size(TA);
|
||||
int thr_idx = TA(thrid);
|
||||
|
||||
int x = k * cell_width + a_base_x;
|
||||
int y = m * cell_height + a_base_y;
|
||||
|
||||
int thr_x = x + cell_width / 2;
|
||||
int thr_y = y + cell_height / 4;
|
||||
int val_x = x + cell_width / 2;
|
||||
int val_y = y + cell_height * 3 / 4;
|
||||
|
||||
printf("<rect x=\"%d\" y=\"%d\" width=\"%d\" height=\"%d\" "
|
||||
"fill=\"rgb(%s)\" stroke=\"black\" />\n",
|
||||
x, y, cell_width, cell_height, color_map[thr_idx % 8]);
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"8\">T%d</text>\n",
|
||||
thr_x, thr_y, thr_idx);
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"8\">V%d</text>\n",
|
||||
val_x, val_y, val_idx);
|
||||
}
|
||||
}
|
||||
|
||||
// B
|
||||
int b_base_x = (size<1>(A) + 2) * cell_width;
|
||||
int b_base_y = cell_height;
|
||||
for (int n = 0; n < size<0>(B); ++n) {
|
||||
for (int k = 0; k < size<1>(B); ++k) {
|
||||
int thrid = B(n, k) % size(TB);
|
||||
int val_idx = B(n, k) / size(TB);
|
||||
int thr_idx = TB(thrid);
|
||||
|
||||
int x = n * cell_width + b_base_x;
|
||||
int y = k * cell_height + b_base_y;
|
||||
|
||||
int thr_x = x + cell_width / 2;
|
||||
int thr_y = y + cell_height / 4;
|
||||
int val_x = x + cell_width / 2;
|
||||
int val_y = y + cell_height * 3 / 4;
|
||||
|
||||
printf("<rect x=\"%d\" y=\"%d\" width=\"%d\" height=\"%d\" "
|
||||
"fill=\"rgb(%s)\" stroke=\"black\" />\n",
|
||||
x, y, cell_width, cell_height, color_map[thr_idx % 8]);
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"8\">T%d</text>\n",
|
||||
thr_x, thr_y, thr_idx);
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"8\">V%d</text>\n",
|
||||
val_x, val_y, val_idx);
|
||||
}
|
||||
}
|
||||
|
||||
// A labels
|
||||
for (int m = 0; m < size<0>(A); ++m) {
|
||||
int x = cell_width / 2;
|
||||
int y = m * cell_height + cell_height / 2 + a_base_y;
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"12\">%d</text>\n",
|
||||
x, y, m);
|
||||
}
|
||||
for (int k = 0; k < size<1>(A); ++k) {
|
||||
int x = cell_width + k * cell_width + cell_width / 2;
|
||||
int y = -cell_height / 2 + a_base_y;
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"12\">%d</text>\n",
|
||||
x, y, k);
|
||||
}
|
||||
|
||||
// B labels
|
||||
for (int n = 0; n < size<0>(B); ++n) {
|
||||
int x = b_base_x + cell_width * n + cell_width / 2;
|
||||
int y = cell_height / 2;
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"12\">%d</text>\n",
|
||||
x, y, n);
|
||||
}
|
||||
for (int k = 0; k < size<1>(B); ++k) {
|
||||
int x = b_base_x - cell_width / 2;
|
||||
int y = cell_height * (k + 1) + cell_height / 2;
|
||||
printf("<text x=\"%d\" y=\"%d\" text-anchor=\"middle\" "
|
||||
"alignment-baseline=\"central\" font-size=\"12\">%d</text>\n",
|
||||
x, y, k);
|
||||
}
|
||||
|
||||
// footer
|
||||
printf("</svg>\n");
|
||||
}
|
||||
|
||||
template <class... Args>
|
||||
CUTE_HOST_DEVICE
|
||||
void
|
||||
print_svg(MMA_Atom<Args...> const &mma_atom) {
|
||||
print_svg(make_tiled_mma(mma_atom));
|
||||
}
|
||||
|
||||
template <class... Args>
|
||||
CUTE_HOST_DEVICE
|
||||
void
|
||||
print_svg(TiledMMA<Args...> const &mma) {
|
||||
auto layout_and_thrid_C = mma.get_layoutC_MN();
|
||||
auto layoutC_MN = get<0>(layout_and_thrid_C);
|
||||
auto thrID_C = get<1>(layout_and_thrid_C);
|
||||
|
||||
auto layout_and_thrid_A = mma.get_layoutA_MK();
|
||||
auto layoutA_MK = get<0>(layout_and_thrid_A);
|
||||
auto thrID_A = get<1>(layout_and_thrid_A);
|
||||
|
||||
auto layout_and_thrid_B = mma.get_layoutB_NK();
|
||||
auto layoutB_NK = get<0>(layout_and_thrid_B);
|
||||
auto thrID_B = get<1>(layout_and_thrid_B);
|
||||
|
||||
print_svg_mma(layoutC_MN, thrID_C, layoutA_MK, thrID_A, layoutB_NK, thrID_B);
|
||||
}
|
||||
|
||||
} // namespace cute
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -1114,7 +698,7 @@ print_svg(TiledMMA<Args...> const &mma) {
|
||||
#include <cute/atom/mma_traits_sm89.hpp>
|
||||
#include <cute/atom/mma_traits_sm90.hpp>
|
||||
#include <cute/atom/mma_traits_sm90_gmma.hpp>
|
||||
#include <cute/atom/mma_traits_sm100.hpp>
|
||||
#include <cute/atom/mma_traits_sm100.hpp>
|
||||
#include <cute/atom/mma_traits_sm120.hpp>
|
||||
#include <cute/atom/mma_traits_sm120_sparse.hpp>
|
||||
|
||||
|
||||
@@ -3844,4 +3844,39 @@ struct MMA_Traits<SM100_MMA_MXF4NVF4_2x1SM_SS_SPARSE<a_type, b_type, c_type, sf_
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Specialization for a vectorized FMA per thread.
|
||||
*/
|
||||
template <>
|
||||
struct MMA_Traits<SM100_2x1x1_F32F32F32F32>
|
||||
{
|
||||
using ValTypeD = float;
|
||||
using ValTypeA = float;
|
||||
using ValTypeB = float;
|
||||
using ValTypeC = float;
|
||||
|
||||
using Shape_MNK = Shape<_2,_1,_1>;
|
||||
using ThrID = Layout<_1>;
|
||||
|
||||
using ALayout = Layout<Shape<_1,_2>>;
|
||||
using BLayout = Layout<Shape<_1,_1>>;
|
||||
using CLayout = Layout<Shape<_1,_2>>;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct MMA_Traits<SM100_1x2x1_F32F32F32F32>
|
||||
{
|
||||
using ValTypeD = float;
|
||||
using ValTypeA = float;
|
||||
using ValTypeB = float;
|
||||
using ValTypeC = float;
|
||||
|
||||
using Shape_MNK = Shape<_1,_2,_1>;
|
||||
using ThrID = Layout<_1>;
|
||||
|
||||
using ALayout = Layout<Shape<_1,_1>>;
|
||||
using BLayout = Layout<Shape<_1,_2>>;
|
||||
using CLayout = Layout<Shape<_1,_2>>;
|
||||
};
|
||||
|
||||
} // end namespace cute
|
||||
|
||||
Reference in New Issue
Block a user