[mlir][spirv] Add parsing and printing support for SpecConstantOperation

[ergawy]

Summary

Adds more support for SpecConstantOperation by defining a custom
syntax for the op and implementing its parsing and printing.

[antiagainst]

This custom assembly format is okay but I think we can further simplify it given we have very rigid structure inside the region. What about something like

%result = spv.SpecConstantOperation wraps "spv.IAdd"(%lhs, %rhs) : (i32, i32) -> i32

?

The spv.mlir.yield can then be implicitly constructed during parsing and omitted during printing.

We have examples for this in TensorFlow: tf_executor.island. See its parsing and printing function here:

• https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc#L357-L371
• https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc#L321-L335

And examples here: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/tests/executor_island_coarsening.mlir#L183-L188

Let me know how do you think. :)

[mravishankar]

Thanks for the patch and sorry for the delay. Thanks @antiagainst for reviewing it. Just a minor comment above, but looks good to me.

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

649–650

Why not add this tag SingleBlockImplicitTerminator<"SPIRV::YieldOp">

[ergawy]

• Add more traits to aid verification.
• Change custom syntax.

[ergawy]

In D92919#2452713, @antiagainst wrote:

This custom assembly format is okay but I think we can further simplify it given we have very rigid structure inside the region. What about something like

%result = spv.SpecConstantOperation wraps "spv.IAdd"(%lhs, %rhs) : (i32, i32) -> i32

?

The spv.mlir.yield can then be implicitly constructed during parsing and omitted during printing.

We have examples for this in TensorFlow: tf_executor.island. See its parsing and printing function here:

• https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc#L357-L371
• https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc#L321-L335

And examples here: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/tests/executor_island_coarsening.mlir#L183-L188

Let me know how do you think. :)

That makes the syntax more concise and almost reduces it to its essence. Thanks for links. Followed a similar approach to TF.

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

649–650

Thanks for bringing my attention to this, didn't know it existed. Add it and HasParent<SpecConstantOperationOp> as well.

[ergawy]

Revert unneeded changes.

[ergawy]

Rebase and erase extra space.

[ergawy]

Probably this hits a CI bug: https://github.com/google/llvm-premerge-checks/issues/268.

[ergawy]

Rebase to see if CI bug is resolved.

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.";
+def SPV_YieldOp : SPV_Op<"mlir.yield", [
+    HasParent<"SpecConstantOperationOp">, 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.
 
def SPV_YieldOp : SPV_Op<"mlir.yield", [
 
   let results = (outs);
 
   let hasOpcode = 0;
 
   let autogenSerialization = 0;
 
   let assemblyFormat = "attr-dict $operand `:` type($operand)";
+
+  let verifier = [{ return success(); }];
 }
 
 def SPV_SpecConstantOperationOp : SPV_Op<"SpecConstantOperation", [
-                                         InFunctionScope, NoSideEffect,
-                                         IsolatedFromAbove]> {
-  let summary = "Declare a new specialization constant that results from doing an operation.";
+       NoSideEffect, InFunctionScope,
+       SingleBlockImplicitTerminator<"YieldOp">]> {
+  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
+    spv-spec-constant-operation-op ::= `spv.SpecConstantOperation` `wraps`
+                                         generic-spirv-op `:` 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.
 
let description = [{
     - OpUGreaterThanEqual,
     - OpSGreaterThanEqual
 
     TODO Add capability-specific ops when supported.
 
     #### Example:
     ```mlir
     %0 = spv.constant 1: i32
+    %1 = spv.constant 1: i32
 
-    %1 = "spv.SpecConstantOperation"(%0) ({
-      %ret = spv.IAdd %0, %0 : i32
-      spv.mlir.yield %ret : i32
-    }) : (i32) -> i32
+    %2 = spv.SpecConstantOperation wraps "spv.IAdd"(%0, %1) : (i32, i32) -> i32
     ```
   }];
 
-  let arguments = (ins Variadic<AnyType>:$operands);
+  let arguments = (ins);
 
-  let results = (outs AnyType:$results);
+  let results = (outs AnyType:$result);
 
   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() !=
              << cType.getElementType(index) << ", but provided "
              << constituentSpecConstOp.default_value().getType();
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// spv.mlir.yield
+// spv.SpecConstantOperation
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(spirv::YieldOp yieldOp) {
-  Operation *parentOp = yieldOp->getParentOp();
+static ParseResult parseSpecConstantOperationOp(OpAsmParser &parser,
+                                                OperationState &state) {
+  Region *body = state.addRegion();
 
-  if (!parentOp || !isa<spirv::SpecConstantOperationOp>(parentOp))
-    return yieldOp.emitOpError(
-        "expected parent op to be 'spv.SpecConstantOperation'");
+  if (parser.parseKeyword("wraps"))
+    return failure();
 
-  Block &block = parentOp->getRegion(0).getBlocks().front();
-  Operation &enclosedOp = block.getOperations().front();
+  body->push_back(new Block);
+  Block &block = body->back();
+  Operation *wrappedOp = parser.parseGenericOperation(&block, block.begin());
 
-  if (yieldOp.getOperand().getDefiningOp() != &enclosedOp)
-    return yieldOp.emitOpError(
-        "expected operand to be defined by preceeding op");
+  if (!wrappedOp)
+    return failure();
 
-  return success();
-}
+  OpBuilder builder(parser.getBuilder().getContext());
+  builder.setInsertionPointToEnd(&block);
+  builder.create<spirv::YieldOp>(wrappedOp->getLoc(), wrappedOp->getResult(0));
+  state.location = wrappedOp->getLoc();
 
-static ParseResult parseSpecConstantOperationOp(OpAsmParser &parser,
-                                                OperationState &state) {
-  // TODO: For now, only generic form is supported.
+  state.addTypes(wrappedOp->getResult(0).getType());
+
+  if (parser.parseOptionalAttrDict(state.attributes))
     return failure();
+
+  return success();
 }
 
 static void print(spirv::SpecConstantOperationOp op, OpAsmPrinter &printer) {
-  // TODO
-  printer.printGenericOp(op);
+  printer << op.getOperationName() << " wraps ";
+  printer.printGenericOp(&op.body().front().front());
 }
 
 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");
 
-  if (enclosedOp.getNumOperands() != constOp.getOperands().size())
-    return constOp.emitOpError("invalid number of operands; expected ")
-           << enclosedOp.getNumOperands() << ", actual "
-           << constOp.getOperands().size();
-
-  if (enclosedOp.getNumOperands() != constOp.getRegion().getNumArguments())
-    return constOp.emitOpError("invalid number of region arguments; expected ")
-           << enclosedOp.getNumOperands() << ", actual "
-           << constOp.getRegion().getNumArguments();
-
-  for (auto operand : constOp.getOperands())
+  for (auto operand : enclosedOp.getOperands())
     if (!isa<spirv::ConstantOp, spirv::SpecConstantOp,
              spirv::SpecConstantCompositeOp, spirv::SpecConstantOperationOp>(
             operand.getDefiningOp()))
-      return constOp.emitOpError("invalid operand");
+      return constOp.emitOpError(
+          "invalid operand, must be defined by a constant operation");
 
   return success();
 }
 
 namespace mlir {
 namespace spirv {
 
 // TableGen'erated operation interfaces for querying versions, extensions, and

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" {
+    // CHECK: [[LHS:%.*]] = spv.constant
     %0 = spv.constant 1: i32
-    %2 = spv.constant 1: i32
-
-    %1 = "spv.SpecConstantOperation"(%0, %0) ({
-    ^bb(%lhs : i32, %rhs : i32):
-      %ret = spv.IAdd %lhs, %rhs : i32
-      spv.mlir.yield %ret : i32
-    }) : (i32, i32) -> i32
-
-    spv.ReturnValue %1 : i32
-  }
-}
-
-// -----
-
-spv.module Logical GLSL450 {
-  spv.func @foo() -> i32 "None" {
-    %0 = spv.constant 1: i32
-    %2 = spv.constant 1: i32
+    // CHECK: [[RHS:%.*]] = spv.constant
+    %1 = spv.constant 1: i32
 
-    // expected-error @+1 {{invalid number of operands; expected 2, actual 1}}
-    %1 = "spv.SpecConstantOperation"(%0) ({
-    ^bb(%lhs : i32, %rhs : i32):
-      %ret = spv.IAdd %lhs, %rhs : i32
-      spv.mlir.yield %ret : i32
-    }) : (i32) -> i32
+    // CHECK: spv.SpecConstantOperation wraps "spv.IAdd"([[LHS]], [[RHS]]) : (i32, i32) -> i32
+    %2 = spv.SpecConstantOperation wraps "spv.IAdd"(%0, %1) : (i32, i32) -> i32
 
-    spv.ReturnValue %1 : i32
+    spv.ReturnValue %2 : i32
   }
 }
 
 // -----
 
 spv.module Logical GLSL450 {
   spv.func @foo() -> i32 "None" {
     %0 = spv.constant 1: i32
-    %2 = spv.constant 1: i32
-
-    // expected-error @+1 {{invalid number of region arguments; expected 2, actual 1}}
-    %1 = "spv.SpecConstantOperation"(%0, %0) ({
-    ^bb(%lhs : i32):
-      %ret = spv.IAdd %lhs, %lhs : i32
-      spv.mlir.yield %ret : i32
-    }) : (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'}}
+    // expected-error @+1 {{op expects parent op '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, %0) ({
-    ^bb(%lhs : i32, %rhs : i32):
-      %ret = spv.ISub %lhs, %rhs : i32
-      // expected-error @+1 {{expected operand to be defined by preceeding op}}
-      spv.mlir.yield %lhs : i32
-    }) : (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, %0) ({
-    ^bb(%lhs : i32, %rhs : i32):
-      %ret = spv.IAdd %lhs, %rhs : i32
-      %ret2 = spv.IAdd %lhs, %rhs : i32
-      spv.mlir.yield %ret : i32
-    }) : (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, %0) ({
-    ^bb(%lhs : i32, %rhs : i32):
-      %ret = spv.IAdd %lhs, %rhs : i32
-      spv.ReturnValue %ret : i32
-    }) : (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) ({
-    ^bb(%arg0 : !spv.ptr<i32, Function>):
-      %ret = spv.Load "Function" %arg0 : i32
-      spv.mlir.yield %ret : i32
-    }) : (!spv.ptr<i32, Function>) -> i32
+    %1 = spv.SpecConstantOperation wraps "spv.Load"(%0) {memory_access = 0 : i32} : (!spv.ptr<i32, Function>) -> i32
+    spv.Return
   }
 }
 
 // -----
 
 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, %1) ({
-    ^bb(%lhs: i32, %rhs: i32):
-      %ret = spv.IAdd %lhs, %lhs : i32
-      spv.mlir.yield %ret : i32
-    }) : (i32, i32) -> i32
+    // expected-error @+1 {{invalid operand, must be defined by a constant operation}}
+    %2 = spv.SpecConstantOperation wraps "spv.IAdd"(%1, %1) : (i32, i32) -> i32
+
+    spv.Return
   }
 }