[mlir][SPIRVToLLVM] Propagate location attribute from spv.GlobalVariable to llvm.mlir.global

[Weiwei-2021]

Summary

This patch is mainly to propogate location attribute from spv.GlobalVariable to llvm.mlir.global. I keep the current parse and print approach of location since I feel that it is the autogen's direction.

It also contains three small changes. I cannot test spirv binary->spirv dialect -> llvm dialect without them. But it seems that we don't have test for spirv binary->mlir, so I did add tests.

1. Remove the restriction on UniformConstant In SPIRVToLLVM.cpp;
2. Remove the errorCheck on relaxedPrecision when deserializering SPIR-V in Deserializer.cpp
3. In SPIRVOps.cpp, let ConstantOp takes signedInteger too.

co-authered: Alan Liu <[e-mail]> and Xinyi Liu <[e-mail]>

[antiagainst]

we don't have test for spirv binary->mlir

We have roundtrip tests for it at https://github.com/llvm/llvm-project/tree/main/mlir/test/Target/SPIRV. It converts MLIR/SPIR-V dialect modules into spirv blobs and then convert back.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVStructureOps.td

398

Nit: if (

401

Nit: can you move this builder to follow the one at L410 right now? It then will form the logical order of more generic -> more specific. :)

407

Nit: if (

mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp

737

This can be a switch now given we have many cases. Also the comment needs to be updated.

749

Comments need to be updated.

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

265

Can you add a test for this in https://github.com/llvm/llvm-project/blob/main/mlir/test/Target/SPIRV/decorations.mlir to prevent regression?

mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir

70

Could you also add tests for using UniformConstant to prevent regression?

73

Hmm, I'm not sure you can have location on storage buffer variables. Should be Input or Output here?

[Weiwei-2021]

Thank you for comments.

mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir

73

You are right. It is a lazy copy paste from the above test and I didn't realize the storage class is storage buffer.....Let me correct it to avoid confusion.

[Weiwei-2021]

mlir/test/Target/SPIRV/decorations.mlir

56

I will add a new line here.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVStructureOps.td

let description = [{
     spv.GlobalVariable @var2 bind(1, 2) : !spv.ptr<f32, Uniform>
     spv.GlobalVariable @var3 built_in("GlobalInvocationId") : !spv.ptr<vector<3xi32>, Input>
     ```
   }];
 
   let arguments = (ins
     TypeAttr:$type,
     StrAttr:$sym_name,
-    OptionalAttr<FlatSymbolRefAttr>:$initializer
+    OptionalAttr<FlatSymbolRefAttr>:$initializer,
+    OptionalAttr<I32Attr>:$location,
+    OptionalAttr<I32Attr>:$binding,
+    OptionalAttr<I32Attr>:$descriptorSet,
+    OptionalAttr<StrAttr>:$builtin
   );
 
   let results = (outs);
 
   let builders = [
+    OpBuilder<(ins "TypeAttr":$type,
+                   "StringAttr":$sym_name,
+                   CArg<"FlatSymbolRefAttr", "nullptr">:$initializer),
+    [{
+      $_state.addAttribute("type", type);
+      $_state.addAttribute(sym_nameAttrName($_state.name), sym_name);
+      if (initializer)
+        $_state.addAttribute(initializerAttrName($_state.name), initializer);
+    }]>,
     OpBuilder<(ins "TypeAttr":$type, "ArrayRef<NamedAttribute>":$namedAttrs),
     [{
       $_state.addAttribute("type", type);
       $_state.addAttributes(namedAttrs);
     }]>,
     OpBuilder<(ins "Type":$type, "StringRef":$name,
       "unsigned":$descriptorSet, "unsigned":$binding)>,
     OpBuilder<(ins "Type":$type, "StringRef":$name,
-      "spirv::BuiltIn":$builtin)>
+      "spirv::BuiltIn":$builtin)>,
+    OpBuilder<(ins "Type":$type,
+                   "StringRef":$sym_name,
+                    CArg<"FlatSymbolRefAttr", "{}">:$initializer),
+    [{
+      $_state.addAttribute("type", TypeAttr::get(type));
+      $_state.addAttribute(sym_nameAttrName($_state.name), $_builder.getStringAttr(sym_name));
+      if (initializer)
+        $_state.addAttribute(initializerAttrName($_state.name), initializer);
+    }]>
   ];
 
   let hasOpcode = 0;
 
   let autogenSerialization = 0;
 
   let extraClassDeclaration = [{
     ::mlir::spirv::StorageClass storageClass() {

mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp

matchAndRewrite(spirv::GlobalVariableOp op, OpAdaptor adaptor,
     auto dstType = typeConverter.convertType(srcType.getPointeeType());
     if (!dstType)
       return failure();
 
     // Limit conversion to the current invocation only or `StorageBuffer`
     // required by SPIR-V runner.
     // This is okay because multiple invocations are not supported yet.
     auto storageClass = srcType.getStorageClass();
-    if (storageClass != spirv::StorageClass::Input &&
-        storageClass != spirv::StorageClass::Private &&
-        storageClass != spirv::StorageClass::Output &&
-        storageClass != spirv::StorageClass::StorageBuffer) {
+    switch (storageClass) {
+    case spirv::StorageClass::Input:
+    case spirv::StorageClass::Private:
+    case spirv::StorageClass::Output:
+    case spirv::StorageClass::StorageBuffer:
+    case spirv::StorageClass::UniformConstant:
+      break;
+    default:
       return failure();
     }
 
     // LLVM dialect spec: "If the global value is a constant, storing into it is
-    // not allowed.". This corresponds to SPIR-V 'Input' storage class that is
-    // read-only.
-    bool isConstant = storageClass == spirv::StorageClass::Input;
+    // not allowed.". This corresponds to SPIR-V 'Input' and 'UniformConstant'
+    // storage class that is read-only.
+    bool isConstant = (storageClass == spirv::StorageClass::Input) ||
+                      (storageClass == spirv::StorageClass::UniformConstant);
     // SPIR-V spec: "By default, functions and global variables are private to a
     // module and cannot be accessed by other modules. However, a module may be
     // written to export or import functions and global (module scope)
     // variables.". Therefore, map 'Private' storage class to private linkage,
     // 'Input' and 'Output' to external linkage.
     auto linkage = storageClass == spirv::StorageClass::Private
                        ? LLVM::Linkage::Private
                        : LLVM::Linkage::External;
-    rewriter.replaceOpWithNewOp<LLVM::GlobalOp>(
+    auto newGlobalOp = rewriter.replaceOpWithNewOp<LLVM::GlobalOp>(
         op, dstType, isConstant, linkage, op.sym_name(), Attribute(),
         /*alignment=*/0);
+
+    // Attach location attribute if applicable
+    if (op.locationAttr())
+      newGlobalOp->setAttr(op.locationAttrName(), op.locationAttr());
+
     return success();
   }
 };
 
 /// Converts SPIR-V cast ops that do not have straightforward LLVM
 /// equivalent in LLVM dialect.
 template <typename SPIRVOp, typename LLVMExtOp, typename LLVMTruncOp>
 class IndirectCastPattern : public SPIRVToLLVMConversion<SPIRVOp> {

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

static LogicalResult extractValueFromConstOp(Operation *op, int32_t &value) {
   if (!constOp) {
     return failure();
   }
   auto valueAttr = constOp.value();
   auto integerValueAttr = valueAttr.dyn_cast<IntegerAttr>();
   if (!integerValueAttr) {
     return failure();
   }
+
+  if (integerValueAttr.getType().isSignlessInteger())
     value = integerValueAttr.getInt();
+  else
+    value = integerValueAttr.getSInt();
+
   return success();
 }
 
 template <typename Ty>
 static ArrayAttr
 getStrArrayAttrForEnumList(Builder &builder, ArrayRef<Ty> enumValues,
                            function_ref<StringRef(Ty)> stringifyFn) {
   if (enumValues.empty()) {

 
 //===----------------------------------------------------------------------===//
 // spv.GlobalVariable
 //===----------------------------------------------------------------------===//
 
 void spirv::GlobalVariableOp::build(OpBuilder &builder, OperationState &state,
                                     Type type, StringRef name,
                                     unsigned descriptorSet, unsigned binding) {
-  build(builder, state, TypeAttr::get(type), builder.getStringAttr(name),
-        nullptr);
+  build(builder, state, TypeAttr::get(type), builder.getStringAttr(name));
   state.addAttribute(
       spirv::SPIRVDialect::getAttributeName(spirv::Decoration::DescriptorSet),
       builder.getI32IntegerAttr(descriptorSet));
   state.addAttribute(
       spirv::SPIRVDialect::getAttributeName(spirv::Decoration::Binding),
       builder.getI32IntegerAttr(binding));
 }
 
 void spirv::GlobalVariableOp::build(OpBuilder &builder, OperationState &state,
                                     Type type, StringRef name,
                                     spirv::BuiltIn builtin) {
-  build(builder, state, TypeAttr::get(type), builder.getStringAttr(name),
-        nullptr);
+  build(builder, state, TypeAttr::get(type), builder.getStringAttr(name));
   state.addAttribute(
       spirv::SPIRVDialect::getAttributeName(spirv::Decoration::BuiltIn),
       builder.getStringAttr(spirv::stringifyBuiltIn(builtin)));
 }
 
 static ParseResult parseGlobalVariableOp(OpAsmParser &parser,
                                          OperationState &state) {
   // Parse variable name.

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

LogicalResult spirv::Deserializer::processDecoration(ArrayRef<uint32_t> words) {
   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:
+  case spirv::Decoration::RelaxedPrecision:
     if (words.size() != 2) {
       return emitError(unknownLoc, "OpDecoration with ")
              << decorationName << "needs a single target <id>";
     }
     // Block decoration does not affect spv.struct type, but is still stored for
     // verification.
     // TODO: Update StructType to contain this information since
     // it is needed for many validation rules.

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

case spirv::Decoration::BuiltIn:
     }
     return emitError(loc, "expected string attribute for ") << attrName;
   case spirv::Decoration::Aliased:
   case spirv::Decoration::Flat:
   case spirv::Decoration::NonReadable:
   case spirv::Decoration::NonWritable:
   case spirv::Decoration::NoPerspective:
   case spirv::Decoration::Restrict:
+  case spirv::Decoration::RelaxedPrecision:
     // For unit attributes, the args list has no values so we do nothing
     if (auto unitAttr = attr.second.dyn_cast<UnitAttr>())
       break;
     return emitError(loc, "expected unit attribute for ") << attrName;
   default:
     return emitError(loc, "unhandled decoration ") << decorationName;
   }
   return emitDecoration(resultID, decoration.getValue(), args);

mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir

spv.module @name Logical GLSL450 {
   //       CHECK:   llvm.mlir.addressof @name_bar_descriptor_set0_binding0 : !llvm.ptr<i32>
   spv.GlobalVariable @bar bind(0, 0) : !spv.ptr<i32, StorageBuffer>
   spv.func @foo() "None" {
     %0 = spv.mlir.addressof @bar : !spv.ptr<i32, StorageBuffer>
     spv.Return
   }
 }
 
+spv.module Logical GLSL450 {
+  // CHECK: llvm.mlir.global external @bar() {location = 1 : i32} : i32
+  // CHECK-LABEL: @foo
+  spv.GlobalVariable @bar {location = 1 : i32} : !spv.ptr<i32, Output>
+  spv.func @foo() "None" {
+    %0 = spv.mlir.addressof @bar : !spv.ptr<i32, Output>
+    spv.Return
+  }
+}
+
+spv.module Logical GLSL450 {
+  // CHECK: llvm.mlir.global external constant @bar() {location = 3 : i32} : f32
+  // CHECK-LABEL: @foo
+  spv.GlobalVariable @bar {descriptor_set = 0 : i32, location = 3 : i32} : !spv.ptr<f32, UniformConstant>
+  spv.func @foo() "None" {
+    %0 = spv.mlir.addressof @bar : !spv.ptr<f32, UniformConstant>
+    spv.Return
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // spv.Load
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @load
 spv.func @load() "None" {
   %0 = spv.Variable : !spv.ptr<f32, Function>
   //  CHECK: llvm.load %{{.*}} : !llvm.ptr<f32>

mlir/test/Target/SPIRV/decorations.mlir

 
 // -----
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   // CHECK: restrict
   spv.GlobalVariable @var bind(0, 0) {restrict} : !spv.ptr<!spv.struct<(!spv.array<4xf32, stride=4>[0])>, StorageBuffer>
 }
 
+// -----
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+  // CHECK: relaxed_precision
+  spv.GlobalVariable @var {location = 0 : i32, relaxed_precision} : !spv.ptr<vector<4xf32>, Output>
+}
+