[mlir][spirv] Add support to import/export Functions and Global variables

[mshahneo]

Summary

"LinkageAttributes" decoration allow a SPIR-V module to import
external functions and global variables, or export functions or
global variables for other SPIR-V modules to link against and use.

Import/export capability is extremely important when using outside
libraries (e.g., intrinsic libraries).

Added decorations:

• LinkageAttributes

[mshahneo]

Add Linkage capabilities to the test cases

[mshahneo]

Remove an unnecessary test

[antiagainst]

Cool, thanks for the contribution! Adding the support for this is certainly fine; but I think we should make it more integrated in the implementation. Two major issues to address before going into details. :)

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

3520

Instead of just attaching this to the function and global variable as a normal attribute, I'd think we want to modify the definition of spirv.func and spirv.GlobalVariable to include it as a native (optional) attribute. That is, update those two ops' ODS to includ this. It would be better for printing/verification and to avoid accidentally dropping this.

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

289

Can we actually have a proper MLIR attribute definition of LinkageAttributes? You can see SPIRVAttribute.td/h/cpp where we have a few examples there. A proper MLIR attribute, like #spirv.linkage<Import, "ExternalName"> would make it easier to read/print/verify.

[mshahneo]

In D148749#4284507, @antiagainst wrote:

Cool, thanks for the contribution! Adding the support for this is certainly fine; but I think we should make it more integrated in the implementation. Two major issues to address before going into details. :)

Absolutely, can do. Yep, that would be nice, I wasn't sure of that would be too invasive. Will update the PR. Thank you so much :)

[mshahneo]

Define LinkageAttributes as a prpoper MLIR attribute
Update the implementation accordingly
Update the test cases

[mshahneo]

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

3520

Updated the definition of spirv.func and spirv.GlobalVariable

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

289

Created a new SPIR-V attribute.

[antiagainst]

Nice! Just a few more nits.

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

280

Strictly this should be 3 + ceildiv(strlen(name), 4) + 1?

281

For unused code please just remove it instead of commenting out.

471

Nit: use . to end the sentence to be consistent with other comments.

Applies to all comments in other places.

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

243

Flip the condition check and do early returns for this case so we can avoid one level of indent?

300

Can we use a local variable for this? Multi-line if conditions is not very nice to readability.

mlir/test/Dialect/SPIRV/IR/function-decorations.mlir

13

I think we can improve the assembly by customizing it to spirv.func @... "Pure" Import["outside.func"], which reads nicer. I won't block it on this though; given this adds the functionality. Would you be interested to handle it later?

15

Nit: indent this to show that linkage_name and linkage_type is nested in linkage_attributes?

Applies to other places too.

[mshahneo]

Added some stylistic change.

[mshahneo]

Thank you so much for pointing out some nit issues :). I tried to address them.

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

280

I think it should be

3 + strlen(name)/4 + 1

or

3 + ceildiv(strlen(name), 4)

, right?

mlir/test/Dialect/SPIRV/IR/function-decorations.mlir

13

I think we can do it later, it does read nice, but do we actually want to do this? It kind of goes away from the normal rule of passing attributes.

[mshahneo]

A 'space' was added on an error message.

[antiagainst]

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

280

Yes actually. (I was thinking the trailing 1 is for "Linkage Type". Sorry for the confusion.)

mlir/test/Dialect/SPIRV/IR/function-decorations.mlir

13

I'm not sure what "goes away from the normal rule of passing attributes" means? It's common to define custom assembly to avoid boilerplate in parsing/printing. But anyway, we can do it later. :)

[mshahneo]

Hi @antiagainst ,
Could you please help me commit the change, I do not have commit access.
Thank you so much in advance :).

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

280

No problem at all :). Thank you so much for looking into it :).

mlir/test/Dialect/SPIRV/IR/function-decorations.mlir

13

Ah, by normal rule I mean, nested way of writing, attribute_name=value. But yes, LinakgeType["LinkageName"] would read nice :).

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

if (auto entryPointOp = dyn_cast<spirv::EntryPointOp>(op)) {
       auto key = std::pair<spirv::FuncOp, spirv::ExecutionModel>(
           funcOp, entryPointOp.getExecutionModel());
       auto entryPtIt = entryPoints.find(key);
       if (entryPtIt != entryPoints.end()) {
         return entryPointOp.emitError("duplicate of a previous EntryPointOp");
       }
       entryPoints[key] = entryPointOp;
     } else if (auto funcOp = dyn_cast<spirv::FuncOp>(op)) {
-      if (funcOp.isExternal())
-        return op.emitError("'spirv.module' cannot contain external functions");
+      // If the function is external and does not have 'Import'
+      // linkage_attributes(LinkageAttributes), throw an error; 'Import'
+      // LinkageAttributes is used to import external functions
+      auto linkageAttr = funcOp->getAttr("linkage_attributes");
+      if (funcOp.isExternal() &&
+          !(linkageAttr && spirv::symbolizeEnum<spirv::LinkageType>(
+                               linkageAttr.dyn_cast<ArrayAttr>()[1]
+                                   .dyn_cast<StringAttr>()
+                                   .strref()) == spirv::LinkageType::Import))
+        return op.emitError(
+            "'spirv.module' cannot contain external functions "
+            "without 'Import' linkage_attributes(LinkageAttributes)");
 
       // TODO: move this check to spirv.func.
       for (auto &block : funcOp)
         for (auto &op : block) {
           if (op.getDialect() != dialect)
             return op.emitError(
                 "functions in 'spirv.module' can only contain spirv.* ops");
         }

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

case spirv::Decoration::BuiltIn:
     break;
   case spirv::Decoration::ArrayStride:
     if (words.size() != 3) {
       return emitError(unknownLoc, "OpDecorate with ")
              << decorationName << " needs a single integer literal";
     }
     typeDecorations[words[0]] = words[2];
     break;
+  case spirv::Decoration::LinkageAttributes: {
+    if (words.size() < 4) {
+      return emitError(unknownLoc, "OpDecorate with ")
+             << decorationName
+             << " needs at least 1 string and 1 integer literal";
+    }
+    // LinkageAttributes has two parameters ["Name", "LinkageType"]
+    // e.g., OpDecorate %imported_func LinkageAttributes "outside.func" Import
+    // "name" is a stringliteral encoded as uint32_t,
+    // hence the size of name is variable length which results in words.size()
+    // being variable length, words.size() = 3 + strlen(name)/4 + 1
+    // uint32_t linkageNameLen = words.size() - 3;
+    unsigned wordIndex = 2;
+    StringRef linkageName = spirv::decodeStringLiteral(words, wordIndex);
+    std::vector<mlir::Attribute> attrElements;
+    attrElements.push_back(opBuilder.getStringAttr(linkageName));
+    attrElements.push_back(opBuilder.getStringAttr(stringifyLinkageType(
+        static_cast<spirv::LinkageType>(words[words.size() - 1]))));
+    ArrayAttr linkageAttrVal = opBuilder.getArrayAttr(attrElements);
+    decorations[words[0]].set(symbol, linkageAttrVal.dyn_cast<Attribute>());
+    break;
+  }
   case spirv::Decoration::Aliased:
   case spirv::Decoration::Block:
   case spirv::Decoration::BufferBlock:
   case spirv::Decoration::Flat:
   case spirv::Decoration::NonReadable:
   case spirv::Decoration::NonWritable:
   case spirv::Decoration::NoPerspective:
   case spirv::Decoration::Restrict:
if ((isVoidType(resultType) && functionType.getNumResults() != 0) ||
        functionType.getResult(0) != resultType)) {
     return emitError(unknownLoc, "mismatch in function type ")
            << functionType << " and return type " << resultType << " specified";
   }
 
   std::string fnName = getFunctionSymbol(fnID);
   auto funcOp = opBuilder.create<spirv::FuncOp>(
       unknownLoc, fnName, functionType, fnControl.value());
+  bool isImportedFunc = false;
+  // Processing other function attributes
+  if (decorations.count(fnID)) {
+    for (auto attr : decorations[fnID].getAttrs()) {
+      funcOp->setAttr(attr.getName(), attr.getValue());
+      if (attr.getName() == "linkage_attributes" &&
+          spirv::symbolizeEnum<spirv::LinkageType>(attr.getValue()
+                                                       .dyn_cast<ArrayAttr>()[1]
+                                                       .dyn_cast<StringAttr>()
+                                                       .strref()) ==
+              spirv::LinkageType::Import)
+        isImportedFunc = true;
+    }
+  }
   curFunction = funcMap[fnID] = funcOp;
   auto *entryBlock = funcOp.addEntryBlock();
   LLVM_DEBUG({
     logger.startLine()
         << "//===-------------------------------------------===//\n";
     logger.startLine() << "[fn] name: " << fnName << "\n";
     logger.startLine() << "[fn] type: " << fnType << "\n";
     logger.startLine() << "[fn] ID: " << fnID << "\n";
for (size_t i = 0, e = functionType.getNumInputs(); i != e; ++i) {
         return emitError(unknownLoc, "duplicate definition of result <id> ")
                << operands[1];
       }
       auto argValue = funcOp.getArgument(i);
       valueMap[operands[1]] = argValue;
     }
   }
 
+  // entryBlock is needed to access the arguments, Once that is done, we can
+  // erase the block for functions with 'Import' LinkageAttributes, since these
+  // are essentially function declarations, so they have no body
+  if (isImportedFunc)
+    funcOp.eraseBody();
   // RAII guard to reset the insertion point to the module's region after
   // deserializing the body of this function.
   OpBuilder::InsertionGuard moduleInsertionGuard(opBuilder);
 
   spirv::Opcode opcode = spirv::Opcode::OpNop;
   ArrayRef<uint32_t> instOperands;
 
   // Special handling for the entry block. We need to make sure it starts with

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

 #include "Serializer.h"
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/IR/RegionGraphTraits.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Target/SPIRV/SPIRVBinaryUtils.h"
 #include "llvm/ADT/DepthFirstIterator.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "spirv-serialization"
 
 using namespace mlir;
 
 /// A pre-order depth-first visitor function for processing basic blocks.
 ///
LogicalResult Serializer::processFuncOp(spirv::FuncOp op) {
   operands.push_back(static_cast<uint32_t>(op.getFunctionControl()));
   operands.push_back(fnTypeID);
   encodeInstructionInto(functionHeader, spirv::Opcode::OpFunction, operands);
 
   // Add function name.
   if (failed(processName(funcID, op.getName()))) {
     return failure();
   }
-
+  // Handle external functions with linkage_attributes(LinkageAttributes)
+  // differently
+  if (op.isExternal()) {
+    auto linkageAttr = op->getAttr("linkage_attributes");
+    if (linkageAttr && spirv::symbolizeEnum<spirv::LinkageType>(
+                           linkageAttr.dyn_cast<ArrayAttr>()[1]
+                               .dyn_cast<StringAttr>()
+                               .strref()) == spirv::LinkageType::Import) {
+      // Add an entry block to set up the block arguments
+      // to match the signature of the function.
+      // This is to generate OpFunctionParameter for functions with
+      // LinkageAttributes
+      // WARNING: This operation has side-effect, it essentially adds a body
+      // to the func. Hence, making it not external anymore (isExternal()
+      // is going to return false for this function from now on)
+      // Hence, we'll remove the body once we are done with the serialization
+      op.addEntryBlock();
+      for (auto arg : op.getArguments()) {
+        uint32_t argTypeID = 0;
+        if (failed(processType(op.getLoc(), arg.getType(), argTypeID))) {
+          return failure();
+        }
+        auto argValueID = getNextID();
+        valueIDMap[arg] = argValueID;
+        encodeInstructionInto(functionHeader,
+                              spirv::Opcode::OpFunctionParameter,
+                              {argTypeID, argValueID});
+      }
+      // Don't need to process the added block, there is nothing to process, the
+      // fake body was added just to get the arguments, remove the body, since
+      // it's use is done
+      op.eraseBody();
+    } else
+      return op.emitError(
+          "'spirv.module' cannot contain external functions "
+          "without 'Import' linkage_attributes(LinkageAttributes)");
+  } else {
     // Declare the parameters.
     for (auto arg : op.getArguments()) {
       uint32_t argTypeID = 0;
       if (failed(processType(op.getLoc(), arg.getType(), argTypeID))) {
         return failure();
       }
       auto argValueID = getNextID();
       valueIDMap[arg] = argValueID;
       encodeInstructionInto(functionHeader, spirv::Opcode::OpFunctionParameter,
                             {argTypeID, argValueID});
     }
 
-  // Process the body.
-  if (op.isExternal()) {
-    return op.emitError("external function is unhandled");
-  }
-
     // Some instructions (e.g., OpVariable) in a function must be in the first
     // block in the function. These instructions will be put in functionHeader.
-  // Thus, we put the label in functionHeader first, and omit it from the first
-  // block.
+    // Thus, we put the label in functionHeader first, and omit it from the
+    // first block.
+    // OpLabel only needs to be added for functions with body (including empty
+    // body) Since, we added a fake body for functions with Import Linkage
+    // attributes, These functions are essentially function delcaration, so they
+    // should not have OpLabel and a terminating instruction
     encodeInstructionInto(functionHeader, spirv::Opcode::OpLabel,
                           {getOrCreateBlockID(&op.front())});
     if (failed(processBlock(&op.front(), /*omitLabel=*/true)))
       return failure();
     if (failed(visitInPrettyBlockOrder(
             &op.front(), [&](Block *block) { return processBlock(block); },
             /*skipHeader=*/true))) {
       return failure();
     }
 
-  // There might be OpPhi instructions who have value references needing to fix.
+    // There might be OpPhi instructions who have value references needing to
+    // fix.
     for (const auto &deferredValue : deferredPhiValues) {
       Value value = deferredValue.first;
       uint32_t id = getValueID(value);
       LLVM_DEBUG(llvm::dbgs() << "[phi] fix reference of value " << value
                               << " to id = " << id << '\n');
       assert(id && "OpPhi references undefined value!");
       for (size_t offset : deferredValue.second)
         functionBody[offset] = id;
     }
     deferredPhiValues.clear();
-
+  }
   LLVM_DEBUG(llvm::dbgs() << "-- completed function '" << op.getName()
                           << "' --\n");
+  // Insert Decorations based on Function Attributes
+  // Only attributes we should be considering for decoration are the
+  // ::mlir::spirv::Decoration attributes
+  for (auto attr : op->getAttrs()) {
+    // Only generate OpDecorate op for spirv::Decoration attributes
+    if (mlir::spirv::symbolizeEnum<spirv::Decoration>(
+            llvm::convertToCamelFromSnakeCase(attr.getName().strref(),
+                                              /*capitalizeFirst=*/true)) !=
+        std::nullopt) {
+      if (failed(processDecoration(op.getLoc(), funcID, attr))) {
+        return failure();
+      }
+    }
+  }
   // Insert OpFunctionEnd.
   encodeInstructionInto(functionBody, spirv::Opcode::OpFunctionEnd, {});
 
   functions.append(functionHeader.begin(), functionHeader.end());
   functions.append(functionBody.begin(), functionBody.end());
   functionHeader.clear();
   functionBody.clear();
 

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

auto am = static_cast<uint32_t>(
           .getValue());
 
   encodeInstructionInto(memoryModel, spirv::Opcode::OpMemoryModel, {am, mm});
 }
 
 LogicalResult Serializer::processDecoration(Location loc, uint32_t resultID,
                                             NamedAttribute attr) {
   auto attrName = attr.getName().strref();
-  auto decorationName = llvm::convertToCamelFromSnakeCase(attrName, true);
+  auto decorationName =
+      llvm::convertToCamelFromSnakeCase(attrName, /*capitalizeFirst=*/true);
   auto decoration = spirv::symbolizeDecoration(decorationName);
   if (!decoration) {
     return emitError(
                loc, "non-argument attributes expected to have snake-case-ified "
                     "decoration name, unhandled attribute with name : ")
            << attrName;
   }
   SmallVector<uint32_t, 1> args;
   switch (*decoration) {
+  case spirv::Decoration::LinkageAttributes: {
+    // Get the value of the Linkage Attributes
+    // e.g., LinkageAttributes=["Name", "LinkageType"]
+    // TODO: check if attribute values are passed in the following format
+    // LinkageAttributes=["Name", "LinkageType"]
+    // At this point we assume, they are passed in this format
+    auto arrayAttrVal = attr.getValue().dyn_cast<ArrayAttr>();
+    if (arrayAttrVal.size() != 2)
+      return emitError(loc, "attribute must have 2 values ") << attrName;
+    // Encode the Linkage Name (string literal to uint32_t)
+    spirv::encodeStringLiteralInto(
+        args, arrayAttrVal[0].dyn_cast<StringAttr>().strref());
+    // Encode LinkageType
+    // Add the Linkagetype to the args
+    auto linkageType = static_cast<uint32_t>(
+        spirv::symbolizeEnum<spirv::LinkageType>(
+            arrayAttrVal[1].dyn_cast<StringAttr>().strref())
+            .value());
+    args.push_back(linkageType);
+    break;
+  }
   case spirv::Decoration::Binding:
   case spirv::Decoration::DescriptorSet:
   case spirv::Decoration::Location:
     if (auto intAttr = attr.getValue().dyn_cast<IntegerAttr>()) {
       args.push_back(intAttr.getValue().getZExtValue());
       break;
     }
     return emitError(loc, "expected integer attribute for ") << attrName;

mlir/test/Dialect/SPIRV/IR/function-decorations.mlir

+// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+    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 = ["outSideGlobalVar1", "Import"]
+    spirv.GlobalVariable @var1 {linkage_attributes=["outSideGlobalVar1", "Import"]} : !spirv.ptr<f32, Private>
+    
+    // CHECK: linkage_attributes = ["outside.func", "Import"]
+    spirv.func @outside.func.with.linkage(%arg0 : i8) -> () "Pure" attributes {linkage_attributes=["outside.func", "Import"]}
+    spirv.func @inside.func() -> () "Pure" attributes {} {spirv.Return}
+}

mlir/test/Dialect/SPIRV/IR/structure-ops.mlir

 
 // CHECK: spirv.func @baz(%{{.+}}: i32) "DontInline" attributes {additional_stuff = 64 : i64}
 spirv.func @baz(%arg: i32) "DontInline" attributes {
   additional_stuff = 64
 } { spirv.Return }
 
 // -----
 
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+    // CHECK: linkage_attributes = ["outside.func", "Import"]
+    spirv.func @outside.func.with.linkage(%arg0 : i8) -> () "Pure" attributes {linkage_attributes=["outside.func", "Import"]}
+    spirv.func @inside.func() -> () "Pure" attributes {} {spirv.Return}
+}
+// -----
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> { 
+  // expected-error @+1 {{'spirv.module' cannot contain external functions without 'Import' linkage_attributes(LinkageAttributes)}}
+  spirv.func @outside.func.with.linkage(%arg0 : i8) -> () "Pure"
+  spirv.func @inside.func() -> () "Pure" attributes {} {spirv.Return}
+}
+
+// -----
+
 // expected-error @+1 {{expected function_control attribute specified as string}}
 spirv.func @missing_function_control() { spirv.Return }
 
 // -----
 
 // expected-error @+1 {{cannot have more than one result}}
 spirv.func @cannot_have_more_than_one_result(%arg: i32) -> (i32, i32) "None"
 

 // Allow in other module-like ops
 module {
   // CHECK: spirv.GlobalVariable
   spirv.GlobalVariable @var0 : !spirv.ptr<f32, Input>
 }
 
 // -----
 
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+  // CHECK: linkage_attributes = ["outSideGlobalVar1", "Import"]
+  spirv.GlobalVariable @var1 {linkage_attributes=["outSideGlobalVar1", "Import"]} : !spirv.ptr<f32, Private>
+}
+
+// -----
+
 spirv.module Logical GLSL450 {
   // expected-error @+1 {{expected spirv.ptr type}}
   spirv.GlobalVariable @var0 : f32
 }
 
 // -----
 
 spirv.module Logical GLSL450 {

 spirv.module Logical GLSL450 {
   spirv.func @foo() "None" {
     // expected-error @+1 {{op must appear in a module-like op's block}}
     spirv.GlobalVariable @var0 : !spirv.ptr<f32, Input>
     spirv.Return
   }
 }
 
+
 // -----
 
 //===----------------------------------------------------------------------===//
 // spirv.module
 //===----------------------------------------------------------------------===//
 
 // Module without capability and extension
 // CHECK: spirv.module Logical GLSL450

mlir/test/Target/SPIRV/decorations.mlir

 }
 
 // -----
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK: relaxed_precision
   spirv.GlobalVariable @var {location = 0 : i32, relaxed_precision} : !spirv.ptr<vector<4xf32>, Output>
 }
+// -----
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+  // CHECK: linkage_attributes = ["outSideGlobalVar1", "Import"]
+  spirv.GlobalVariable @var1 {linkage_attributes=["outSideGlobalVar1", "Import"]} : !spirv.ptr<f32, Private>
+}
 

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

+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, Linkage], []> {
+    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 = ["outside.func", "Import"]
+    spirv.func @outside.func.with.linkage(%arg0 : i8) -> () "Pure" attributes {linkage_attributes=["outside.func", "Import"]}
+    spirv.func @inside.func() -> () "Pure" attributes {} {spirv.Return}
+}