[MLIR][SPIRV] Add initial support for OpSpecConstantOp.

[ergawy]

Summary

This commit adds initial support for SPIR-V OpSpecConstantOp
instruction. The following is introdcued:

• A new spv.specConstantOperation operation consisting of a single

region and of 2 operations within that regions (more details in the
docs of the op itself).

• A new spv.yield instruction that acts a terminator for

spv.specConstantOperation.

For now, the generic form of the new op is supported (i.e. no custom
parsing or printing). This will be done in a follow up patch.

[ergawy]

This is an initial patch to make sure I didn't anything disastrously wrong :D.

[ergawy]

This comment was removed by ergawy.

[ergawy]

Rename to spv.mlir.yield.

[antiagainst]

Nice! Just a few nits.

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

612

Yields the result computed in .. region back to .. op

643

Could you use the symbol SpecConstantOperation? (So capitalize the first letter.) (We have a bunch other ops starting with small cases but we should change them to all capitalized to be consistent.)

643

Wrap around 80 characters

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

3411

We've checked this in the enclosing op. We can drop the check here?

[antiagainst]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

3442

This list is fine for now; but it would be nice that if we can have a following up CL to introduce an op trait, like spirv::UsableInSpecConstantOp and define its counterparts in TableGen so that we can mark those op themselves.

You can find examples of op traits in https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/IR/OpDefinition.h#L684 and the corresonding TableGen definition in https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/IR/OpBase.td#L1749.

[ergawy]

Handle review comments.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

3411

The verifier for SpecCosntantOperation below doesn't have that check actually. I think doing this here will provide a more localized error message.

3442

Added a TODO.

[ergawy]

Rebase.

[ergawy]

• Add IsolatedFromAbove trait to spv.SpecConstantOperation.

[ergawy]

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

646

Please feel free to disagree with adding this trait. I added this to make sure we have tight control on what values get used inside the SpecConstantOperation's region and avoid capturing values from outside that aren't not constants.

Verification logic is simpler and more robust this way, I guess.

[antiagainst]

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

646

Looks good!

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

def SPV_SpecConstantCompositeOp : SPV_Op<"specConstantComposite", [InModuleScope, Symbol]> {
   );
 
   let results = (outs);
 
   let hasOpcode = 0;
 
   let autogenSerialization = 0;
 }
+
+def SPV_YieldOp : SPV_Op<"mlir.yield", [NoSideEffect, Terminator]> {
+  let summary = "Yields the result computed in `spv.SpecConstantOperation`'s"
+                "region back to the parent op.";
+
+  let description = [{
+    This op is a special terminator whose only purpose is to terminate
+    an `spv.SpecConstantOperation`'s enclosed region. It accepts a
+    single operand produced by the preceeding (and only other) instruction
+    in its parent block (see SPV_SpecConstantOperation for further
+    details). This op has no corresponding SPIR-V instruction.
+
+    ```
+    spv.mlir.yield ::= `spv.mlir.yield` ssa-id : spirv-type
+    ```
+
+    #### Example:
+    ```mlir
+    %0 = ... (some op supported by SPIR-V OpSpecConstantOp)
+    spv.mlir.yield %0
+    ```
+  }];
+
+  let arguments = (ins AnyType:$operand);
+
+  let results = (outs);
+
+  let hasOpcode = 0;
+
+  let autogenSerialization = 0;
+
+  let assemblyFormat = "attr-dict $operand `:` type($operand)";
+}
+
+def SPV_SpecConstantOperationOp : SPV_Op<"SpecConstantOperation", [
+    InFunctionScope, NoSideEffect
+  ]> {
+  let summary = "Declare a new specialization constant that results from doing an operation.";
+
+  let description = [{
+    This op declares a SPIR-V specialization constant that results from
+    doing an operation on other constants (specialization or otherwise).
+
+    In the `spv` dialect, this op is modelled as follows:
+
+    ```
+    spv-spec-constant-operation-op ::= `"spv.SpecConstantOperation"`
+                                         `(`ssa-id (`, ` ssa-id)`)`
+                                       `({`
+                                         ssa-id = spirv-op
+                                         `spv.mlir.yield` ssa-id
+                                       `})` `:` function-type
+    ```
+
+    In particular, an `spv.SpecConstantOperation` contains exactly one
+    region. In turn, that region, contains exactly 2 instructions:
+    - One of SPIR-V's instructions that are allowed within an
+    OpSpecConstantOp.
+    - An `spv.mlir.yield` instruction as the terminator.
+
+    The following SPIR-V instructions are valid:
+    - OpSConvert,
+    - OpUConvert,
+    - OpFConvert,
+    - OpSNegate,
+    - OpNot,
+    - OpIAdd,
+    - OpISub,
+    - OpIMul,
+    - OpUDiv,
+    - OpSDiv,
+    - OpUMod,
+    - OpSRem,
+    - OpSMod
+    - OpShiftRightLogical,
+    - OpShiftRightArithmetic,
+    - OpShiftLeftLogical
+    - OpBitwiseOr,
+    - OpBitwiseXor,
+    - OpBitwiseAnd
+    - OpVectorShuffle,
+    - OpCompositeExtract,
+    - OpCompositeInsert
+    - OpLogicalOr,
+    - OpLogicalAnd,
+    - OpLogicalNot,
+    - OpLogicalEqual,
+    - OpLogicalNotEqual
+    - OpSelect
+    - OpIEqual,
+    - OpINotEqual
+    - OpULessThan,
+    - OpSLessThan
+    - OpUGreaterThan,
+    - OpSGreaterThan
+    - OpULessThanEqual,
+    - OpSLessThanEqual
+    - OpUGreaterThanEqual,
+    - OpSGreaterThanEqual
+
+    TODO Add capability-specific ops when supported.
+
+    #### Example:
+    ```mlir
+    %0 = spv.constant 1: i32
+
+    %1 = "spv.SpecConstantOperation"(%0) ({
+      %ret = spv.IAdd %0, %0 : i32
+      spv.mlir.yield %ret : i32
+    }) : (i32) -> i32
+    ```
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$operands);
+
+  let results = (outs Variadic<AnyType>:$results);
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasOpcode = 0;
+
+  let autogenSerialization = 0;
+}
+
 // -----
 
 #endif // SPIRV_STRUCTURE_OPS

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

if (constituentSpecConstOp.default_value().getType() !=
       return constOp.emitError("has incorrect types of operands: expected ")
              << cType.getElementType(index) << ", but provided "
              << constituentSpecConstOp.default_value().getType();
   }
 
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// spv.mlir.yield
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(spirv::YieldOp yieldOp) {
+  Operation *parentOp = yieldOp.getParentOp();
+
+  if (!parentOp || !isa<spirv::SpecConstantOperationOp>(parentOp))
+    return yieldOp.emitOpError(
+        "expected parent op to be 'spv.SpecConstantOperation'");
+
+  Block &block = parentOp->getRegion(0).getBlocks().front();
+  Operation &enclosedOp = block.getOperations().front();
+
+  if (yieldOp.getOperand().getDefiningOp() != &enclosedOp)
+    return yieldOp.emitOpError(
+        "expected operand to be defined by preceeding op");
+
+  return success();
+}
+
+static ParseResult parseSpecConstantOperationOp(OpAsmParser &parser,
+                                                OperationState &state) {
+  // TODO: For now, only generic form is supported.
+  return failure();
+}
+
+static void print(spirv::SpecConstantOperationOp op, OpAsmPrinter &printer) {
+  // TODO
+  printer.printGenericOp(op);
+}
+
+static LogicalResult verify(spirv::SpecConstantOperationOp constOp) {
+  Block &block = constOp.getRegion().getBlocks().front();
+
+  if (block.getOperations().size() != 2)
+    return constOp.emitOpError("expected exactly 2 nested ops");
+
+  Operation &yieldOp = block.getOperations().back();
+
+  if (!isa<spirv::YieldOp>(yieldOp))
+    return constOp.emitOpError("expected terminator to be a yield op");
+
+  Operation &enclosedOp = block.getOperations().front();
+
+  // TODO Add a `UsableInSpecConstantOp` trait and mark ops from the list below
+  // with it instead.
+  if (!isa<spirv::SConvertOp, spirv::UConvertOp, spirv::FConvertOp,
+           spirv::SNegateOp, spirv::NotOp, spirv::IAddOp, spirv::ISubOp,
+           spirv::IMulOp, spirv::UDivOp, spirv::SDivOp, spirv::UModOp,
+           spirv::SRemOp, spirv::SModOp, spirv::ShiftRightLogicalOp,
+           spirv::ShiftRightArithmeticOp, spirv::ShiftLeftLogicalOp,
+           spirv::BitwiseOrOp, spirv::BitwiseXorOp, spirv::BitwiseAndOp,
+           spirv::CompositeExtractOp, spirv::CompositeInsertOp,
+           spirv::LogicalOrOp, spirv::LogicalAndOp, spirv::LogicalNotOp,
+           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 : constOp.getOperands()) {
+    if (!isa<spirv::ConstantOp, spirv::SpecConstantOp,
+             spirv::SpecConstantCompositeOp, spirv::SpecConstantOperationOp>(
+            operand.getDefiningOp()))
+      return constOp.emitOpError("invalid operand");
+  }
+
+  return success();
+}
+
 namespace mlir {
 namespace spirv {
 
 // TableGen'erated operation interfaces for querying versions, extensions, and
 // capabilities.
 #include "mlir/Dialect/SPIRV/SPIRVAvailability.cpp.inc"
 } // namespace spirv
 } // namespace mlir

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

 
 // -----
 
 spv.module Logical GLSL450 {
   spv.specConstant @sc1 = 1.5 : f32
   // expected-error @+1 {{unsupported composite type}}
   spv.specConstantComposite @scc (@sc1) : !spv.coopmatrix<8x16xf32, Device>
 }
+//===----------------------------------------------------------------------===//
+// spv.SpecConstantOperation
+//===----------------------------------------------------------------------===//
+
+// -----
+
+spv.module Logical GLSL450 {
+  spv.func @foo() -> i32 "None" {
+    %0 = spv.constant 1: i32
+
+    %1 = "spv.SpecConstantOperation"(%0) ({
+      %ret = spv.IAdd %0, %0 : i32
+      spv.mlir.yield %ret : i32
+    }) : (i32) -> i32
+
+    spv.ReturnValue %1 : i32
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 {
+  spv.func @foo() -> i32 "None" {
+    %0 = spv.constant 1: i32
+    // expected-error @+1 {{expected parent op to be 'spv.SpecConstantOperation'}}
+    spv.mlir.yield %0 : i32
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 {
+  spv.func @foo() -> i32 "None" {
+    %0 = spv.constant 1: i32
+
+    %1 = "spv.SpecConstantOperation"(%0) ({
+      %ret = spv.ISub %0, %0 : i32
+      // expected-error @+1 {{expected operand to be defined by preceeding op}}
+      spv.mlir.yield %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 {{expected exactly 2 nested ops}}
+    %1 = "spv.SpecConstantOperation"(%0) ({
+      %ret = spv.IAdd %0, %0 : i32
+      %ret2 = spv.IAdd %0, %0 : i32
+      spv.mlir.yield %ret : i32
+    }) : (i32) -> i32
+
+    spv.ReturnValue %1 : i32
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 {
+  spv.func @foo() -> i32 "None" {
+    %0 = spv.constant 1: i32
+
+    // expected-error @+1 {{expected terminator to be a yield op}}
+    %1 = "spv.SpecConstantOperation"(%0) ({
+      %ret = spv.IAdd %0, %0 : i32
+      spv.ReturnValue %ret : i32
+    }) : (i32) -> i32
+
+    spv.ReturnValue %1 : i32
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 {
+  spv.func @foo() -> () "None" {
+    %0 = spv.Variable : !spv.ptr<i32, Function>
+
+    // expected-error @+1 {{invalid enclosed op}}
+    %2 = "spv.SpecConstantOperation"(%0) ({
+      %ret = spv.Load "Function" %0 : i32
+      spv.mlir.yield %ret : i32
+    }) : (!spv.ptr<i32, Function>) -> i32
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 {
+  spv.func @foo() -> () "None" {
+    %0 = spv.Variable : !spv.ptr<i32, Function>
+    %1 = spv.Load "Function" %0 : i32
+
+    // expected-error @+1 {{invalid operand}}
+    %2 = "spv.SpecConstantOperation"(%1) ({
+      %ret = spv.IAdd %1, %1 : i32
+      spv.mlir.yield %ret : i32
+    }) : (i32) -> i32
+  }
+}