[MLIR][SPIRV] Extend _reference_of to support SpecConstantCompositeOp.

[ergawy]

Summary

Adds support for SPIR-V composite speciailization constants to spv._reference_of.

[ergawy]

• Rebase (testing how stacked reviews work).

[ergawy]

Looks like I misunderstood how stacked reviews work and that's why the CI fails. Is there a workflow that you can share for creating such stacks of reviews?

[ergawy]

• Merge remote-tracking branch 'origin/master' into mlir/spirv/support_op_spec_constant_composite

[antiagainst]

Thanks!

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

 }
 
 // -----
 
 def SPV_ReferenceOfOp : SPV_Op<"_reference_of", [NoSideEffect]> {
   let summary = "Reference a specialization constant.";
 
   let description = [{
-    Specialization constant in module scope are defined using symbol names.
+    Specialization constants in module scope are defined using symbol names.
     This op generates an SSA value that can be used to refer to the symbol
     within function scope for use in ops that expect an SSA value.
     This operation has no corresponding SPIR-V instruction; it's merely used
     for modelling purpose in the SPIR-V dialect. This op's return type is
     the same as the specialization constant.
 
     <!-- End of AutoGen section -->
 

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

static LogicalResult verify(spirv::ModuleOp moduleOp) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // spv._reference_of
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(spirv::ReferenceOfOp referenceOfOp) {
-  auto specConstOp = dyn_cast_or_null<spirv::SpecConstantOp>(
-      SymbolTable::lookupNearestSymbolFrom(referenceOfOp.getParentOp(),
-                                           referenceOfOp.spec_const()));
-  if (!specConstOp) {
-    return referenceOfOp.emitOpError("expected spv.specConstant symbol");
-  }
-  if (referenceOfOp.reference().getType() !=
-      specConstOp.default_value().getType()) {
+  auto *specConstSym = SymbolTable::lookupNearestSymbolFrom(
+      referenceOfOp.getParentOp(), referenceOfOp.spec_const());
+  Type constType;
+
+  auto specConstOp = dyn_cast_or_null<spirv::SpecConstantOp>(specConstSym);
+  if (specConstOp)
+    constType = specConstOp.default_value().getType();
+
+  auto specConstCompositeOp =
+      dyn_cast_or_null<spirv::SpecConstantCompositeOp>(specConstSym);
+  if (specConstCompositeOp)
+    constType = specConstCompositeOp.type();
+
+  if (!specConstOp && !specConstCompositeOp)
+    return referenceOfOp.emitOpError(
+        "expected spv.specConstant or spv.SpecConstantComposite symbol");
+
+  if (referenceOfOp.reference().getType() != constType)
     return referenceOfOp.emitOpError("result type mismatch with the referenced "
                                      "specialization constant's type");
-  }
+
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // spv.Return
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(spirv::ReturnOp returnOp) {

mlir/lib/Dialect/SPIRV/Serialization/Deserializer.cpp

private:
   /// exists; otherwise creates one based on the <id>.
   std::string getSpecConstantSymbol(uint32_t id);
 
   /// Gets the specialization constant with the given result <id>.
   spirv::SpecConstantOp getSpecConstant(uint32_t id) {
     return specConstMap.lookup(id);
   }
 
+  /// Gets the composite specialization constant with the given result <id>.
+  spirv::SpecConstantCompositeOp getSpecConstantComposite(uint32_t id) {
+    return specConstCompositeMap.lookup(id);
+  }
+
   /// Creates a spirv::SpecConstantOp.
   spirv::SpecConstantOp createSpecConstant(Location loc, uint32_t resultID,
                                            Attribute defaultValue);
 
   /// Processes the OpVariable instructions at current `offset` into `binary`.
   /// It is expected that this method is used for variables that are to be
   /// defined at module scope and will be deserialized into a spv.globalVariable
   /// instruction.
private:
   /// In the SPIR-V binary format, all constants are placed in the module and
   /// shared by instructions at module level and in subsequent functions. But in
   /// the SPIR-V dialect, we materialize the constant to where it's used in the
   /// function. So when seeing a constant instruction in the binary format, we
   /// don't immediately emit a constant op into the module, we keep its value
   /// (and type) here. Later when it's used, we materialize the constant.
   DenseMap<uint32_t, std::pair<Attribute, Type>> constantMap;
 
-  // Result <id> to variable mapping.
+  // Result <id> to spec constant mapping.
   DenseMap<uint32_t, spirv::SpecConstantOp> specConstMap;
 
+  // Result <id> to composite spec constant mapping.
+  DenseMap<uint32_t, spirv::SpecConstantCompositeOp> specConstCompositeMap;
+
   // Result <id> to variable mapping.
   DenseMap<uint32_t, spirv::GlobalVariableOp> globalVariableMap;
 
   // Result <id> to function mapping.
   DenseMap<uint32_t, spirv::FuncOp> funcMap;
 
   // Result <id> to block mapping.
   DenseMap<uint32_t, Block *> blockMap;
Deserializer::processSpecConstantComposite(ArrayRef<uint32_t> operands) {
   }
 
   Type resultType = getType(operands[0]);
   if (!resultType) {
     return emitError(unknownLoc, "undefined result type from <id> ")
            << operands[0];
   }
 
-  auto symName = opBuilder.getStringAttr(getSpecConstantSymbol(operands[1]));
+  auto resultID = operands[1];
+  auto symName = opBuilder.getStringAttr(getSpecConstantSymbol(resultID));
 
   SmallVector<Attribute, 4> elements;
   elements.reserve(operands.size() - 2);
   for (unsigned i = 2, e = operands.size(); i < e; ++i) {
     auto elementInfo = getSpecConstant(operands[i]);
     elements.push_back(opBuilder.getSymbolRefAttr(elementInfo));
   }
 
-  opBuilder.create<spirv::SpecConstantCompositeOp>(
+  auto op = opBuilder.create<spirv::SpecConstantCompositeOp>(
       unknownLoc, TypeAttr::get(resultType), symName,
       opBuilder.getArrayAttr(elements));
+  specConstCompositeMap[resultID] = op;
 
   return success();
 }
 
 LogicalResult Deserializer::processConstantNull(ArrayRef<uint32_t> operands) {
   if (operands.size() != 2) {
     return emitError(unknownLoc,
                      "OpConstantNull must have type <id> and result <id>");
if (auto varOp = getGlobalVariable(id)) {
     return addressOfOp.pointer();
   }
   if (auto constOp = getSpecConstant(id)) {
     auto referenceOfOp = opBuilder.create<spirv::ReferenceOfOp>(
         unknownLoc, constOp.default_value().getType(),
         opBuilder.getSymbolRefAttr(constOp.getOperation()));
     return referenceOfOp.reference();
   }
+  if (auto constCompositeOp = getSpecConstantComposite(id)) {
+    auto referenceOfOp = opBuilder.create<spirv::ReferenceOfOp>(
+        unknownLoc, constCompositeOp.type(),
+        opBuilder.getSymbolRefAttr(constCompositeOp.getOperation()));
+    return referenceOfOp.reference();
+  }
   if (auto undef = getUndefType(id)) {
     return opBuilder.create<spirv::UndefOp>(unknownLoc, undef);
   }
   return valueMap.lookup(id);
 }
 
 LogicalResult
 Deserializer::sliceInstruction(spirv::Opcode &opcode,

mlir/test/Dialect/SPIRV/Serialization/spec-constant.mlir

 // RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   // CHECK: spv.specConstant @sc_true = true
   spv.specConstant @sc_true = true
   // CHECK: spv.specConstant @sc_false spec_id(1) = false
   spv.specConstant @sc_false spec_id(1) = false
 
   // CHECK: spv.specConstant @sc_int = -5 : i32
   spv.specConstant @sc_int = -5 : i32
 
   // CHECK: spv.specConstant @sc_float spec_id(5) = 1.000000e+00 : f32
   spv.specConstant @sc_float spec_id(5) = 1. : f32
 
+  // CHECK: spv.specConstantComposite @scc (@sc_int, @sc_int) : !spv.array<2 x i32>
+  spv.specConstantComposite @scc (@sc_int, @sc_int) : !spv.array<2 x i32>
+
   // CHECK-LABEL: @use
   spv.func @use() -> (i32) "None" {
     // We materialize a `spv._reference_of` op at every use of a
     // specialization constant in the deserializer. So two ops here.
     // CHECK: %[[USE1:.*]] = spv._reference_of @sc_int : i32
     // CHECK: %[[USE2:.*]] = spv._reference_of @sc_int : i32
     // CHECK: spv.IAdd %[[USE1]], %[[USE2]]
 
     %0 = spv._reference_of @sc_int : i32
     %1 = spv.IAdd %0, %0 : i32
     spv.ReturnValue %1 : i32
   }
+
+  // CHECK-LABEL: @use
+  spv.func @use_composite() -> (i32) "None" {
+    // We materialize a `spv._reference_of` op at every use of a
+    // specialization constant in the deserializer. So two ops here.
+    // CHECK: %[[USE1:.*]] = spv._reference_of @scc : !spv.array<2 x i32>
+    // CHECK: %[[ITM0:.*]] = spv.CompositeExtract %[[USE1]][0 : i32] : !spv.array<2 x i32>
+    // CHECK: %[[USE2:.*]] = spv._reference_of @scc : !spv.array<2 x i32>
+    // CHECK: %[[ITM1:.*]] = spv.CompositeExtract %[[USE2]][1 : i32] : !spv.array<2 x i32>
+    // CHECK: spv.IAdd %[[ITM0]], %[[ITM1]]
+
+    %0 = spv._reference_of @scc : !spv.array<2 x i32>
+    %1 = spv.CompositeExtract %0[0 : i32] : !spv.array<2 x i32>
+    %2 = spv.CompositeExtract %0[1 : i32] : !spv.array<2 x i32>
+    %3 = spv.IAdd %1, %2 : i32
+    spv.ReturnValue %3 : i32
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+
+  spv.specConstant @sc_f32_1 = 1.5 : f32
+  spv.specConstant @sc_f32_2 = 2.5 : f32
+  spv.specConstant @sc_f32_3 = 3.5 : f32
+
+  spv.specConstant @sc_i32_1 = 1   : i32
+
+  // CHECK: spv.specConstantComposite @scc_array (@sc_f32_1, @sc_f32_2, @sc_f32_3) : !spv.array<3 x f32>
+  spv.specConstantComposite @scc_array (@sc_f32_1, @sc_f32_2, @sc_f32_3) : !spv.array<3 x f32>
+
+  // CHECK: spv.specConstantComposite @scc_struct (@sc_i32_1, @sc_f32_2, @sc_f32_3) : !spv.struct<i32, f32, f32>
+  spv.specConstantComposite @scc_struct (@sc_i32_1, @sc_f32_2, @sc_f32_3) : !spv.struct<i32, f32, f32>
+
+  // CHECK: spv.specConstantComposite @scc_vector (@sc_f32_1, @sc_f32_2, @sc_f32_3) : vector<3xf32>
+  spv.specConstantComposite @scc_vector (@sc_f32_1, @sc_f32_2, @sc_f32_3) : vector<3 x f32>
 }
 
 // -----
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
 
   spv.specConstant @sc_f32_1 = 1.5 : f32
   spv.specConstant @sc_f32_2 = 2.5 : f32

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

 // spv._reference_of
 //===----------------------------------------------------------------------===//
 
 spv.module Logical GLSL450 {
   spv.specConstant @sc1 = false
   spv.specConstant @sc2 = 42 : i64
   spv.specConstant @sc3 = 1.5 : f32
 
+  spv.specConstantComposite @scc (@sc1, @sc2, @sc3) : !spv.struct<i1, i64, f32>
+
   // CHECK-LABEL: @reference
   spv.func @reference() -> i1 "None" {
     // CHECK: spv._reference_of @sc1 : i1
     %0 = spv._reference_of @sc1 : i1
     spv.ReturnValue %0 : i1
   }
 
+  // CHECK-LABEL: @reference_composite
+  spv.func @reference_composite() -> i1 "None" {
+    // CHECK: spv._reference_of @scc : !spv.struct<i1, i64, f32>
+    %0 = spv._reference_of @scc : !spv.struct<i1, i64, f32>
+    %1 = spv.CompositeExtract %0[0 : i32] : !spv.struct<i1, i64, f32>
+    spv.ReturnValue %1 : i1
+  }
+
   // CHECK-LABEL: @initialize
   spv.func @initialize() -> i64 "None" {
     // CHECK: spv._reference_of @sc2 : i64
     %0 = spv._reference_of @sc2 : i64
     %1 = spv.Variable init(%0) : !spv.ptr<i64, Function>
     %2 = spv.Load "Function" %1 : i64
     spv.ReturnValue %2 : i64
   }

 func @reference_of() {
   // CHECK: spv._reference_of @sc
   %0 = spv._reference_of @sc : i32
   return
 }
 
 // -----
 
+spv.specConstant @sc = 5 : i32
+spv.specConstantComposite @scc (@sc) : !spv.array<1 x i32>
+
+func @reference_of_composite() {
+  // CHECK: spv._reference_of @scc : !spv.array<1 x i32>
+  %0 = spv._reference_of @scc : !spv.array<1 x i32>
+  %1 = spv.CompositeExtract %0[0 : i32] : !spv.array<1 x i32>
+  return
+}
+
+// -----
+
 spv.module Logical GLSL450 {
   spv.func @foo() -> () "None" {
-    // expected-error @+1 {{expected spv.specConstant symbol}}
+    // expected-error @+1 {{expected spv.specConstant or spv.SpecConstantComposite symbol}}
     %0 = spv._reference_of @sc : i32
     spv.Return
   }
 }
 
 // -----
 
 spv.module Logical GLSL450 {
   spv.specConstant @sc = 42 : i32
   spv.func @foo() -> () "None" {
     // expected-error @+1 {{result type mismatch with the referenced specialization constant's type}}
     %0 = spv._reference_of @sc : f32
     spv.Return
   }
 }
 
 // -----
 
+spv.module Logical GLSL450 {
+  spv.specConstant @sc = 42 : i32
+  spv.specConstantComposite @scc (@sc) : !spv.array<1 x i32>
+  spv.func @foo() -> () "None" {
+    // expected-error @+1 {{result type mismatch with the referenced specialization constant's type}}
+    %0 = spv._reference_of @scc : f32
+    spv.Return
+  }
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spv.specConstant
 //===----------------------------------------------------------------------===//
 
 spv.module Logical GLSL450 {
   // CHECK: spv.specConstant @sc1 = false
   spv.specConstant @sc1 = false
   // CHECK: spv.specConstant @sc2 spec_id(5) = 42 : i64