(DRAFT) [mlir][ods] Don't allow trailing optional operands in assembly formats

[Mogball]

Summary

An optional operand at the end of an assembly format creates a format ambiguity with the SSA names of the results of the next operation, e.g.

// test.foo
$optional attr-dict

test.foo
%0 = test.bar

Prevent this by adding a verifier to check for these cases.

[Mogball]

make it compile at least

[lattner]

Cool, great catch!

mlir/include/mlir/Dialect/GPU/GPUOps.td

613

This should be ok: it isn't a trailing argument, it has a trailing type.

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

475

Does the parser correctly disambiguate whether it is parsing replValues or replOperation here? That requires it to look for a colon after .. with (%4 to decode which path it is in

[Mogball]

Basically, I threw this patch up because I don't think, as things stand, there's a great way to detect this ambiguity. A lot of ops use the format $variadicOperands attr-dict require at least 1 operand or have some other invariant that make the format not actually ambiguous. Enforcing this check would make a lot of existing formats uglier than they need to be...

mlir/include/mlir/Dialect/GPU/GPUOps.td

613

Right, but the guard on the optional group is parseOptionalOperand, so if it's empty, the parser will read into the results of the next op.

It's the same situation as you reported here https://github.com/llvm/llvm-project/issues/44860 (why I tagged you)

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

475

Uh, it won't. Good catch.

In reality, however, this operation will have either replValues or replOperation, so the format as it was is not ambiguous, but there's no way to communicate this to ODS.

[lattner]

mlir/include/mlir/Dialect/GPU/GPUOps.td

613

Ah right, got it!

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

475

Right, I was wondering if ODS would catch the ambiguity. :-)

mlir/include/mlir/Dialect/GPU/GPUOps.td

def GPU_PrintfOp : GPU_Op<"printf", [MemoryEffects<[MemWrite]>]>,
   let description = [{
     `gpu.printf` takes a literal format string `format` and an arbitrary number of
     scalar arguments that should be printed.
 
     The format string is a C-style printf string, subject to any restrictions
     imposed by one's target platform.
   }];
   let assemblyFormat = [{
-    $format attr-dict ($args^ `:` type($args))?
+    $format attr-dict (`(` $args^ `:` type($args) `)`)?
   }];
 }
 
 def GPU_ReturnOp : GPU_Op<"return", [HasParent<"GPUFuncOp">, NoSideEffect,
                                      Terminator]>,
     Arguments<(ins Variadic<AnyType>:$operands)>, Results<(outs)> {
   let summary = "Terminator for GPU functions.";
   let description = [{

mlir/include/mlir/Dialect/NVGPU/NVGPU.td

def NVGPU_DeviceAsyncCreateGroupOp : NVGPU_Op<"device_async_create_group", []> {
   This is meant to be used with `gpu.device_async_wait` to synchronize copies
   as explained in those ops descriptions.
 
   Groups are executed in the order they are created.
 
   Example:
 
   ```mlir
-  %0 = gpu.device_async_create_group
+  %0 = gpu.device_async_create_group(%cp0)
   ```
   }];
   let results = (outs NVGPU_DeviceAsyncToken:$asyncToken);
   let arguments = (ins Variadic<NVGPU_DeviceAsyncToken>:$inputTokens);
   let assemblyFormat = [{
-    $inputTokens attr-dict
+    `(` $inputTokens `)` attr-dict
   }];
 }
 
 def NVGPU_DeviceAsyncWaitOp : NVGPU_Op<"device_async_wait", []> {
   let summary = "Wait for async gpu ops to complete.";
   let description = [{
   The `gpu.device_async_wait` op will block the execution thread until the group
   associated with the source token is fully completed.

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

let description = [{
     pdl.replace %root with %otherOp
     ```
   }];
   let arguments = (ins PDL_Operation:$operation,
                        Optional<PDL_Operation>:$replOperation,
                        Variadic<PDL_InstOrRangeOf<PDL_Value>>:$replValues);
   let assemblyFormat = [{
     $operation `with` (`(` $replValues^ `:` type($replValues) `)`)?
-    ($replOperation^)? attr-dict
+    (`(` $replOperation^ `)`)? attr-dict
   }];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::ResultOp
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

let description = [{
     nothing else. They should not exist after a program is fully lowered and
     ready to execute.
 
     Example:
     ```mlir
     %w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Passing
     %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
     %w2 = shape.cstr_eq [1,2], [1,2], [1,2] // Passing
-    %wf = shape.assuming_all %w0, %w1 // Failure
-    %wt = shape.assuming_all %w0, %w2 // Passing
+    %wf = shape.assuming_all(%w0, %w1) // Failure
+    %wt = shape.assuming_all(%w0, %w2) // Passing
     ```
   }];
 
   let arguments = (ins Variadic<Shape_WitnessType>:$inputs);
   let results = (outs Shape_WitnessType:$result);
 
-  let assemblyFormat = "$inputs attr-dict";
+  let assemblyFormat = "`(` $inputs `)` attr-dict";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def Shape_AssumingOp : Shape_Op<"assuming", [
     SingleBlockImplicitTerminator<"AssumingYieldOp">,

mlir/test/mlir-pdll/CodeGen/MLIR/expr.pdll

 }
 
 // -----
 
 // CHECK: pdl.pattern @TupleMemberAccessNumber
 // CHECK: %[[FIRST:.*]] = operation "test.first"
 // CHECK: %[[SECOND:.*]] = operation "test.second"
 // CHECK: rewrite %[[FIRST]] {
-// CHECK:   replace %[[FIRST]] with %[[SECOND]]
+// CHECK:   replace %[[FIRST]] with (%[[SECOND]])
 Pattern TupleMemberAccessNumber {
   let firstOp = op<test.first>;
   let secondOp = op<test.second>(firstOp);
   let tuple = (firstOp, secondOp);
   replace tuple.0 with tuple.1;
 }
 
 // -----
 
 // CHECK: pdl.pattern @TupleMemberAccessName
 // CHECK: %[[FIRST:.*]] = operation "test.first"
 // CHECK: %[[SECOND:.*]] = operation "test.second"
 // CHECK: rewrite %[[FIRST]] {
-// CHECK:   replace %[[FIRST]] with %[[SECOND]]
+// CHECK:   replace %[[FIRST]] with (%[[SECOND]])
 Pattern TupleMemberAccessName {
   let firstOp = op<test.first>;
   let secondOp = op<test.second>(firstOp);
   let tuple = (first = firstOp, second = secondOp);
   replace tuple.first with tuple.second;
 }
 
 // -----

mlir/test/mlir-tblgen/op-format-verify-attribute-colon.td

 // RUN: mlir-tblgen -gen-op-decls -asmformat-error-is-fatal=false -I %S/../../include %s -o=%t 2>&1 | FileCheck %s
 
 include "mlir/IR/OpBase.td"
 
 def TestDialect : Dialect {
   let name = "test";
 }
 class TestFormat_Op<string fmt, list<Trait> traits = []>
     : Op<TestDialect, "format_op", traits> {
   let assemblyFormat = fmt;
 }
 
 //===----------------------------------------------------------------------===//
 // Format ambiguity caused by attribute followed by colon literal
 //===----------------------------------------------------------------------===//
 
 // Test attribute followed by a colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeA : TestFormat_Op<[{
   $attr `:` attr-dict
 }]>, Arguments<(ins AnyAttr:$attr)>;
 
 // Test optional attribute followed by colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeB : TestFormat_Op<[{
   (`foo` $attr^)? `:` attr-dict
 }]>, Arguments<(ins OptionalAttr<AnyAttr>:$attr)>;
 
 // Test attribute followed by whitespace and then colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeC : TestFormat_Op<[{
   $attr ` ` `:` attr-dict
 }]>, Arguments<(ins AnyAttr:$attr)>;
 
 // Test attribute followed by optional dictionary and then colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeD : TestFormat_Op<[{
   $attr attr-dict `:`
 }]>, Arguments<(ins AnyAttr:$attr)>;
 
 // Test attribute followed by optional group and then colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeE : TestFormat_Op<[{
   $attr ($a^)? `:` attr-dict type($a)
 }]>, Arguments<(ins AnyAttr:$attr, Optional<I32>:$a)>;
 
 // Test attribute followed by optional group with literals and then colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeF : TestFormat_Op<[{
   $attr (`(` $a^ `)`)? `:` attr-dict (`(` type($a)^ `)`)?
 }]>, Arguments<(ins AnyAttr:$attr, Optional<I32>:$a)>;
 
 // Test attribute followed by optional group with else group.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeG : TestFormat_Op<[{
   $attr (`(` $a^ `)`) : (`foo`)? `:` attr-dict (`(` type($a)^ `)`)?
 }]>, Arguments<(ins AnyAttr:$attr, Optional<I32>:$a)>;
 
 // Test attribute followed by optional group with colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeH : TestFormat_Op<[{
   $attr (`:` $a^ `)`)? attr-dict (`(` type($a)^ `)`)?
 }]>, Arguments<(ins AnyAttr:$attr, Optional<I32>:$a)>;
 
 // Test attribute followed by optional group with colon in else group.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeI : TestFormat_Op<[{
   $attr (`(` $a^ `)`) : (`:`)? attr-dict (`(` type($a)^ `)`)?
 }]>, Arguments<(ins AnyAttr:$attr, Optional<I32>:$a)>;
 
 // Test attribute followed by two optional groups and then a colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeJ : TestFormat_Op<[{
   $attr (`(` $a^ type($a) `)`) : (`foo`)? ` ` attr-dict (`(` $b^ type($b) `)`)?
   `:`
 }], [AttrSizedOperandSegments]>,
   Arguments<(ins AnyAttr:$attr, Optional<I32>:$a, Optional<I32>:$b)>;
 
 // Test attribute followed by two optional groups and then a colon in the else
 // group.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeK : TestFormat_Op<[{
   $attr (`(` $a^ type($a) `)`) : (`foo`)? ` ` attr-dict
   (`(` $b^ type($b) `)`) : (`:`)?
 }], [AttrSizedOperandSegments]>,
   Arguments<(ins AnyAttr:$attr, Optional<I32>:$a, Optional<I32>:$b)>;
 
 // Test attribute followed by two optional groups with guarded colons but then a
 // colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeL : TestFormat_Op<[{
   $attr (`(` $a^ `:` type($a) `)`) : (`foo` `:`)? ` ` attr-dict
   (`(` $b^ `:` type($b) `)`) : (`foo` `:`)? `:`
 }], [AttrSizedOperandSegments]>,
   Arguments<(ins AnyAttr:$attr, Optional<I32>:$a, Optional<I32>:$b)>;
 
 // Test optional attribute followed by optional groups with a colon along one
 // path.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeM : TestFormat_Op<[{
   (`(` $attr^ ` `)? (`(` $a^ `:` type($a) `)`) : (`foo` `:`)? ` ` attr-dict
   (`(` $b^ `:` type($b) `)`) : (`foo` `:`)? `:`
 }], [AttrSizedOperandSegments]>,
   Arguments<(ins OptionalAttr<AnyAttr>:$attr, Optional<I32>:$a,
                  Optional<I32>:$b)>;
 
 // Test optional attribute followed by optional groups with a colon along one
 // path inside an optional group.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeN : TestFormat_Op<[{
   (`(` $attr^ ` `)? (`(` $a^ `:` type($a) `)`) : (`foo` `:`)? ` ` attr-dict
   (`(` $b^ `:` type($b) `)`) : (`:`)?
 }], [AttrSizedOperandSegments]>,
   Arguments<(ins OptionalAttr<AnyAttr>:$attr, Optional<I32>:$a,
                  Optional<I32>:$b)>;
 
 // Test attribute followed by optional attribute, operand, successor, region,
 // and a colon.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `attr`
 def AmbiguousTypeO : TestFormat_Op<[{
   $attr attr-dict $a $b $c $d $e `:`
 }], [AttrSizedOperandSegments]> {
   let arguments = (ins AnyAttr:$attr, OptionalAttr<I32Attr>:$a,
                        Optional<I32>:$b, Variadic<I32>:$c);
   let successors = (successor VariadicSuccessor<AnySuccessor>:$d);
   let regions = (region VariadicRegion<AnyRegion>:$e);
 }
 
 // Test two attributes, where the second one is ambiguous.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `b`
 def AmbiguousTypeP : TestFormat_Op<[{
   $a attr-dict `(` `:` $b (`:` $c^)?
 }]>, Arguments<(ins AnyAttr:$a, AnyAttr:$b, Optional<I32>:$c)>;
 
 // Test two attributes, where the second one is ambiguous.
 // CHECK: error: format ambiguity caused by `:` literal found after attribute `b`
 def AmbiguousTypeQ : TestFormat_Op<[{
   $a attr-dict (`(` $c^ `:`)? `(` `:` $b `:`
 }]>, Arguments<(ins AnyAttr:$a, AnyAttr:$b, Optional<I32>:$c)>;
 
 // CHECK-NOT: error
 
 // Test attribute followed by two optional groups with guarded colons.
 def ValidTypeA : TestFormat_Op<[{
   $attr (`(` $a^ `:` type($a) `)`) : (`foo` `:`)? ` ` attr-dict
   (`(` $b^ `:` type($b) `)`) : (`foo` `:`)? ` ` `(` `:`
 }], [AttrSizedOperandSegments]>,
   Arguments<(ins AnyAttr:$attr, Optional<I32>:$a, Optional<I32>:$b)>;
 
 // Test optional attribute followed by two optional groups with guarded colons.
 def ValidTypeB : TestFormat_Op<[{
   (`(` $attr^ ` `)? (`(` $a^ `:` type($a) `)`) : (`foo` `:`)? ` ` attr-dict
   (`(` $b^ `:` type($b) `)`) : (`foo` `:`)? ` ` `(` `:`
 }], [AttrSizedOperandSegments]>,
   Arguments<(ins OptionalAttr<AnyAttr>:$attr, Optional<I32>:$a,
                  Optional<I32>:$b)>;
 
 // Test optional attribute guarded colon along within segment.
 def ValidTypeC : TestFormat_Op<[{
   (`(` $attr^ `)`) : (`:`)? attr-dict `:`
 }]>, Arguments<(ins OptionalAttr<AnyAttr>:$attr)>;
 
 // Test optional group guard blocks colon.
 def ValidTypeD : TestFormat_Op<[{
   $a attr-dict ($c^ `:`)?
 }]>, Arguments<(ins AnyAttr:$a, Optional<I32>:$c)>;

mlir/test/mlir-tblgen/op-format-verify-trailing-optional-operand.td

+// RUN: mlir-tblgen -gen-op-decls -asmformat-error-is-fatal=false -I %S/../../include %s -o=%t 2>&1 | FileCheck %s
+
+include "mlir/IR/OpBase.td"
+
+def TestDialect : Dialect {
+  let name = "test";
+}
+class TestFormat_Op<string fmt, list<Trait> traits = []>
+    : Op<TestDialect, "format_op", traits> {
+  let assemblyFormat = fmt;
+}
+
+//===----------------------------------------------------------------------===//
+// Format ambiguity caused by trailing optional operand(s)
+//===----------------------------------------------------------------------===//
+
+// Test trailing optional operand.
+// CHECK: error: format ambiguity caused by trailing
+def AmbiguousTypeA : TestFormat_Op<"attr-dict $a">,
+                     Arguments<(ins Optional<I32>:$a)>;
+
+// Test trailing optional operand followed by optional elements.
+// CHECK: error: format ambiguity caused by trailing
+def AmbiguousTypeB : TestFormat_Op<"$a attr-dict">,
+                     Arguments<(ins Optional<I32>:$a)>;

mlir/tools/mlir-tblgen/OpFormatGen.cpp

private:
 
   /// Verify the state of operation attributes within the format.
   LogicalResult verifyAttributes(SMLoc loc, ArrayRef<FormatElement *> elements);
 
   /// Verify that attributes elements aren't followed by colon literals.
   LogicalResult verifyAttributeColonType(SMLoc loc,
                                          ArrayRef<FormatElement *> elements);
 
+  /// Verify that does not end with an optional operand, which can cause
+  /// ambiguities with the next operation's results.
+  LogicalResult
+  verifyTrailingOptionalOperand(SMLoc loc, ArrayRef<FormatElement *> elements);
+
   /// Verify the state of operation operands within the format.
   LogicalResult
   verifyOperands(SMLoc loc,
                  llvm::StringMap<TypeResolutionInstance> &variableTyResolver);
 
   /// Verify the state of operation regions within the format.
   LogicalResult verifyRegions(SMLoc loc);
 
OpFormatParser::verifyAttributes(SMLoc loc,
                                  ArrayRef<FormatElement *> elements) {
   // Check that there are no `:` literals after an attribute without a constant
   // type. The attribute grammar contains an optional trailing colon type, which
   // can lead to unexpected and generally unintended behavior. Given that, it is
   // better to just error out here instead.
   if (failed(verifyAttributeColonType(loc, elements)))
     return failure();
 
+  // Check for format ambiguities.
+  if (failed(verifyTrailingOptionalOperand(loc, elements)))
+    return failure();
+
   // Check for VariadicOfVariadic variables. The segment attribute of those
   // variables will be infered.
   for (const NamedTypeConstraint *var : seenOperands) {
     if (var->constraint.isVariadicOfVariadic()) {
       fmt.inferredAttributes.insert(
           var->constraint.getVariadicOfVariadicSegmentSizeAttr());
     }
   }

 
 LogicalResult
 OpFormatParser::verifyAttributeColonType(SMLoc loc,
                                          ArrayRef<FormatElement *> elements) {
   return verifyAttributeColon(
       [&](const Twine &msg) { return emitError(loc, msg); }, elements);
 }
 
+static LogicalResult
+verifyTrailingOperands(function_ref<LogicalResult(const Twine &msg)> emitError,
+                       ArrayRef<FormatElement *> elements,
+                       FormatElement *anchor = nullptr) {
+  for (FormatElement *element : llvm::reverse(elements)) {
+    // The anchor element is always required.
+    if (element == anchor)
+      return success();
+
+    // Check for a variable-length operand.
+    if (auto *operandVar = dyn_cast<OperandVariable>(element)) {
+      const NamedTypeConstraint *operand = operandVar->getVar();
+      if (!operand->isVariableLength())
+        return success();
+      return emitError("format ambiguity caused by trailing optional or "
+                       "variadic operand in assembly format");
+    }
+
+    // Recurse on optional groups.
+    if (auto *optional = dyn_cast<OptionalElement>(element)) {
+      if (failed(verifyTrailingOperands(
+              emitError, optional->getThenElements().drop_front(),
+              optional->getAnchor())) ||
+          // The guard element cannot be an optional operand.
+          failed(verifyTrailingOperands(
+              emitError, optional->getThenElements().take_front(1))) ||
+          failed(
+              verifyTrailingOperands(emitError, optional->getElseElements())))
+        return failure();
+    } else if (!isOptionallyParsed(element)) {
+      // Found a non-optionally parsed element that isn't an operand.
+      return success();
+    }
+  }
+  // Iterated back to the start.
+  return success();
+}
+
+LogicalResult OpFormatParser::verifyTrailingOptionalOperand(
+    SMLoc loc, ArrayRef<FormatElement *> elements) {
+  // Terminators can have trailing operands.
+  if (op.getTrait("::mlir::OpTrait::IsTerminator"))
+    return success();
+  return verifyTrailingOperands(
+      [&](const Twine &msg) { return emitError(loc, msg); }, elements);
+}
+
 LogicalResult OpFormatParser::verifyOperands(
     SMLoc loc, llvm::StringMap<TypeResolutionInstance> &variableTyResolver) {
   // Check that all of the operands are within the format, and their types can
   // be inferred.
   auto &buildableTypes = fmt.buildableTypes;
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
     NamedTypeConstraint &operand = op.getOperand(i);