[hal] Add ObjectLike trait to hal.executable.source

compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td

@@ -2616,6 +2616,7 @@ def HAL_ExecutableSourceOp : HAL_Op<"executable.source", [
   SingleBlockImplicitTerminator<"IREE::HAL::ExecutableSourceEndOp">,
   Symbol,
   SymbolTable,
+  Util_ObjectLike,
 ]> {
   let summary = [{Generic source contents of an executable op.}];
   let description = [{

tools/test/BUILD.bazel

@@ -41,6 +41,7 @@ iree_lit_test_suite(
             "iree-dump-module.mlir",
             "iree-dump-parameters.txt",
             "iree-link-bundle.mlir",
+            "iree-link-executable-source.mlir",
             "iree-link.mlir",
             "iree-run-mlir.mlir",
             "iree-run-module-expected.mlir",
@@ -71,6 +72,7 @@ iree_lit_test_suite(
             "iree-link-globals-a.mlir",
             "iree-link-globals-b.mlir",
             "iree-link-globals-c.mlir",
+            "iree-link-executable-source-module-d.mlir",
             "iree-link-module-a.mlir",
             "iree-link-module-b.mlir",
             "iree-link-module-c.mlir",
@@ -82,6 +84,7 @@ iree_lit_test_suite(
         "iree-link-anonymous.mlir",
         "iree-tokenize.json",
         "iree-tokenize.tiktoken",
+        "iree-link-executable-source-module-d.mlir",
         "iree-link-globals-a.mlir",
         "iree-link-globals-b.mlir",
         "iree-link-globals-c.mlir",

tools/test/CMakeLists.txt

@@ -29,6 +29,7 @@ iree_lit_test_suite(
     "iree-dump-module.mlir"
     "iree-dump-parameters.txt"
     "iree-link-bundle.mlir"
+    "iree-link-executable-source.mlir"
     "iree-link.mlir"
     "iree-run-mlir.mlir"
     "iree-run-module-expected.mlir"
@@ -68,6 +69,7 @@ iree_lit_test_suite(
     MLIROpBaseTdFiles
     echo_npy.py
     iree-link-anonymous.mlir
+    iree-link-executable-source-module-d.mlir
     iree-link-globals-a.mlir
     iree-link-globals-b.mlir
     iree-link-globals-c.mlir

tools/test/iree-link-executable-source-module-d.mlir

@@ -0,0 +1,27 @@
+// Source module with hal.executable.source (extern kernel dispatch).
+// Used by iree-link-executable-source.mlir.
+
+module @module_d {
+  // Extern kernel executable (e.g., a pre-compiled GPU kernel object).
+  hal.executable.source private @extern_kernel attributes {
+    objects = #hal.executable.objects<{}>
+  } {
+    hal.executable.export public @entry ordinal(0)
+        layout(#hal.pipeline.layout<constants = 2, bindings = [
+          #hal.pipeline.binding<storage_buffer, ReadOnly>,
+          #hal.pipeline.binding<storage_buffer>
+        ]>) count(%device: !hal.device) -> (index, index, index) {
+      %c1 = arith.constant 1 : index
+      hal.return %c1, %c1, %c1 : index, index, index
+    } attributes {workgroup_size = [64 : index, 1 : index, 1 : index]}
+  }
+
+  // Function that dispatches to the extern kernel.
+  util.func public @transform(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+    %c4 = arith.constant 4 : i32
+    %c1 = arith.constant 1 : index
+    %0 = flow.dispatch @extern_kernel::@entry(%c4, %arg0)
+        : (i32, tensor<4xf32>) -> tensor<4xf32>
+    util.return %0 : tensor<4xf32>
+  }
+}

tools/test/iree-link-executable-source.mlir

@@ -0,0 +1,46 @@
+// Tests linking modules containing hal.executable.source ops.
+// hal.executable.source is used for hand-authored extern dispatch kernels
+// (e.g., pre-compiled GPU kernel objects). The linker must copy these
+// alongside the functions that reference them via flow.dispatch.
+
+// RUN: iree-link \
+// RUN:     --link-module=%p/iree-link-executable-source-module-d.mlir \
+// RUN:     %s | FileCheck %s
+
+// External reference to module D's transform function.
+util.func private @module_d.transform(%arg0: tensor<4xf32>) -> tensor<4xf32>
+
+// Main entry point that calls the linked function.
+util.func public @main(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = util.call @module_d.transform(%arg0) : (tensor<4xf32>) -> tensor<4xf32>
+  util.return %0 : tensor<4xf32>
+}
+
+// After linking:
+// 1. The function body is filled in (not just a declaration).
+// CHECK-DAG: util.func public @main
+// CHECK-DAG: util.func private @module_d.transform
+
+// 2. The flow.dispatch reference in the linked function body is intact.
+// CHECK-DAG: flow.dispatch @extern_kernel::@entry
+
+// 3. The hal.executable.source is copied with its body intact.
+// CHECK-DAG: hal.executable.source private @extern_kernel
+// CHECK-DAG: hal.executable.export public @entry ordinal(0)
+// CHECK-DAG: workgroup_size = [64 : index, 1 : index, 1 : index]
+
+// -----
+
+// Test that hal.executable.source internal symbols are NOT treated as
+// external dependencies. The nested hal.executable.export should not
+// appear as a linkable symbol.
+
+// RUN: iree-link --list-symbols \
+// RUN:     %p/iree-link-executable-source-module-d.mlir \
+// RUN:     | FileCheck %s --check-prefix=SYMBOLS
+
+// Only the public function should be listed, not the executable or its exports.
+// SYMBOLS: Public symbols in
+// SYMBOLS: @transform
+// SYMBOLS-NOT: @extern_kernel
+// SYMBOLS-NOT: @entry