[Codegen][Tuner] skip linking based on the default entry point attribute (iree-org#19603)

This PR generalizes the cases in which the linking pass can be skipped
based on the presence of the default entry point attribute.

---------

Signed-off-by: Bangtian Liu <[email protected]>

Author: bangtianliu

compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp

@@ -197,14 +197,30 @@ struct MaterializeTuningSpecsPass final
       return;
     }
 
-    // If only the default tuning spec is available, use it directly and skip
-    // the linking stage.
-    if (!hasUserTuningSpec) {
-      if (failed(dumpFinalTuningSpecToDir(*defaultTuningSpec))) {
+    // When the user tuning spec and default spec are available, link all
+    // available libraries into a single module. We insert the default tuning
+    // spec last, so that any user-specified tuning configurations take
+    // precedence.
+    SmallVector<ModuleOp, 2> allSpecs;
+    if (hasUserTuningSpec) {
+      allSpecs.push_back(*userTuningSpec);
+    }
+    if (hasDefaultTuningSpec) {
+      allSpecs.push_back(*defaultTuningSpec);
+    }
+
+    // Determine if the linking pass should be skipped.
+    // Skip if there is only one tuning spec (either user-provided or default)
+    // with the default attribute.
+    if (allSpecs.size() == 1 &&
+        allSpecs[0]->hasAttr(kTuningSpecDefaultEntrypointAttrName)) {
+      // Use the appropriate tuning spec (user or default).
+      ModuleOp tuningSpecWithDefaultAttr = allSpecs[0];
+      if (failed(dumpFinalTuningSpecToDir(tuningSpecWithDefaultAttr))) {
         return signalPassFailure();
       }
       FailureOr<DenseElementsAttr> serializedSpec =
-          serializeTuningSpecToAttr(*defaultTuningSpec);
+          serializeTuningSpecToAttr(tuningSpecWithDefaultAttr);
       if (failed(serializedSpec)) {
         module->emitError("Failed to serialize default tuning specs");
         return signalPassFailure();
@@ -213,14 +229,6 @@ struct MaterializeTuningSpecsPass final
       return;
     }
 
-    // When the user tuning spec is available, link all available libraries into
-    // a single module. We insert the default tuning spec last, so that any
-    // user-specified tuning configurations take precedence.
-    SmallVector<ModuleOp, 2> allSpecs = {*userTuningSpec};
-    if (hasDefaultTuningSpec) {
-      allSpecs.push_back(*defaultTuningSpec);
-    }
-
     Location loc =
         FusedLoc::get(ctx, llvm::map_to_vector<2>(allSpecs, [](ModuleOp m) {
                         return m.getLoc();

compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel

@@ -107,6 +107,7 @@ iree_lit_test_suite(
             "reductions_codegen_spec.mlir",
             "reductions_match_spec.mlir",
             "tuning_spec.mlir",
+            "tuning_spec_default.mlir",
         ],
     ),
     cfg = "//compiler:lit.cfg.py",
@@ -118,6 +119,7 @@ iree_lit_test_suite(
         "reductions_codegen_spec.mlir",
         "reductions_match_spec.mlir",
         "tuning_spec.mlir",
+        "tuning_spec_default.mlir",
     ],
     tools = [
         "//tools:iree-opt",

compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt

@@ -104,6 +104,7 @@ iree_lit_test_suite(
     reductions_codegen_spec.mlir
     reductions_match_spec.mlir
     tuning_spec.mlir
+    tuning_spec_default.mlir
 )
 
 ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###

compiler/src/iree/compiler/Codegen/Common/test/materialize_tuning_specs.mlir

@@ -3,6 +3,11 @@
 // RUN:   --iree-codegen-dump-tuning-specs-to=- \
 // RUN:   --mlir-disable-threading --no-implicit-module %s | FileCheck %s
 
+// RUN: iree-opt --pass-pipeline='builtin.module(iree-codegen-materialize-tuning-specs)' \
+// RUN:   --iree-codegen-tuning-spec-path=%p/tuning_spec_default.mlir \
+// RUN:   --iree-codegen-dump-tuning-specs-to=- \
+// RUN:   --mlir-disable-threading --no-implicit-module %s | FileCheck %s --check-prefix=SKIPLINK
+
 // Check that the final tuning spec is as expected when the user tuning spec is provided.
 
 // CHECK-LABEL: module @iree_linked_tuning_spec
@@ -19,6 +24,17 @@
 // CHECK-SAME:     iree_codegen.tuning_spec_mlirbc = dense<{{.+}}> : vector<{{[0-9]+}}xi8>
 // CHECK-LABEL:    func.func @main_0
 
+
+// CHECK that the user-provided tuning spec is materized without linking when default tuing spec
+// is missing and the user-provided tuning spec is marked the default attribute.
+
+// SKIPLINK-LABEL: module  @user_spec
+// SKIPLINK-SAME:    iree_codegen.tuning_spec_with_default_entrypoint
+// SKIPLINK-SAME:    transform.with_named_sequence
+// SKIPLINK-NOT:    module @{{.+}}
+// SKIPLINK:        module attributes
+// SKIPLINK-SAME:     iree_codegen.tuning_spec_mlirbc = dense<{{.+}}> : vector<{{[0-9]+}}xi8>
+// SKIPLINK-LABEL:    func.func @main_0
 module {
   func.func @main_0() {
     return

compiler/src/iree/compiler/Codegen/Common/test/tuning_spec_default.mlir

@@ -0,0 +1,9 @@
+// RUN: iree-opt %s
+
+module @user_spec attributes { transform.with_named_sequence, iree_codegen.tuning_spec_with_default_entrypoint } {
+  transform.named_sequence @__kernel_config(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op
+    attributes { iree_codegen.tuning_spec_entrypoint } {
+    transform.print {name = "Hello Tuning Spec", skip_regions}
+    transform.yield %arg0 : !transform.any_op
+  }
+}