[MLIR][SPIRV] Add (de-)serialization support for SpecConstantOpeation.

[ergawy]

Summary

This commit adds support for (de-)serializing SpecConstantOpeation. One
thing worth noting is that during deserialization, we assign a fake ID to
enclosed ops inside SpecConstantOpeation. We need to do this in order
for deserialization logic to properly update ID to value map and to
later reference the created value from the sibling 'spv::YieldOp'.

[ergawy]

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

3460–3461

That was an implementation error in my previous patch. Apologies for the no-so-thorough testing. Fixed and added a test.

[ergawy]

Add more comments.

[mravishankar]

mlir/lib/Target/SPIRV/Deserialization.cpp

1791

This works fine I guess, but what do you think about just "creating a new stack"

Basically valueMap is the current defined values. You can create an new map and use that to collect values from within the region

DenseMap<uint32_t, Value> newValueMap;
std::swap(valueMap, newValueMap);

<... process the body ...>

std::swap(newValueMap, valueMap)

Sorry if this is totally off here. I might be missing some context (this has been flushed from cache now)

[ergawy]

mlir/lib/Target/SPIRV/Deserialization.cpp

1791

Thanks for your comment and review.

I might have missed what you mean but the problem here is that auto-generated deserialization methods depend on the list of input operands having an ID for the corresponding value in SPIR-V input module. For OpSpecConstantOp, we have a special case: there no SPIR-V ID for the result of the wrapped/enclosed op. So what is represented as a single ID (SSA value) in the SPIR-V module needs to be represented using 2 MLIR SSA values in the spv one:

• The MLIR ID for the result of spv.SpecConstantOperation.
• The MLIR ID for the result of wrapped/enclosed op within its region.

Let me know if I misunderstood.

[ergawy]

Rebase.

[ergawy]

Rebase and bing! :).

[antiagainst]

Thanks for the contribution and sorry about the delay; I was out of office for quite some time. :)

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

3460–3461

No worries. I didn't catch this properly either. :)

mlir/lib/Target/SPIRV/Deserialization.cpp

1791

You can still have the fake ID for the enclosed instruction. We are constructing a transient instruction for the enclosed instruction and invoke the normal instruction processing functionality for it. What Mahesh means is swapping the "context" before processing the transient instruction. The context includes the valueMap and insertion point. So before parsing the transient instruction, if we swap the valueMap with a temporary specOpValueMap, then we are still inserting the fake ID into the specOpValueMap. It's just that the specOpValueMap will only ever contain one fake ID and there is no chance we collide. The specOpValueMap can just live in this particular function. :)

1792

s/spv/SPIR-V/. Similarly for the following occurances.

In comment we prefer the fully-spelled-out name as we don't suffer from symbol limitations as in code. :)

1812

s/materialzied/materialized/

1820

What about first creating the spec op, set the insertion point to its body (with proper guard), and then process the transient instruction? This can avoid us from doing the block manipulation right? Following the above comment, this is part of swapping the "context".

mlir/lib/Target/SPIRV/Serialization.cpp

753

This isn't correct. Spec constant is a kind of constant; spv.SpecConstantOperation also follows in this category. So it should be in typesGlobalValues. This becomes a bit tricky and sorry I didn't catch this earlier. The deserialization will become more complicated because the spv.SpecConstantOperation will be parsed from module scope and we want it to be in function scope. This isn't something entirely new though; it's how spv.constant are handled. They are also parsed in the module scope and materialized in function scope when use. You can see how it's handled in getValue and constantMap. For spv.SpecConstantOperation, we need to keep track of the enclosed opcode and its operands' result IDs. When materializing, use getValue to automatically materialize operand normal/spec constant, and then construct the spv.SpecConstantOperation.

mlir/test/Target/SPIRV/spec-constant.mlir

104

Change one of the operand to be a normal constant?

[ergawy]

Handle review comments:

• Use different value map for emitting SpecConstOperation.
• Small change in testing code.

[ergawy]

mlir/lib/Target/SPIRV/Deserialization.cpp

1791

Ah, I see what Mahesh and you mean. I guess that works since OpSpecConstantOp takes operands whose instructions are materialized while the consuming instruction is being emitted.

1820

If we do this, since the transient instruction has operands that are constants, global vars, or spec constants, then the instructions to materialize these references will be emitted in the spec op's region. This is the main reason I think we need such manipulation here; we first need to materialize all references, if any, in the parent region and then isolate the spec op's enclosed op.

Let me know if what you are saying is going over my head :D.

[antiagainst]

Awesome! I just have a few nits left now. Marked as blocking because we need to avoid change Clang code.

clang-tools-extra/clangd/ClangdLSPServer.cpp

623

Accidental change or rebase leftover?

Anyway we need to revert the change here. :)

mlir/lib/Target/SPIRV/Deserialization.cpp

557

Just define a struct instead of using std::tuple? It will make the code more readable. Also we should use SmallVector instead of ArrayRef here. I guess ArrayRef is okay because we have access to the underlying SPIR-V blob right now; but it can be a potential hazard given ArrayRef does not owns the data.

1811

s/creaed/created/

1816

What about using llvm::SaveAndRestore to avoid manually swap?

1820

You are right. Sorry I left this comment earlier than the comment about serializing into typesGlobalValues.

[ergawy]

Handle review comments.

[ergawy]

clang-tools-extra/clangd/ClangdLSPServer.cpp

623

No idea where that came from. Rebasing made it go away.

[antiagainst]

Awesome, thanks!

clang-tools-extra/clangd/ClangdLSPServer.cpp

llvm::json::Object Result{
                  {"resolveProvider", false},
              }},
             {"hoverProvider", true},
             {"renameProvider", std::move(RenameProvider)},
             {"selectionRangeProvider", true},
             {"documentSymbolProvider", true},
             {"workspaceSymbolProvider", true},
             {"referencesProvider", true},
-            {"astProvider", true},
+            {"astProvider", true}, // clangd extension
             {"executeCommandProvider",
              llvm::json::Object{
                  {"commands",
                   {ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND,
                    ExecuteCommandParams::CLANGD_APPLY_TWEAK}},
              }},
             {"typeHierarchyProvider", true},
-            {"memoryUsageProvider", true}, // clangd extension.
+            {"memoryUsageProvider", true}, // clangd extension
+            {"compilationDatabase",        // clangd extension
+             llvm::json::Object{{"automaticReload", true}}},
             {"callHierarchyProvider", true},
         }}}};
   if (Opts.Encoding)
     Result["offsetEncoding"] = *Opts.Encoding;
   if (Opts.TheiaSemanticHighlighting)
     Result.getObject("capabilities")
         ->insert(
             {"semanticHighlighting",

clang-tools-extra/clangd/test/initialize-params.test

 # RUN: clangd -lit-test < %s | FileCheck -strict-whitespace %s
 # Test initialize request parameters with rootUri
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootUri":"test:///workspace","capabilities":{},"trace":"off"}}
 #      CHECK:  "id": 0,
 # CHECK-NEXT:  "jsonrpc": "2.0",
 # CHECK-NEXT:  "result": {
 # CHECK-NEXT:    "capabilities": {
 # CHECK-NEXT:      "astProvider": true,
 # CHECK-NEXT:      "callHierarchyProvider": true,
 # CHECK-NEXT:      "codeActionProvider": true,
+# CHECK-NEXT:      "compilationDatabase": {
+# CHECK-NEXT:        "automaticReload": true
+# CHECK-NEXT:      },
 # CHECK-NEXT:      "completionProvider": {
 # CHECK-NEXT:        "allCommitCharacters": [
 # CHECK-NEXT:          " ",
 # CHECK-NEXT:          "\t",
 # CHECK-NEXT:          "(",
 # CHECK-NEXT:          ")",
 # CHECK-NEXT:          "[",
 # CHECK-NEXT:          "]",

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

 def SPV_OC_OpConstantFalse             : I32EnumAttrCase<"OpConstantFalse", 42>;
 def SPV_OC_OpConstant                  : I32EnumAttrCase<"OpConstant", 43>;
 def SPV_OC_OpConstantComposite         : I32EnumAttrCase<"OpConstantComposite", 44>;
 def SPV_OC_OpConstantNull              : I32EnumAttrCase<"OpConstantNull", 46>;
 def SPV_OC_OpSpecConstantTrue          : I32EnumAttrCase<"OpSpecConstantTrue", 48>;
 def SPV_OC_OpSpecConstantFalse         : I32EnumAttrCase<"OpSpecConstantFalse", 49>;
 def SPV_OC_OpSpecConstant              : I32EnumAttrCase<"OpSpecConstant", 50>;
 def SPV_OC_OpSpecConstantComposite     : I32EnumAttrCase<"OpSpecConstantComposite", 51>;
+def SPV_OC_OpSpecConstantOperation     : I32EnumAttrCase<"OpSpecConstantOperation", 52>;
 def SPV_OC_OpFunction                  : I32EnumAttrCase<"OpFunction", 54>;
 def SPV_OC_OpFunctionParameter         : I32EnumAttrCase<"OpFunctionParameter", 55>;
 def SPV_OC_OpFunctionEnd               : I32EnumAttrCase<"OpFunctionEnd", 56>;
 def SPV_OC_OpFunctionCall              : I32EnumAttrCase<"OpFunctionCall", 57>;
 def SPV_OC_OpVariable                  : I32EnumAttrCase<"OpVariable", 59>;
 def SPV_OC_OpLoad                      : I32EnumAttrCase<"OpLoad", 61>;
 def SPV_OC_OpStore                     : I32EnumAttrCase<"OpStore", 62>;
 def SPV_OC_OpCopyMemory                : I32EnumAttrCase<"OpCopyMemory", 63>;
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpMemoryModel, SPV_OC_OpEntryPoint, SPV_OC_OpExecutionMode,
       SPV_OC_OpCapability, SPV_OC_OpTypeVoid, SPV_OC_OpTypeBool, SPV_OC_OpTypeInt,
       SPV_OC_OpTypeFloat, SPV_OC_OpTypeVector, SPV_OC_OpTypeMatrix,
       SPV_OC_OpTypeArray, SPV_OC_OpTypeRuntimeArray, SPV_OC_OpTypeStruct,
       SPV_OC_OpTypePointer, SPV_OC_OpTypeFunction, SPV_OC_OpTypeForwardPointer,
       SPV_OC_OpConstantTrue, SPV_OC_OpConstantFalse, SPV_OC_OpConstant,
       SPV_OC_OpConstantComposite, SPV_OC_OpConstantNull, SPV_OC_OpSpecConstantTrue,
       SPV_OC_OpSpecConstantFalse, SPV_OC_OpSpecConstant,
-      SPV_OC_OpSpecConstantComposite, SPV_OC_OpFunction, SPV_OC_OpFunctionParameter,
+      SPV_OC_OpSpecConstantComposite, SPV_OC_OpSpecConstantOperation,
+      SPV_OC_OpFunction, SPV_OC_OpFunctionParameter,
       SPV_OC_OpFunctionEnd, SPV_OC_OpFunctionCall, SPV_OC_OpVariable, SPV_OC_OpLoad,
       SPV_OC_OpStore, SPV_OC_OpCopyMemory, SPV_OC_OpAccessChain, SPV_OC_OpDecorate,
       SPV_OC_OpMemberDecorate, SPV_OC_OpVectorExtractDynamic,
       SPV_OC_OpVectorInsertDynamic, SPV_OC_OpCompositeConstruct,
       SPV_OC_OpCompositeExtract, SPV_OC_OpCompositeInsert, SPV_OC_OpTranspose,
       SPV_OC_OpConvertFToU, SPV_OC_OpConvertFToS, SPV_OC_OpConvertSToF,
       SPV_OC_OpConvertUToF, SPV_OC_OpUConvert, SPV_OC_OpSConvert, SPV_OC_OpFConvert,
       SPV_OC_OpBitcast, SPV_OC_OpSNegate, SPV_OC_OpFNegate, SPV_OC_OpIAdd,

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

if (!isa<spirv::SConvertOp, spirv::UConvertOp, spirv::FConvertOp,
            spirv::LogicalEqualOp, spirv::LogicalNotEqualOp, spirv::SelectOp,
            spirv::IEqualOp, spirv::INotEqualOp, spirv::ULessThanOp,
            spirv::SLessThanOp, spirv::UGreaterThanOp, spirv::SGreaterThanOp,
            spirv::ULessThanEqualOp, spirv::SLessThanEqualOp,
            spirv::UGreaterThanEqualOp, spirv::SGreaterThanEqualOp>(enclosedOp))
     return constOp.emitOpError("invalid enclosed op");
 
   for (auto operand : enclosedOp.getOperands())
-    if (!isa<spirv::ConstantOp, spirv::SpecConstantOp,
-             spirv::SpecConstantCompositeOp, spirv::SpecConstantOperationOp>(
-            operand.getDefiningOp()))
+    if (!isa<spirv::ConstantOp, spirv::ReferenceOfOp,
+             spirv::SpecConstantOperationOp>(operand.getDefiningOp()))
       return constOp.emitOpError(
           "invalid operand, must be defined by a constant operation");
 
   return success();
 }
 
 namespace mlir {
 namespace spirv {

mlir/lib/Target/SPIRV/Deserialization.cpp

 //===- Deserializer.cpp - MLIR SPIR-V Deserialization ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the SPIR-V binary to MLIR SPIR-V module deserialization.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Target/SPIRV/Deserialization.h"
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Target/SPIRV/SPIRVBinaryUtils.h"
private:
   /// Processes OpFunctionEnd and finalizes function. This wires up block
   /// argument created from OpPhi instructions and also structurizes control
   /// flow.
   LogicalResult processFunctionEnd(ArrayRef<uint32_t> operands);
 
   /// Gets the constant's attribute and type associated with the given <id>.
   Optional<std::pair<Attribute, Type>> getConstant(uint32_t id);
 
-  /// Gets the constant's integer attribute with the given <id>. Returns a null
-  /// IntegerAttr if the given is not registered or does not correspond to an
-  /// integer constant.
+  /// Gets the info needed to materialize the spec constant operation op
+  /// associated with the given <id>.
+  Optional<std::tuple<spirv::Opcode, uint32_t, ArrayRef<uint32_t>>>
+  getSpecConstantOperation(uint32_t id);
+
+  /// Gets the constant's integer attribute with the given <id>. Returns a
+  /// null IntegerAttr if the given is not registered or does not correspond
+  /// to an integer constant.
   IntegerAttr getConstantInt(uint32_t id);
 
   /// Returns a symbol to be used for the function name with the given
   /// result <id>. This tries to use the function's OpName if
   /// exists; otherwise creates one based on the <id>.
   std::string getFunctionSymbol(uint32_t id);
 
   /// Returns a symbol to be used for the specialization constant with the given
private:
   /// constant.
   LogicalResult processConstantBool(bool isTrue, ArrayRef<uint32_t> operands,
                                     bool isSpec);
 
   /// Processes a SPIR-V OpConstantComposite instruction with the given
   /// `operands`.
   LogicalResult processConstantComposite(ArrayRef<uint32_t> operands);
 
+  /// Processes a SPIR-V OpSpecConstantComposite instruction with the given
+  /// `operands`.
   LogicalResult processSpecConstantComposite(ArrayRef<uint32_t> operands);
 
+  /// Processes a SPIR-V OpSpecConstantOperation instruction with the given
+  /// `operands`.
+  LogicalResult processSpecConstantOperation(ArrayRef<uint32_t> operands);
+
+  /// Materializes/emits an OpSpecConstantOperation instruction.
+  Value materializeSpecConstantOperation(uint32_t resultID,
+                                         spirv::Opcode enclosedOpcode,
+                                         uint32_t resultTypeID,
+                                         ArrayRef<uint32_t> enclosedOpOperands);
+
   /// Processes a SPIR-V OpConstantNull instruction with the given `operands`.
   LogicalResult processConstantNull(ArrayRef<uint32_t> operands);
 
   //===--------------------------------------------------------------------===//
   // Debug
   //===--------------------------------------------------------------------===//
 
   /// Discontinues any source-level location information that might be active
private:
   DenseMap<uint32_t, std::pair<Attribute, Type>> constantMap;
 
   // 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 info needed to materialize an OpSpecConstantOperation
+  // mapping.
+  DenseMap<uint32_t, std::tuple<spirv::Opcode, uint32_t, ArrayRef<uint32_t>>>
+      specConstOperationMap;
+
   // 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;

 
 Optional<std::pair<Attribute, Type>> Deserializer::getConstant(uint32_t id) {
   auto constIt = constantMap.find(id);
   if (constIt == constantMap.end())
     return llvm::None;
   return constIt->getSecond();
 }
 
+Optional<std::tuple<spirv::Opcode, uint32_t, ArrayRef<uint32_t>>>
+Deserializer::getSpecConstantOperation(uint32_t id) {
+  auto constIt = specConstOperationMap.find(id);
+  if (constIt == specConstOperationMap.end())
+    return llvm::None;
+  return constIt->getSecond();
+}
+
 std::string Deserializer::getFunctionSymbol(uint32_t id) {
   auto funcName = nameMap.lookup(id).str();
   if (funcName.empty()) {
     funcName = "spirv_fn_" + std::to_string(id);
   }
   return funcName;
 }
 
Deserializer::processSpecConstantComposite(ArrayRef<uint32_t> operands) {
   auto op = opBuilder.create<spirv::SpecConstantCompositeOp>(
       unknownLoc, TypeAttr::get(resultType), symName,
       opBuilder.getArrayAttr(elements));
   specConstCompositeMap[resultID] = op;
 
   return success();
 }
 
+LogicalResult
+Deserializer::processSpecConstantOperation(ArrayRef<uint32_t> operands) {
+  if (operands.size() < 3)
+    return emitError(unknownLoc, "OpConstantOperation must have type <id>, "
+                                 "result <id>, and operand opcode");
+
+  uint32_t resultTypeID = operands[0];
+
+  if (!getType(resultTypeID))
+    return emitError(unknownLoc, "undefined result type from <id> ")
+           << resultTypeID;
+
+  uint32_t resultID = operands[1];
+  spirv::Opcode enclosedOpcode = static_cast<spirv::Opcode>(operands[2]);
+  specConstOperationMap.try_emplace(
+      resultID, enclosedOpcode, resultTypeID,
+      ArrayRef<uint32_t>{operands.begin() + 3, operands.end()});
+
+  return success();
+}
+
+Value Deserializer::materializeSpecConstantOperation(
+    uint32_t resultID, spirv::Opcode enclosedOpcode, uint32_t resultTypeID,
+    ArrayRef<uint32_t> enclosedOpOperands) {
+
+  Type resultType = getType(resultTypeID);
+
+  // Instructions wrapped by OpSpecConstantOp need an ID for their
+  // Deserializer::processOp<op_name>(...) to emit the corresponding SPIR-V
+  // dialect wrapped op. For that purpose, a new value map is created and "fake"
+  // ID in that map is assigned to the result of the enclosed instruction. Note
+  // that there is no need to update this fake ID since we only need to
+  // reference the creaed Value for the enclosed op from the spv::YieldOp
+  // created later in this method (both of which are the only values in their
+  // region: the SpecConstantOperation's region). If we encounter another
+  // SpecConstantOperation in the module, we simply re-use the fake ID since the
+  // previous Value assigned to it isn't visible in the current scope anyway.
+  DenseMap<uint32_t, Value> newValueMap;
+  std::swap(valueMap, newValueMap);
+  constexpr uint32_t fakeID = static_cast<uint32_t>(-3);
+
+  SmallVector<uint32_t, 4> enclosedOpResultTypeAndOperands;
+  enclosedOpResultTypeAndOperands.push_back(resultTypeID);
+  enclosedOpResultTypeAndOperands.push_back(fakeID);
+  enclosedOpResultTypeAndOperands.append(enclosedOpOperands.begin(),
+                                         enclosedOpOperands.end());
+
+  // Process enclosed instruction before creating the enclosing
+  // specConstantOperation (and its region). This way, references to constants,
+  // global variables, and spec constants will be materialized outside the new
+  // op's region. For more info, see Deserializer::getValue's implementation.
+  if (failed(
+          processInstruction(enclosedOpcode, enclosedOpResultTypeAndOperands)))
+    return Value();
+
+  // Since the enclosed op is emitted in the current block, split it in a
+  // separate new block.
+  Block *enclosedBlock = curBlock->splitBlock(&curBlock->back());
+
+  auto loc = createFileLineColLoc(opBuilder);
+  auto specConstOperationOp =
+      opBuilder.create<spirv::SpecConstantOperationOp>(loc, resultType);
+
+  Region &body = specConstOperationOp.body();
+  // Move the new block into SpecConstantOperation's body.
+  body.getBlocks().splice(body.end(), curBlock->getParent()->getBlocks(),
+                          Region::iterator(enclosedBlock));
+  Block &block = body.back();
+
+  // RAII guard to reset the insertion point to the module's region after
+  // deserializing the body of the specConstantOperation.
+  OpBuilder::InsertionGuard moduleInsertionGuard(opBuilder);
+  opBuilder.setInsertionPointToEnd(&block);
+
+  opBuilder.create<spirv::YieldOp>(loc, block.front().getResult(0));
+
+  // Switch back to the original "context".
+  std::swap(newValueMap, valueMap);
+
+  return specConstOperationOp.getResult();
+}
+
 LogicalResult Deserializer::processConstantNull(ArrayRef<uint32_t> operands) {
   if (operands.size() != 2) {
     return emitError(unknownLoc,
                      "OpConstantNull must have type <id> and result <id>");
   }
 
   Type resultType = getType(operands[0]);
   if (!resultType) {
if (auto constOp = getSpecConstant(id)) {
     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 specConstOperationInfo = getSpecConstantOperation(id)) {
+    return materializeSpecConstantOperation(
+        id, std::get<0>(*specConstOperationInfo),
+        std::get<1>(*specConstOperationInfo),
+        std::get<2>(*specConstOperationInfo));
+  }
   if (auto undef = getUndefType(id)) {
     return opBuilder.create<spirv::UndefOp>(unknownLoc, undef);
   }
   return valueMap.lookup(id);
 }
 
 LogicalResult
 Deserializer::sliceInstruction(spirv::Opcode &opcode,
LogicalResult Deserializer::processInstruction(spirv::Opcode opcode,
   case spirv::Opcode::OpConstant:
     return processConstant(operands, /*isSpec=*/false);
   case spirv::Opcode::OpSpecConstant:
     return processConstant(operands, /*isSpec=*/true);
   case spirv::Opcode::OpConstantComposite:
     return processConstantComposite(operands);
   case spirv::Opcode::OpSpecConstantComposite:
     return processSpecConstantComposite(operands);
+  case spirv::Opcode::OpSpecConstantOperation:
+    return processSpecConstantOperation(operands);
   case spirv::Opcode::OpConstantTrue:
     return processConstantBool(/*isTrue=*/true, operands, /*isSpec=*/false);
   case spirv::Opcode::OpSpecConstantTrue:
     return processConstantBool(/*isTrue=*/true, operands, /*isSpec=*/true);
   case spirv::Opcode::OpConstantFalse:
     return processConstantBool(/*isTrue=*/false, operands, /*isSpec=*/false);
   case spirv::Opcode::OpSpecConstantFalse:
     return processConstantBool(/*isTrue=*/false, operands, /*isSpec=*/true);

mlir/lib/Target/SPIRV/Serialization.cpp

private:
 
   LogicalResult processConstantOp(spirv::ConstantOp op);
 
   LogicalResult processSpecConstantOp(spirv::SpecConstantOp op);
 
   LogicalResult
   processSpecConstantCompositeOp(spirv::SpecConstantCompositeOp op);
 
+  LogicalResult
+  processSpecConstantOperationOp(spirv::SpecConstantOperationOp op);
+
   /// SPIR-V dialect supports OpUndef using spv.UndefOp that produces a SSA
   /// value to use with other operations. The SPIR-V spec recommends that
   /// OpUndef be generated at module level. The serialization generates an
   /// OpUndef for each type needed at module level.
   LogicalResult processUndefOp(spirv::UndefOp op);
 
   /// Emit OpName for the given `resultID`.
   LogicalResult processName(uint32_t resultID, StringRef name);
Serializer::processSpecConstantCompositeOp(spirv::SpecConstantCompositeOp op) {
 
   encodeInstructionInto(typesGlobalValues,
                         spirv::Opcode::OpSpecConstantComposite, operands);
   specConstIDMap[op.sym_name()] = resultID;
 
   return processName(resultID, op.sym_name());
 }
 
+LogicalResult
+Serializer::processSpecConstantOperationOp(spirv::SpecConstantOperationOp op) {
+  uint32_t typeID = 0;
+  if (failed(processType(op.getLoc(), op.getType(), typeID))) {
+    return failure();
+  }
+
+  auto resultID = getNextID();
+
+  SmallVector<uint32_t, 8> operands;
+  operands.push_back(typeID);
+  operands.push_back(resultID);
+
+  Block &block = op.getRegion().getBlocks().front();
+  Operation &enclosedOp = block.getOperations().front();
+
+  std::string enclosedOpName;
+  llvm::raw_string_ostream rss(enclosedOpName);
+  rss << "Op" << enclosedOp.getName().stripDialect();
+  auto enclosedOpcode = spirv::symbolizeOpcode(rss.str());
+
+  if (!enclosedOpcode) {
+    op.emitError("Couldn't find op code for op ")
+        << enclosedOp.getName().getStringRef();
+    return failure();
+  }
+
+  operands.push_back(static_cast<uint32_t>(enclosedOpcode.getValue()));
+
+  // Append operands to the enclosed op to the list of operands.
+  for (Value operand : enclosedOp.getOperands()) {
+    uint32_t id = getValueID(operand);
+    assert(id && "use before def!");
+    operands.push_back(id);
+  }
+
+  encodeInstructionInto(typesGlobalValues,
+                        spirv::Opcode::OpSpecConstantOperation, operands);
+  valueIDMap[op.getResult()] = resultID;
+
+  return success();
+}
+
 LogicalResult Serializer::processUndefOp(spirv::UndefOp op) {
   auto undefType = op.getType();
   auto &id = undefValIDMap[undefType];
   if (!id) {
     id = getNextID();
     uint32_t typeID = 0;
     if (failed(processType(op.getLoc(), undefType, typeID)) ||
         failed(encodeInstructionInto(typesGlobalValues, spirv::Opcode::OpUndef,
return TypeSwitch<Operation *, LogicalResult>(opInst)
       .Case([&](spirv::LoopOp op) { return processLoopOp(op); })
       .Case([&](spirv::ModuleEndOp) { return success(); })
       .Case([&](spirv::ReferenceOfOp op) { return processReferenceOfOp(op); })
       .Case([&](spirv::SelectionOp op) { return processSelectionOp(op); })
       .Case([&](spirv::SpecConstantOp op) { return processSpecConstantOp(op); })
       .Case([&](spirv::SpecConstantCompositeOp op) {
         return processSpecConstantCompositeOp(op);
       })
+      .Case([&](spirv::SpecConstantOperationOp op) {
+        return processSpecConstantOperationOp(op);
+      })
       .Case([&](spirv::UndefOp op) { return processUndefOp(op); })
       .Case([&](spirv::VariableOp op) { return processVariableOp(op); })
 
       // Then handle all the ops that directly mirror SPIR-V instructions with
       // auto-generated methods.
       .Default(
           [&](Operation *op) { return dispatchToAutogenSerialization(op); });
 }

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

spv.func @foo() -> i32 "None" {
 
     spv.ReturnValue %2 : i32
   }
 }
 
 // -----
 
 spv.module Logical GLSL450 {
+  spv.specConstant @sc = 42 : i32
+
+  spv.func @foo() -> i32 "None" {
+    // CHECK: [[SC:%.*]] = spv.mlir.referenceof @sc
+    %0 = spv.mlir.referenceof @sc : i32
+    // CHECK: spv.SpecConstantOperation wraps "spv.ISub"([[SC]], [[SC]]) : (i32, i32) -> i32
+    %1 = spv.SpecConstantOperation wraps "spv.ISub"(%0, %0) : (i32, i32) -> i32
+    spv.ReturnValue %1 : i32
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 {
   spv.func @foo() -> i32 "None" {
     %0 = spv.constant 1: i32
     // expected-error @+1 {{op expects parent op 'spv.SpecConstantOperation'}}
     spv.mlir.yield %0 : i32
   }
 }
 
 // -----

mlir/test/Target/SPIRV/spec-constant.mlir

spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   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_i32_1 = 1 : i32
+
+  spv.func @use_composite() -> (i32) "None" {
+    // CHECK: [[USE1:%.*]] = spv.mlir.referenceof @sc_i32_1 : i32
+    // CHECK: [[USE2:%.*]] = spv.constant 0 : i32
+
+    // CHECK: [[RES1:%.*]] = spv.SpecConstantOperation wraps "spv.ISub"([[USE1]], [[USE2]]) : (i32, i32) -> i32
+
+    // CHECK: [[USE3:%.*]] = spv.mlir.referenceof @sc_i32_1 : i32
+    // CHECK: [[USE4:%.*]] = spv.constant 0 : i32
+
+    // CHECK: [[RES2:%.*]] = spv.SpecConstantOperation wraps "spv.ISub"([[USE3]], [[USE4]]) : (i32, i32) -> i32
+
+    %0 = spv.mlir.referenceof @sc_i32_1 : i32
+    %1 = spv.constant 0 : i32
+    %2 = spv.SpecConstantOperation wraps "spv.ISub"(%0, %1) : (i32, i32) -> i32
+
+    // CHECK: [[RES3:%.*]] = spv.SpecConstantOperation wraps "spv.IMul"([[RES1]], [[RES2]]) : (i32, i32) -> i32
+    %3 = spv.SpecConstantOperation wraps "spv.IMul"(%2, %2) : (i32, i32) -> i32
+
+    // Make sure deserialization continues from the right place after creating
+    // the previous op.
+    // CHECK: spv.ReturnValue [[RES3]]
+    spv.ReturnValue %3 : i32
+  }
+}