Skip to content

Commit

Permalink
check shared memory predicate in matmul tests (#3102)
Browse files Browse the repository at this point in the history 
Adding check to matmul tests since shared memory access in matmul
kernels shouldn't be predicated.
Helps #2339
  • Loading branch information
@liqiangxl
liqiangxl authored Oct 7, 2024
1 parent c7c5de8 commit 61a77e0
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 6 deletions.
76 changes: 70 additions & 6 deletions tests/cpp/test_matmul.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmul) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -184,6 +186,8 @@ TEST_P(MatmulTestWithLayout, AmperePrologueFusionBroadcast) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -241,6 +245,8 @@ TEST_P(MatmulTestWithLayout, AmpereProloguePointwise) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.sin().to(at::kFloat),
Expand Down Expand Up @@ -298,6 +304,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulBFloat16) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -357,6 +365,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulPipelineGmem) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -437,6 +447,8 @@ TEST_P(MatmulTestWithLayout, AmpereSwizzle) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -561,6 +573,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulRegCircularBuffer) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -849,7 +863,8 @@ TEST_F(MatmulTest, MatmulMatmulAmpere) {
fe.compileFusion(&fusion, {t0, t1, t2}, LaunchParams(), matmul_cparams));

auto cg_outputs = fe.runFusion({t0, t1, t2});

ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
// relaxed check for now, err accumulation is significant.
NVF_CHECK(cg_outputs[0].allclose(tref, 0.1, 0.1));
}
Expand Down Expand Up @@ -1228,7 +1243,8 @@ TEST_F(MatmulTest, MatmulSoftmaxMatmulAmpere) {
fe.compileFusion(&fusion, {t0, t1, t2}, LaunchParams(), matmul_cparams));

auto cg_outputs = fe.runFusion({t0, t1, t2});

ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto g1 = t0.to(at::kFloat).matmul(t1.t().to(at::kFloat));
auto sg1 = at::_softmax(g1, -1, false);
auto gsg1 = sg1.matmul(t2.t().to(at::kFloat));
Expand Down Expand Up @@ -1276,6 +1292,8 @@ TEST_P(MatmulTestWithLayout, TuringMatmul) {
NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
7, 5, fe.compileFusion(&fusion, {inputs.first, inputs.second}));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -1421,7 +1439,8 @@ TEST_F(MatmulTest, AmpereMatmulTNCpAsync) {
fe.compileFusion(&fusion, {t0, t1}, LaunchParams(), matmul_cparams));

auto cg_outputs = fe.runFusion({t0, t1});

ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto tref = t0.to(at::kFloat).matmul(t1.t().to(at::kFloat));

NVF_CHECK(cg_outputs[0].allclose(tref, 0.0001, 0.0001));
Expand Down Expand Up @@ -1589,7 +1608,8 @@ TEST_F(MatmulTest, AmpereStridedBatchedMatmulTN) {
fe.compileFusion(&fusion, {t0, t1}, LaunchParams(), matmul_cparams));

auto cg_outputs = fe.runFusion({t0, t1});

ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
// ref implementation:
auto ref_t0 = t0.permute({0, 2, 1, 3})
.contiguous()
Expand Down Expand Up @@ -1761,7 +1781,8 @@ TEST_F(MatmulTest, AmpereViewMatmulTN) {
fe.compileFusion(&fusion, {t0, t1}, LaunchParams(), matmul_cparams));

auto cg_outputs = fe.runFusion({t0, t1});

ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto tref =
at::native::view(t0, {M, K}).to(at::kFloat).matmul(t1.t().to(at::kFloat));

Expand Down Expand Up @@ -1943,7 +1964,8 @@ TEST_F(MatmulTest, AmpereMatmulTNSwizzled) {
FusionExecutor fe;
fe.compileFusion(&fusion, {t0, t1}, LaunchParams(), matmul_cparams);
auto cg_outputs = fe.runFusion({t0, t1});

ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto tref = t0.to(at::kFloat).matmul(t1.t().to(at::kFloat));

NVF_CHECK(cg_outputs[0].allclose(tref, 0.0001, 0.0001));
Expand Down Expand Up @@ -1998,6 +2020,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulLargeLoad) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -2050,6 +2074,8 @@ TEST_P(MatmulTestWithLayout, TuringMatmulLargeLoad) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -2119,6 +2145,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulTileCheck4warp) {
matmul_cparams));
EXPECT_TRUE(getBankConflictInfo(fe.kernel()).empty());
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
NVF_CHECK(
Expand Down Expand Up @@ -2195,6 +2223,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulTileCheck8warp) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -2262,6 +2292,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulTileCheck6warp) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -2318,6 +2350,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulLargeLoadLargeK) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -2368,6 +2402,8 @@ TEST_P(MatmulTestWithLayout, AmpereSplitKLikeStridedBatchedMatmul) {
0,
fe.compileFusion(&fusion, {t0, t1}, LaunchParams(), matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({t0, t1});
auto tref = splitkLikeAtMatmul(t0.to(at::kFloat), t1.to(at::kFloat), layout);
NVF_CHECK(cg_outputs[0].allclose(tref, 0.0001, 0.0001));
Expand Down Expand Up @@ -2464,6 +2500,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulSmemEpilogue) {

// check bank conflicts
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
// (0.001, 0.001) passed on local A100 but failed on CI A100
NVF_CHECK(
cg_outputs[0].allclose(tref, 0.01, 0.01),
Expand Down Expand Up @@ -2602,6 +2640,8 @@ TEST_F(MatmulTest, AmpereMatmulSmemEpiloguePromotionRequiredA100) {

// check bank conflicts
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
// (0.001, 0.001) passed on local A100 but failed on CI A100
NVF_CHECK(
cg_outputs[0].allclose(tref, 0.01, 0.01),
Expand Down Expand Up @@ -2700,6 +2740,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulSmemEpilogueCast) {
tref = tref.to(at::kHalf);
// check bank conflicts
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
// (0.001, 0.001) passed on local A100 but failed on CI A100
NVF_CHECK(
cg_outputs[0].allclose(tref, 0.01, 0.01),
Expand Down Expand Up @@ -2795,6 +2837,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulSmemEpilogueRelu) {

// check bank conflicts
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
// (0.001, 0.001) passed on local A100 but failed on CI A100
NVF_CHECK(
cg_outputs[0].allclose(tref, 0.01, 0.01),
Expand Down Expand Up @@ -2874,6 +2918,8 @@ TEST_P(MatmulTestWithLayout, FusionAmpereMatmulSplitK_CUDA) {
7, 5, fe.compileFusion(&fusion, {inputs.first, inputs.second}));
EXPECT_TRUE(getBankConflictInfo(fe.kernel()).empty());
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);

Expand Down Expand Up @@ -2940,6 +2986,8 @@ TEST_P(MatmulTestWithLayout, FusionAmpereMatmulSplitKBias_CUDA) {
7, 5, fe.compileFusion(&fusion, inputs));
EXPECT_TRUE(getBankConflictInfo(fe.kernel()).empty());
auto cg_outputs = fe.runFusion(inputs);
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto tref = atBiasEpilogue(
atMatmul(aten_a.to(at::kFloat), aten_b.to(at::kFloat), layout),
aten_bias);
Expand Down Expand Up @@ -3003,6 +3051,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulBatchSplitK) {
NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
7, 5, fe.compileFusion(&fusion, inputs));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion(inputs);
auto tref =
atMatmul(aten_a.to(at::kFloat), aten_b.to(at::kFloat), layout);
Expand Down Expand Up @@ -3071,6 +3121,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulBatchSplitKBias) {
NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
7, 5, fe.compileFusion(&fusion, inputs));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion(inputs);
auto tref = atBiasEpilogue(
atMatmul(aten_a.to(at::kFloat), aten_b.to(at::kFloat), layout),
Expand Down Expand Up @@ -3134,6 +3186,8 @@ TEST_F(MatmulTest, ReproIssue1808) {
LaunchParams(),
matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
auto tref = atMatmul(
inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
Expand Down Expand Up @@ -3284,6 +3338,8 @@ TEST_P(MatmulTestWithLayout, MisalignedVectorization) {
fe.compileFusion(
fusion.get(), inputs, LaunchParams(), matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto outputs = fe.runFusion(inputs);

EXPECT_TRUE(outputs[0].allclose(tref, 0.001, 0.001));
Expand Down Expand Up @@ -3339,6 +3395,8 @@ TEST_F(MatmulTest, MultipleConsecutiveDims) {
NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
8, 0, fe.compileFusion(&fusion, inputs, LaunchParams(), matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion(inputs);
auto tref = at::reshape(
at::linear(
Expand Down Expand Up @@ -3403,6 +3461,8 @@ TEST_F(MatmulTest, DISABLED_MultipleNonConsecutiveMDims) {
NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
8, 0, fe.compileFusion(&fusion, inputs, LaunchParams(), matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion(inputs);
auto Apermuted = A.permute({{1, 2}}).reshape({M1 * M2, K});
auto tref = at::linear(Apermuted.to(at::kFloat), B.to(at::kFloat))
Expand Down Expand Up @@ -3467,6 +3527,8 @@ TEST_F(MatmulTest, DISABLED_MultipleNonConsecutiveNDims) {
NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
8, 0, fe.compileFusion(&fusion, inputs, LaunchParams(), matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion(inputs);
auto Bpermuted = B.permute({{1, 2}}).reshape({N1 * N2, K});
auto tref = at::linear(A.to(at::kFloat), Bpermuted.to(at::kFloat))
Expand Down Expand Up @@ -3523,6 +3585,8 @@ TEST_F(MatmulTest, MultipleMDimsBatch) {
NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
8, 0, fe.compileFusion(&fusion, inputs, LaunchParams(), matmul_cparams));
ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
ASSERT_FALSE(
PredicatedChecker::isCpAsyncMmaPredicatedByIfThenElse(fe.kernel()));
auto cg_outputs = fe.runFusion(inputs);
auto tref =
at::matmul(A.to(at::kFloat), at::permute(B.to(at::kFloat), {0, 2, 1}));
Expand Down
27 changes: 27 additions & 0 deletions tests/cpp/utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,33 @@ class PredicatedChecker : public kir::IrVisitor {
return isPredicated(tv->name(), kernel);
}

static bool isPredicatedByIfThenElse(
StmtNameType tv_name,
kir::Kernel* kernel) {
PredicatedChecker checker(tv_name, kernel->topLevelExprs());
return checker.predicated_ite_;
}

// If CpAsync from gmem to smem, then loaded from smem to registers using
// ldmatrix, then it is used in mma and should not use if-then-else predicate.
// If just CpAsync from gmem to smem, without further copy to register, then
// it is not used in mma and can use if-then-else predicate.
static bool isCpAsyncMmaPredicatedByIfThenElse(kir::Kernel* kernel) {
for (auto tv : kernel->allTvs()) {
if (tv->definition() != nullptr &&
ir_utils::isCpAsyncOp(tv->definition())) {
const auto& consumers = ir_utils::consumerTvsOf(tv);
if (std::any_of(
consumers.begin(), consumers.end(), [&](TensorView* tv) {
return ir_utils::isLdMatrixOp(tv->definition());
})) {
return isPredicatedByIfThenElse(tv->name(), kernel);
}
}
}
return false;
}

private:
PredicatedChecker() = delete;

Expand Down

0 comments on commit 61a77e0

Please sign in to comment.