[transform] add a default schedule (#899)

alibaba · Dec 23, 2022 · 59359ff · 59359ff
1 parent 490b7a1
commit 59359ff
Show file tree

Hide file tree

Showing 19 changed files with 802 additions and 107 deletions.
diff --git a/scripts/python/tao_build.py b/scripts/python/tao_build.py
@@ -391,7 +391,7 @@ def bazel_test(target, flag=""):
         logger.info("Testing bazel target: " + target)
         flag += " --experimental_ui_max_stdouterr_bytes=-1 "
         execute(" ".join([BAZEL_BUILD_CMD, flag, target]))
-        execute(" ".join([BAZEL_TEST_CMD, flag + ' --test_env=TF_CPP_VMODULE=disc_compiler=1 --test_env=TF_ENABLE_ONEDNN_OPTS=0' , target]))
+        execute(" ".join([BAZEL_TEST_CMD, flag + ' --test_env=TF_CPP_VMODULE=disc_compiler=1,disc_transform_legalize_to_loop=1 --test_env=TF_ENABLE_ONEDNN_OPTS=0' , target]))
 
     with cwd(tf_root_dir(root)), gcc_env(args.compiler_gcc):
         execute(

diff --git a/tao_compiler/mlir/disc/BUILD b/tao_compiler/mlir/disc/BUILD
@@ -2039,8 +2039,9 @@ disc_cc_library(
 )
 
 cc_library(
-    name = "disc_transform_legalize_to_loop",
-    srcs = ["transforms/disc_transform_legalize_to_loop.cc"],
+    name = "disc_transform_schedule",
+    srcs = ["transforms/disc_transform_schedule.cc"],
+    hdrs = ["transforms/disc_transform_schedule.h"],
     deps = [
         ":codegen_utils",
         ":disc_lhlo_elemental_utils",
@@ -2053,6 +2054,7 @@ cc_library(
         "//tensorflow/compiler/xla/mlir_hlo:lhlo",
         "//tensorflow/compiler/mlir/disc/tools/disc-transform:all_passes",
         "//tensorflow/compiler/mlir/disc/tools/disc-transform:DISCLinalgExtDialect",
+        "//tensorflow/tsl/platform:logging",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:FuncDialect",
@@ -2069,6 +2071,15 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "disc_transform_legalize_to_loop",
+    srcs = ["transforms/disc_transform_legalize_to_loop.cc"],
+    deps = [
+        ":disc_transform_schedule",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "all_passes",
     hdrs = [

diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/packed_matmul_nn_p_f32_large_schedule.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/packed_matmul_nn_p_f32_large_schedule.mlir
@@ -54,7 +54,7 @@ transform.structured.canonicalized_sequence failures(propagate) {
     multireduction_lowering = "innerparallel",
     split_transfers = "linalg-copy",
     // stages = [0, 1, 2, 3, 4, 5, 6, 7],
-    stages = [0, 1, 2, 3],
+    stages = [0, 1, 2, 3, 4],
     transpose_avx2_lowering = false,
     transpose_lowering = "eltwise",
     unroll_vector_transfers = true

diff --git a/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
@@ -28,9 +28,10 @@ static bool init_threads = []() {
 }();
 
 TEST(SimpleTest, MatMulF32_11x13x12) {
-  EnvSetting setting = {{"DISC_TRANSFORM_SCHEDULE_FILE",
-                         {c_ft_path + "matmul_nn_d_f32_schedule.mlir", false}},
-                        {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_schedule.mlir", false}},
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
   EnvSettingContext ctx(setting);
   EXPECT_TRUE(feature_test_main(
       /*mlir_file_path*/ c_ft_path + "matmul_nn_d_f32.mlir",
@@ -42,9 +43,10 @@ TEST(SimpleTest, MatMulF32_11x13x12) {
 }
 
 TEST(SimpleTest, MatMulF32_111x131x121) {
-  EnvSetting setting = {{"DISC_TRANSFORM_SCHEDULE_FILE",
-                         {c_ft_path + "matmul_nn_d_f32_schedule.mlir", false}},
-                        {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_schedule.mlir", false}},
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
   EnvSettingContext ctx(setting);
   EXPECT_TRUE(feature_test_main(
       /*mlir_file_path*/ c_ft_path + "matmul_nn_d_f32.mlir",
@@ -58,7 +60,7 @@ TEST(SimpleTest, MatMulF32_111x131x121) {
 TEST(SimpleTest, MatMulF32_304x1024x256) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -78,7 +80,7 @@ TEST(SimpleTest, MatMulF32_304x1024x256) {
 TEST(SimpleTest, MatMulF32_1024x1024x1024) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule.mlir", false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -98,7 +100,8 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024) {
 TEST(SimpleTest, MatMulF32_304x1024x256_2) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_2.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_2.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -118,7 +121,8 @@ TEST(SimpleTest, MatMulF32_304x1024x256_2) {
 TEST(SimpleTest, MatMulF32_1024x1024x1024_2) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_2.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_2.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -138,7 +142,8 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024_2) {
 TEST(SimpleTest, MatMulF32_304x256x256_3) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -158,7 +163,8 @@ TEST(SimpleTest, MatMulF32_304x256x256_3) {
 TEST(SimpleTest, MatMulF32_304x512x256_3) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -178,7 +184,8 @@ TEST(SimpleTest, MatMulF32_304x512x256_3) {
 TEST(SimpleTest, MatMulF32_304x1024x256_3) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -198,7 +205,8 @@ TEST(SimpleTest, MatMulF32_304x1024x256_3) {
 TEST(SimpleTest, MatMulF32_304x1024x512_3) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -218,7 +226,8 @@ TEST(SimpleTest, MatMulF32_304x1024x512_3) {
 TEST(SimpleTest, MatMulF32_1024x1024x1024_3) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -238,7 +247,8 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024_3) {
 TEST(SimpleTest, MatMulF32_304x1024x512_4) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_4.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_4.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -258,7 +268,8 @@ TEST(SimpleTest, MatMulF32_304x1024x512_4) {
 TEST(SimpleTest, MatMulF32_1024x1024x1024_4) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_4.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_4.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -278,7 +289,8 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024_4) {
 TEST(SimpleTest, MatMulF32_1026x1024x1024_4) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_4.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_4.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -298,7 +310,8 @@ TEST(SimpleTest, MatMulF32_1026x1024x1024_4) {
 TEST(SimpleTest, MatMulF32_304x1024x512_5) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_nn_d_f32_large_schedule_5.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_nn_d_f32_large_schedule_5.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};

diff --git a/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc b/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc
@@ -30,7 +30,8 @@ static bool init_threads = []() {
 TEST(SimpleMTTest, MatMulF32_111x131x121_Thread_8) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_multithread_nn_d_f32_schedule.mlir", false}},
+       {"kGEMM::" + c_ft_path + "matmul_multithread_nn_d_f32_schedule.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
   EnvSettingContext ctx(setting);
   EXPECT_TRUE(feature_test_main(
@@ -43,10 +44,11 @@ TEST(SimpleMTTest, MatMulF32_111x131x121_Thread_8) {
 }
 
 TEST(SimpleTest, MatMulF32_304x1024x256) {
-  EnvSetting setting = {
-      {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "matmul_multithread_nn_d_f32_large_schedule.mlir", false}},
-      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
+  EnvSetting setting = {{"DISC_TRANSFORM_SCHEDULE_FILE",
+                         {"kGEMM::" + c_ft_path +
+                              "matmul_multithread_nn_d_f32_large_schedule.mlir",
+                          false}},
+                        {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}}};
   EnvSettingContext ctx(setting);
   EXPECT_TRUE(feature_test_main(
       /*mlir_file_path*/ c_ft_path + "matmul_multithread_nn_d_f32.mlir",

diff --git a/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc
@@ -30,7 +30,8 @@ static bool init_threads = []() {
 TEST(PackedMatmul, F32_304x1024x512) {
   EnvSetting setting = {
       {"DISC_TRANSFORM_SCHEDULE_FILE",
-       {c_ft_path + "packed_matmul_nn_p_f32_large_schedule.mlir", false}},
+       {"kGEMM::" + c_ft_path + "packed_matmul_nn_p_f32_large_schedule.mlir",
+        false}},
       {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
       {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
       {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
@@ -47,4 +48,21 @@ TEST(PackedMatmul, F32_304x1024x512) {
       /*profiling*/ true));
 }
 
+TEST(PackedMatmul, F32_304x1024x512_Using_Default_Schedule) {
+  EnvSetting setting = {{"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+                        {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
+                        {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", false}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "packed_matmul_nn_p_512x1024_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 1,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x512xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
 }  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/BUILD b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
@@ -319,6 +319,7 @@ cc_library(
         "@llvm-project//mlir:PDLInterpDialect",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:SCFTransformOps",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:TransformDialect",

diff --git a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -198,9 +198,11 @@ DiagnosedSilenceableFailure DISCBufferizeOp::apply(
       }));
   pm.addNestedPass<func::FuncOp>(bufferization::createBufferDeallocationPass());
 
-  if (failed(pm.run(state.getTopLevel())))
+  if (failed(pm.run(moduleOp)))
     return DiagnosedSilenceableFailure::definiteFailure();
 
+  results.set(getResult().cast<OpResult>(), {payload});
+
   return DiagnosedSilenceableFailure::success();
 }
 

diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/passes.h b/tao_compiler/mlir/disc/tools/disc-transform/transforms/passes.h
@@ -49,6 +49,10 @@ std::unique_ptr<OperationPass<ModuleOp>>
 createDiscTransformDialectInterpreterPass(const std::string& fileName = "",
                                           bool enableExpensiveChecks = false);
 
+// Erases transform dialect schedule from the IR
+std::unique_ptr<OperationPass<ModuleOp>>
+createDiscTransformDialectEraseSchedulePass();
+
 // Converts the transformed payload IR to be suitable for RAL.
 std::unique_ptr<OperationPass<ModuleOp>> createDiscRewritePayloadIRForRALPass(
     bool gpuEnabled = false);