From 3e983a1439c8efd9ff863f1576a65a8e3b9a32ae Mon Sep 17 00:00:00 2001 From: 6clc Date: Mon, 27 Mar 2023 16:21:26 +0800 Subject: [PATCH 1/2] tests(auto-unroll): Check whether the Auto-unroll is correct --- .../search_space/auto_gen_rule/CMakeLists.txt | 2 +- .../auto_gen_rule/auto_unroll_test.cc | 77 ++++++++++++++++++- 2 files changed, 76 insertions(+), 3 deletions(-) diff --git a/cinn/auto_schedule/search_space/auto_gen_rule/CMakeLists.txt b/cinn/auto_schedule/search_space/auto_gen_rule/CMakeLists.txt index 5105ed66dc..510d445adf 100644 --- a/cinn/auto_schedule/search_space/auto_gen_rule/CMakeLists.txt +++ b/cinn/auto_schedule/search_space/auto_gen_rule/CMakeLists.txt @@ -21,4 +21,4 @@ endif() cc_test(test_auto_inline SRCS auto_inline_test.cc DEPS cinncore auto_gen_rule_test_helper) cc_test(test_skip_rule SRCS skip_rule_test.cc DEPS cinncore) -cc_test(test_auto_unroll SRCS auto_unroll_test.cc DEPS cinncore) +cc_test(test_auto_unroll SRCS auto_unroll_test.cc DEPS cinncore auto_gen_rule_test_helper test_program_builder) diff --git a/cinn/auto_schedule/search_space/auto_gen_rule/auto_unroll_test.cc b/cinn/auto_schedule/search_space/auto_gen_rule/auto_unroll_test.cc index 99688a2da6..3531cf9a96 100644 --- a/cinn/auto_schedule/search_space/auto_gen_rule/auto_unroll_test.cc +++ b/cinn/auto_schedule/search_space/auto_gen_rule/auto_unroll_test.cc @@ -17,8 +17,11 @@ #include #include +#include "cinn/auto_schedule/search_space/auto_gen_rule/test_helper.h" #include "cinn/cinn.h" #include "cinn/lang/lower.h" +#include "tests/program_builder.h" +#include "tests/subgraph_program_builder.h" namespace cinn { namespace auto_schedule { @@ -73,7 +76,7 @@ TEST(AutoUnroll, UnrollableApply) { auto* init_schedule_block = init_block_realize->schedule_block.As(); ASSERT_NE(init_schedule_block, nullptr); ASSERT_TRUE(init_schedule_block->attrs.empty()); - VLOG(6) << "Before auto-unroll:\n" << ast_expr; + VLOG(-6) << "Before auto-unroll:\n" << ast_expr; AutoUnroll test_rule(target); ir::IRSchedule ir_schedule(ir::ModuleExpr({ast_expr})); @@ -96,12 +99,82 @@ TEST(AutoUnroll, UnrollableApply) { const int* max_step = absl::get_if(&attr_value); EXPECT_NE(max_step, nullptr); EXPECT_LE(*max_step, 128); - VLOG(6) << "After auto-unroll:max_step=" << *max_step << ", Ast:\n" << ir_sch->GetModule().GetExprs().front(); + VLOG(-6) << "After auto-unroll:max_step=" << *max_step << ", Ast:\n" << ir_sch->GetModule().GetExprs().front(); }; test_func(&ir_schedule); test_func(&states[0]->ir_schedule); } +#ifdef CINN_WITH_CUDA +class TestAutoUnroll : public TestAutoGenRuleBase { + public: + std::vector default_input_names = {"X", "Y"}; + std::vector default_output_names = {"temp_matmul_out"}; +}; +TEST_F(TestAutoUnroll, ApplyOnMatmulWithTiling) { + frontend::Program matmul_op = tests::OpBuilder("matmul").Build({{"X", {32, 4}}, {"Y", {4, 32}}}); + Initialize(common::DefaultNVGPUTarget()); + ir::IRSchedule ir_schedule = MakeIRSchedule(matmul_op); + std::vector func_bodys = ir_schedule.GetModule().GetExprs(); + ASSERT_EQ(func_bodys.size(), 1UL); + VLOG(6) << "Original Expr:\n" << func_bodys[0]; + + AutoUnroll auto_unroll(target_); + SearchState state(ir_schedule, 0, {}); + const std::string& applied_block_name = default_output_names.back(); + EXPECT_EQ(auto_unroll.AnalyseApplyType(state, applied_block_name), RuleApplyType::kApplyAndPruneOtherRules); + auto new_states = auto_unroll.ApplyOnBlock(state, applied_block_name); + std::vector exprs = new_states[0]->ir_schedule.GetModule().GetExprs(); + EXPECT_EQ(exprs.size(), 1UL); + + // Check if the block has an 'auto_unroll_max_step' attribute + auto* applied_block_realize = exprs.front().As()->stmts.front().As(); + auto* applied_schedule_block = applied_block_realize->schedule_block.As(); + ASSERT_FALSE(applied_schedule_block->attrs.empty()); + EXPECT_EQ(applied_schedule_block->attrs.count(ir::attr::auto_unroll_max_step), 1); + const auto& attr_value = applied_schedule_block->attrs.at(ir::attr::auto_unroll_max_step); + const int* max_step = absl::get_if(&attr_value); + EXPECT_NE(max_step, nullptr); + EXPECT_LE(*max_step, 128); + VLOG(6) << "Expr after AutoUnroll applied on block:max_step=" << *max_step << ", Ast:\n" << exprs.front(); + + // build ir::Module and debug source code + auto build_module = BuildIRModule(new_states[0]->ir_schedule); + auto source_code = GenSourceCode(build_module); + VLOG(6) << " auto-schedule source code:\n" << source_code; + // execute and check precision + CheckResult(GenExecutableKernel(build_module), + GenExecutableKernel(BuildIRModule(MakeIRSchedule(matmul_op, /* apply_manual_schedule */ true))), + default_input_names, + default_output_names, + {{4, 4}, {4, 4}}, + {{4, 4}}, + target_); +} + +TEST_F(TestAutoUnroll, PureSpatial) { + Target target = common::DefaultNVGPUTarget(); + Initialize(target); + std::vector input_names = {"x", "y"}; + std::vector output_names = { + "var_6", "var_4", "constant_idx_last", "constant_idx_first", "var_2", "var_5"}; + std::vector input_shape{256, 256}; + std::vector inputs_varinfo({{"x", input_shape}, {"y", input_shape}}); + + Context::Global().ResetNameId(); + ir::IRSchedule ir_schedule = MakeIRSchedule(tests::GatherAddSubSubGraphBuilder().Build(inputs_varinfo)); + SearchState state(ir_schedule, 0, {}); + std::vector func_bodys = ir_schedule.GetModule().GetExprs(); + ASSERT_EQ(func_bodys.size(), 1UL); + VLOG(6) << "Original Expr:\n" << func_bodys[0]; + + AutoUnroll auto_unroll(target_); + for (const auto& applied_block_name : output_names) { + EXPECT_EQ(auto_unroll.AnalyseApplyType(state, applied_block_name), RuleApplyType::kCannotApply); + } +} +#endif + } // namespace auto_schedule } // namespace cinn From 0240813050147299a3dcab6d06f67fde4b85a7db Mon Sep 17 00:00:00 2001 From: 6clc Date: Wed, 29 Mar 2023 20:26:37 +0800 Subject: [PATCH 2/2] tests(auto-unroll): fix code according to merge --- .../auto_gen_rule/auto_unroll_test.cc | 33 +++++++++++++++---- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/cinn/auto_schedule/search_space/auto_gen_rule/auto_unroll_test.cc b/cinn/auto_schedule/search_space/auto_gen_rule/auto_unroll_test.cc index 3531cf9a96..a671c8e0d1 100644 --- a/cinn/auto_schedule/search_space/auto_gen_rule/auto_unroll_test.cc +++ b/cinn/auto_schedule/search_space/auto_gen_rule/auto_unroll_test.cc @@ -20,8 +20,8 @@ #include "cinn/auto_schedule/search_space/auto_gen_rule/test_helper.h" #include "cinn/cinn.h" #include "cinn/lang/lower.h" +#include "tests/concrete_program_builder.h" #include "tests/program_builder.h" -#include "tests/subgraph_program_builder.h" namespace cinn { namespace auto_schedule { @@ -76,7 +76,7 @@ TEST(AutoUnroll, UnrollableApply) { auto* init_schedule_block = init_block_realize->schedule_block.As(); ASSERT_NE(init_schedule_block, nullptr); ASSERT_TRUE(init_schedule_block->attrs.empty()); - VLOG(-6) << "Before auto-unroll:\n" << ast_expr; + VLOG(6) << "Before auto-unroll:\n" << ast_expr; AutoUnroll test_rule(target); ir::IRSchedule ir_schedule(ir::ModuleExpr({ast_expr})); @@ -99,7 +99,7 @@ TEST(AutoUnroll, UnrollableApply) { const int* max_step = absl::get_if(&attr_value); EXPECT_NE(max_step, nullptr); EXPECT_LE(*max_step, 128); - VLOG(-6) << "After auto-unroll:max_step=" << *max_step << ", Ast:\n" << ir_sch->GetModule().GetExprs().front(); + VLOG(6) << "After auto-unroll:max_step=" << *max_step << ", Ast:\n" << ir_sch->GetModule().GetExprs().front(); }; test_func(&ir_schedule); @@ -112,14 +112,30 @@ class TestAutoUnroll : public TestAutoGenRuleBase { std::vector default_input_names = {"X", "Y"}; std::vector default_output_names = {"temp_matmul_out"}; }; -TEST_F(TestAutoUnroll, ApplyOnMatmulWithTiling) { - frontend::Program matmul_op = tests::OpBuilder("matmul").Build({{"X", {32, 4}}, {"Y", {4, 32}}}); + +/* Before AutoUnroll: + * for (i=0; i < 4; i++): + * for(j=0; j < 4; j++): + * for(k=0; k < 4; k++): + * C(i, j) = C(i, j) + A(i, k) * B(k, j) + * + * After AutoUnroll on 'k', the third loop is unrolled. + * for(i=0; i < 4; i++): + * for(j=0; j < 4; j++): + * C(i, j) = C(i, j) + A(i, 0) * B(0, j) + * C(i, j) = C(i, j) + A(i, 1) * B(1, j) + * C(i, j) = C(i, j) + A(i, 2) * B(2, j) + * C(i, j) = C(i, j) + A(i, 3) * B(3, j) + */ +TEST_F(TestAutoUnroll, ApplyOnMatmulWithUnroll) { + frontend::Program matmul_op = tests::OpBuilder("matmul").Build({{"X", {4, 4}}, {"Y", {4, 4}}}); Initialize(common::DefaultNVGPUTarget()); ir::IRSchedule ir_schedule = MakeIRSchedule(matmul_op); std::vector func_bodys = ir_schedule.GetModule().GetExprs(); ASSERT_EQ(func_bodys.size(), 1UL); VLOG(6) << "Original Expr:\n" << func_bodys[0]; + // Construct the computation graph and convert it to ir::Expr AutoUnroll auto_unroll(target_); SearchState state(ir_schedule, 0, {}); const std::string& applied_block_name = default_output_names.back(); @@ -145,7 +161,7 @@ TEST_F(TestAutoUnroll, ApplyOnMatmulWithTiling) { VLOG(6) << " auto-schedule source code:\n" << source_code; // execute and check precision CheckResult(GenExecutableKernel(build_module), - GenExecutableKernel(BuildIRModule(MakeIRSchedule(matmul_op, /* apply_manual_schedule */ true))), + GenExecutableKernel(BuildIRModule(MakeIRSchedule(matmul_op, -1, true))), default_input_names, default_output_names, {{4, 4}, {4, 4}}, @@ -153,6 +169,7 @@ TEST_F(TestAutoUnroll, ApplyOnMatmulWithTiling) { target_); } +/* Operators of type elementwise or injective can not be auto-unrolled.*/ TEST_F(TestAutoUnroll, PureSpatial) { Target target = common::DefaultNVGPUTarget(); Initialize(target); @@ -162,13 +179,15 @@ TEST_F(TestAutoUnroll, PureSpatial) { std::vector input_shape{256, 256}; std::vector inputs_varinfo({{"x", input_shape}, {"y", input_shape}}); + // Construct the computation graph and convert it to ir::Expr Context::Global().ResetNameId(); - ir::IRSchedule ir_schedule = MakeIRSchedule(tests::GatherAddSubSubGraphBuilder().Build(inputs_varinfo)); + ir::IRSchedule ir_schedule = MakeIRSchedule(tests::GatherAddSubBuilder().Build(inputs_varinfo)); SearchState state(ir_schedule, 0, {}); std::vector func_bodys = ir_schedule.GetModule().GetExprs(); ASSERT_EQ(func_bodys.size(), 1UL); VLOG(6) << "Original Expr:\n" << func_bodys[0]; + // Analyzes whether the block can be unrolled AutoUnroll auto_unroll(target_); for (const auto& applied_block_name : output_names) { EXPECT_EQ(auto_unroll.AnalyseApplyType(state, applied_block_name), RuleApplyType::kCannotApply);