[mlir][spirv] Ensure function declarations precede definitions (#164956)

SPIR-V spec requires that any calls to external functions are preceded
by declarations of those external functions. To simplify the
implementation, we sort functions in the serializer using a stronger
condition: any functions declarations are moved before any functions
definitions - this ensures that external functions are always declared
before being used.

Author: IgWod-IMG

mlir/lib/Target/SPIRV/Serialization/Serializer.cpp

@@ -89,6 +89,22 @@ static bool isZeroValue(Attribute attr) {
   return false;
 }
 
+/// Move all functions declaration before functions definitions. In SPIR-V
+/// "declarations" are functions without a body and "definitions" functions
+/// with a body. This is stronger than necessary. It should be sufficient to
+/// ensure any declarations precede their uses and not all definitions, however
+/// this allows to avoid analysing every function in the module this way.
+static void moveFuncDeclarationsToTop(spirv::ModuleOp moduleOp) {
+  Block::OpListType &ops = moduleOp.getBody()->getOperations();
+  if (ops.empty())
+    return;
+  Operation &firstOp = ops.front();
+  for (Operation &op : llvm::drop_begin(ops))
+    if (auto funcOp = dyn_cast<spirv::FuncOp>(op))
+      if (funcOp.getBody().empty())
+        funcOp->moveBefore(&firstOp);
+}
+
 namespace mlir {
 namespace spirv {
 
@@ -119,6 +135,8 @@ LogicalResult Serializer::serialize() {
   processMemoryModel();
   processDebugInfo();
 
+  moveFuncDeclarationsToTop(module);
+
   // Iterate over the module body to serialize it. Assumptions are that there is
   // only one basic block in the moduleOp
   for (auto &op : *module.getBody()) {

mlir/test/Target/SPIRV/function-decorations.mlir

@@ -1,14 +1,22 @@
 // RUN: mlir-translate --no-implicit-module --test-spirv-roundtrip --split-input-file %s | FileCheck %s
 
-spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+// RUN: %if spirv-tools %{ rm -rf %t %}
+// RUN: %if spirv-tools %{ mkdir %t %}
+// RUN: %if spirv-tools %{ mlir-translate --no-implicit-module --serialize-spirv --split-input-file --spirv-save-validation-files-with-prefix=%t/module %s %}
+// RUN: %if spirv-tools %{ spirv-val %t %}
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage, Int8, Int16], []> {
+    // CHECK: spirv.func @outside.func.with.linkage(i8) "Pure" attributes
+    // CHECK: linkage_attributes = #spirv.linkage_attributes<linkage_name = "outside.func", linkage_type = <Import>>
+    // CHECK: spirv.func @linkage_attr_test_kernel() "DontInline" {
+    // CHECK: spirv.func @inside.func() "Pure" {
     spirv.func @linkage_attr_test_kernel()  "DontInline"  attributes {}  {
         %uchar_0 = spirv.Constant 0 : i8
         %ushort_1 = spirv.Constant 1 : i16
         %uint_0 = spirv.Constant 0 : i32
         spirv.FunctionCall @outside.func.with.linkage(%uchar_0):(i8) -> ()
         spirv.Return
     }
-    // CHECK: linkage_attributes = #spirv.linkage_attributes<linkage_name = "outside.func", linkage_type = <Import>>
     spirv.func @outside.func.with.linkage(%arg0 : i8) -> () "Pure" attributes {
       linkage_attributes=#spirv.linkage_attributes<
         linkage_name="outside.func",
@@ -21,7 +29,7 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
 // -----
 
 spirv.module PhysicalStorageBuffer64 GLSL450 requires #spirv.vce<v1.0,
-    [Shader, PhysicalStorageBufferAddresses], [SPV_KHR_physical_storage_buffer]> {
+    [Shader, PhysicalStorageBufferAddresses, Linkage], [SPV_KHR_physical_storage_buffer]> {
   // CHECK-LABEL: spirv.func @func_arg_decoration_aliased(%{{.*}}: !spirv.ptr<i32, PhysicalStorageBuffer> {spirv.decoration = #spirv.decoration<Aliased>})
   spirv.func @func_arg_decoration_aliased(
       %arg0 : !spirv.ptr<i32, PhysicalStorageBuffer> { spirv.decoration = #spirv.decoration<Aliased> }
@@ -33,7 +41,7 @@ spirv.module PhysicalStorageBuffer64 GLSL450 requires #spirv.vce<v1.0,
 // -----
 
 spirv.module PhysicalStorageBuffer64 GLSL450 requires #spirv.vce<v1.0,
-    [Shader, PhysicalStorageBufferAddresses], [SPV_KHR_physical_storage_buffer]> {
+    [Shader, PhysicalStorageBufferAddresses, Linkage], [SPV_KHR_physical_storage_buffer]> {
   // CHECK-LABEL: spirv.func @func_arg_decoration_restrict(%{{.*}}: !spirv.ptr<i32, PhysicalStorageBuffer> {spirv.decoration = #spirv.decoration<Restrict>})
   spirv.func @func_arg_decoration_restrict(
       %arg0 : !spirv.ptr<i32,PhysicalStorageBuffer> { spirv.decoration = #spirv.decoration<Restrict> }
@@ -45,7 +53,7 @@ spirv.module PhysicalStorageBuffer64 GLSL450 requires #spirv.vce<v1.0,
 // -----
 
 spirv.module PhysicalStorageBuffer64 GLSL450 requires #spirv.vce<v1.0,
-    [Shader, PhysicalStorageBufferAddresses], [SPV_KHR_physical_storage_buffer]> {
+    [Shader, PhysicalStorageBufferAddresses, Linkage, GenericPointer], [SPV_KHR_physical_storage_buffer]> {
   // CHECK-LABEL: spirv.func @func_arg_decoration_aliased_pointer(%{{.*}}: !spirv.ptr<!spirv.ptr<i32, PhysicalStorageBuffer>, Generic> {spirv.decoration = #spirv.decoration<AliasedPointer>})
   spirv.func @func_arg_decoration_aliased_pointer(
       %arg0 : !spirv.ptr<!spirv.ptr<i32,PhysicalStorageBuffer>, Generic> { spirv.decoration = #spirv.decoration<AliasedPointer> }
@@ -57,7 +65,7 @@ spirv.module PhysicalStorageBuffer64 GLSL450 requires #spirv.vce<v1.0,
 // -----
 
 spirv.module PhysicalStorageBuffer64 GLSL450 requires #spirv.vce<v1.0,
-    [Shader, PhysicalStorageBufferAddresses], [SPV_KHR_physical_storage_buffer]> {
+    [Shader, PhysicalStorageBufferAddresses, Linkage, GenericPointer], [SPV_KHR_physical_storage_buffer]> {
   // CHECK-LABEL: spirv.func @func_arg_decoration_restrict_pointer(%{{.*}}: !spirv.ptr<!spirv.ptr<i32, PhysicalStorageBuffer>, Generic> {spirv.decoration = #spirv.decoration<RestrictPointer>})
   spirv.func @func_arg_decoration_restrict_pointer(
       %arg0 : !spirv.ptr<!spirv.ptr<i32,PhysicalStorageBuffer>, Generic> { spirv.decoration = #spirv.decoration<RestrictPointer> }
@@ -69,7 +77,7 @@ spirv.module PhysicalStorageBuffer64 GLSL450 requires #spirv.vce<v1.0,
 // -----
 
 spirv.module PhysicalStorageBuffer64 GLSL450 requires #spirv.vce<v1.0,
-    [Shader, PhysicalStorageBufferAddresses], [SPV_KHR_physical_storage_buffer]> {
+    [Shader, PhysicalStorageBufferAddresses, Linkage], [SPV_KHR_physical_storage_buffer]> {
   // CHECK-LABEL: spirv.func @fn1(%{{.*}}: i32, %{{.*}}: !spirv.ptr<i32, PhysicalStorageBuffer> {spirv.decoration = #spirv.decoration<Aliased>})
   spirv.func @fn1(
       %arg0: i32,