[mlir][RFC] Allow to skip operands type specification in ASM form

[vinograd47]

Summary

Based on the following discussion:
https://llvm.discourse.group/t/declarative-assembly-format-requirement-for-type-presence/4399

Relax checks in OperationParser - it allows to skip value type specification,
if the value was already defined in the same block.

[mehdi_amini]

Leaving aside your syntax issue, I don't understand a use-case where it makes sense in MLIR to have the verifyUsedOnlyInTrivialSSACFG restriction applied on the parent regions transitively. Can you elaborate on this?

[vinograd47]

In D111650#3059024, @mehdi_amini wrote:

Leaving aside your syntax issue, I don't understand a use-case where it makes sense in MLIR to have the verifyUsedOnlyInTrivialSSACFG restriction applied on the parent regions transitively. Can you elaborate on this?

I though about the following case:

^bb2:
  // But the parent of parent - not and we still have forware %0 value declaration 
  scf.if (...) { // It's direct parent region satisfies the restrictions  (single block in region)
    %1 = other.op(%0) // Our Operation with UsedOnlyInTrivialSSACFG trait
  }
  ...

^bb1:
  %0 = some.op() -> tensor<1x2x3x4xf32>
  some.br ^bb2

[rriddle]

(Marking as general -1 for now)

[mehdi_amini]

In D111650#3060617, @vinograd47 wrote:

In D111650#3059024, @mehdi_amini wrote:

Leaving aside your syntax issue, I don't understand a use-case where it makes sense in MLIR to have the verifyUsedOnlyInTrivialSSACFG restriction applied on the parent regions transitively. Can you elaborate on this?

I though about the following case:

^bb2:
  // But the parent of parent - not and we still have forware %0 value declaration 
  scf.if (...) { // It's direct parent region satisfies the restrictions  (single block in region)
    %1 = other.op(%0) // Our Operation with UsedOnlyInTrivialSSACFG trait
  }
  ...

^bb1:
  %0 = some.op() -> tensor<1x2x3x4xf32>
  some.br ^bb2

Right, I understand why is it a problem for your syntax. But I don't think this is a valid reason by itself to add such a trait, so I was wondering if this would make sense in another context / for another use-case?

[vinograd47]

In D111650#3060619, @rriddle wrote:

(Marking as general -1 for now)

@rriddle Could you please be a little more specific about this -1? I've tried to preserve original behavior and restrictions as a default state. But just allow users to explicitly switch off this restrictions in their own projects, if needed.

In D111650#3060694, @mehdi_amini wrote:

Right, I understand why is it a problem for your syntax. But I don't think this is a valid reason by itself to add such a trait, so I was wondering if this would make sense in another context / for another use-case?

I can implement this in another way - just a flag in ODS (bit strictOperandsTypeCheck = 1?). Will it be OK for you?

[rriddle]

In D111650#3060702, @vinograd47 wrote:

In D111650#3060619, @rriddle wrote:

(Marking as general -1 for now)

@rriddle Could you please be a little more specific about this -1? I've tried to preserve original behavior and restrictions as a default state. But just allow users to explicitly switch off this restrictions in their own projects, if needed.

In D111650#3060694, @mehdi_amini wrote:

Right, I understand why is it a problem for your syntax. But I don't think this is a valid reason by itself to add such a trait, so I was wondering if this would make sense in another context / for another use-case?

I can implement this in another way - just a flag in ODS (bit strictOperandsTypeCheck = 1?). Will it be OK for you?

Mehdi is doing a better job than I, but I'm not yet convinced that this complexity is worth it.

[vinograd47]

@rriddle @mehdi_amini

I've redone the change. Now it doesn't touch the ODS generation logic,
instead I've used Custom Directives in assembly format to implement the desired behavior.
The only change in the core part - relaxed check in the OperationParser.

Could you please take a look one more time?

[mehdi_amini]

Maybe you missed my last answer on Discourse, can you take another look there?

[vinograd47]

In D111650#3067285, @mehdi_amini wrote:

Maybe you missed my last answer on Discourse, can you take another look there?

@mehdi_amini Are you talking about the following question?

make it such that the type can’t be elided during printing if the producer isn’t in the same block

I assume I've addressed this in the latest change in the custom directive. So this eliding logic is performed on user side, not in the core part.

[mehdi_amini]

In D111650#3069325, @vinograd47 wrote:

In D111650#3067285, @mehdi_amini wrote:

Maybe you missed my last answer on Discourse, can you take another look there?

@mehdi_amini Are you talking about the following question?

make it such that the type can’t be elided during printing if the producer isn’t in the same block

I assume I've addressed this in the latest change in the custom directive. So this eliding logic is performed on user side, not in the core part.

I don't think that is really what I was saying though: this does not provide a feature to elide the type only if the producer is in the same block. Right now you just poked a hole through the declarative assembly syntax which allows to make it unsafe.

[vinograd47]

Add more checks to OperationParser::resolveSSAUse to handle the misuse of the feature.

[vinograd47]

In D111650#3070737, @mehdi_amini wrote:

I don't think that is really what I was saying though: this does not provide a feature to elide the type only if the producer is in the same block. Right now you just poked a hole through the declarative assembly syntax which allows to make it unsafe.

@mehdi_amini Please take a look one more time. I've added more checks to OperationParser::resolveSSAUse to catch the wrong usage and add extra tests for invalid cases.

[vinograd47]

@rriddle @mehdi_amini Gentle remainder :) Could you please take a look at this patch one more time? I've added more checks to OperationParser::resolveSSAUse to catch the wrong usage and add extra tests for invalid cases.

mlir/lib/Parser/Parser.cpp

auto maybeRecordUse = [&](Value value) {
     if (state.asmState)
       state.asmState->addUses(value, useInfo.loc);
     return value;
   };
 
   // If we have already seen a value of this name, return it.
   if (useInfo.number < entries.size() && entries[useInfo.number].value) {
     Value result = entries[useInfo.number].value;
-    // Check that the type matches the other uses.
-    if (result.getType() == type)
-      return maybeRecordUse(result);
 
+    // Check that the type matches the other uses.
+    if (type && result.getType() != type) {
       emitError(useInfo.loc, "use of value '")
           .append(useInfo.name,
                   "' expects different type than prior uses: ", type, " vs ",
                   result.getType())
           .attachNote(getEncodedSourceLocation(entries[useInfo.number].loc))
           .append("prior use here");
       return nullptr;
     }
 
+    // Explicit type specification can be omitted, only if the Value was already
+    // defined prior in the same block.
+    if (!type && result.getParentBlock() != opBuilder.getBlock()) {
+      emitError(useInfo.loc, "value '")
+          .append(useInfo.name, "' was defined in separate block and requires "
+                                "explicit type definition")
+          .attachNote(getEncodedSourceLocation(entries[useInfo.number].loc))
+          .append("defined here");
+      return nullptr;
+    }
+
+    return maybeRecordUse(result);
+  }
+
+  if (!type) {
+    emitError(useInfo.loc, "forward reference of value '")
+        .append(useInfo.name, "' requires explicit type specification");
+    return nullptr;
+  }
+
   // Make sure we have enough slots for this.
   if (entries.size() <= useInfo.number)
     entries.resize(useInfo.number + 1);
 
   // If the value has already been defined and this is an overly large result
   // number, diagnose that.
   if (entries[0].value && !isForwardRefPlaceholder(entries[0].value))
     return (emitError(useInfo.loc, "reference to invalid result number"),

mlir/test/IR/invalid.mlir

test.graph_region {
     }) : () -> ()
   }
   return
 }
 
 // -----
 
 func @foo() {} // expected-error {{expected non-empty function body}}
+
+// -----
+
+// expected-note@+1 {{defined here}}
+func @foo(%arg0: i64, %arg1: memref<1xf64>) {
+    br ^bb1
+
+  ^bb1:
+    // expected-error@+1 {{value '%arg1' was defined in separate block and requires explicit type definition}}
+    test.format_operand_optional_type_op %arg0, %arg1
+    return
+}
+
+// -----
+
+func @foo() {
+    br ^bb2
+
+  ^bb1:
+    // expected-error@+1 {{forward reference of value '%1' requires explicit type specification}}
+    test.format_operand_optional_type_op %0, %1
+    return
+
+  ^bb2:
+    %0 = arith.constant 0 : i64
+    %1 = memref.alloc() : memref<1xf64>
+    br ^bb1
+}

mlir/test/IR/parser.mlir

"unregistered_func_might_have_graph_region"() ( {
   %1 = "foo"(%1, %2) : (i64, i64) -> i64
   %2 = "bar"(%1) : (i64) -> i64
   "unregistered_terminator"() : () -> ()
 }) {sym_name = "unregistered_op_dominance_violation_ok", type = () -> i1} : () -> ()
 
 // This is an unregister operation, the printing/parsing is handled by the dialect.
 // CHECK: test.dialect_custom_printer custom_format
 test.dialect_custom_printer custom_format
+
+// Can skip type definition for operands, if they are already defined in the same block
+// CHECK-LABEL: func @optional_operand_types
+func @optional_operand_types(%arg0: i64, %arg1: memref<1xf64>) {
+  // CHECK: test.format_operand_optional_type_op %arg0, %arg1
+  test.format_operand_optional_type_op %arg0, %arg1
+  // CHECK: test.format_operand_optional_type_op %arg0, %arg1
+  test.format_operand_optional_type_op %arg0, %arg1 : memref<1xf64>
+  return
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

 #include "TestInterfaces.h"
 #include "TestTypes.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/StringSwitch.h"
 
 // Include this before the using namespace lines below to

 static void print(SingleNoTerminatorCustomAsmOp op, OpAsmPrinter &printer) {
   printer.printRegion(
       op.getRegion(), /*printEntryBlockArgs=*/false,
       // This op has a single block without terminators. But explicitly mark
       // as not printing block terminators for testing.
       /*printBlockTerminators=*/false);
 }
 
+//===----------------------------------------------------------------------===//
+// FormatOperandOptionalTypeOp
+//===----------------------------------------------------------------------===//
+
+static bool isDefinedAbove(Value val, Operation *op) {
+  if (val.isa<BlockArgument>())
+    return true;
+
+  return val.getDefiningOp()->getBlock() == op->getBlock() &&
+         val.getDefiningOp()->isBeforeInBlock(op);
+}
+
+static void printOptionalType(OpAsmPrinter &printer,
+                              FormatOperandOptionalTypeOp op, Type type) {
+  if (isDefinedAbove(op.operand(), op))
+    return;
+
+  printer << ":";
+  printer.printType(type);
+}
+
+static ParseResult parseOptionalType(OpAsmParser &parser, Type &type) {
+  if (parser.parseOptionalColon())
+    return success();
+
+  return parser.parseType(type);
+}
+
 #include "TestOpEnums.cpp.inc"
 #include "TestOpInterfaces.cpp.inc"
 #include "TestOpStructs.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

mlir/test/lib/Dialect/Test/TestOps.td

def FormatOperandCOp : FormatOperandBase<"c", [{
   $buildable `,` $operand `:` type(operands) attr-dict
 }]>;
 def FormatOperandDOp : FormatOperandBase<"d", [{
   $buildable `,` $operand `:` type($operand) attr-dict
 }]>;
 def FormatOperandEOp : FormatOperandBase<"e", [{
   $buildable `,` $operand `:` type($buildable) `,` type($operand) attr-dict
 }]>;
+def FormatOperandOptionalTypeOp : FormatOperandBase<"optional_type", [{
+  $buildable `,` $operand custom<OptionalType>(type($operand)) attr-dict
+}]>;
 
 def FormatSuccessorAOp : TEST_Op<"format_successor_a_op", [Terminator]> {
   let successors = (successor VariadicSuccessor<AnySuccessor>:$targets);
   let assemblyFormat = "$targets attr-dict";
 }
 
 def FormatVariadicOperand : TEST_Op<"format_variadic_operand"> {
   let arguments = (ins Variadic<I64>:$operand);