[mlir][tosa] Add assembly format validation for COND_IF op (#142254)

COND_IF's simplified form - where redundant operand notations are
omitted - is not conformant to the specification. According to the
specification, all operands passed into an operation must be explicitly
declared at each operation's structure. Add optional check to verify if
the given form complies with the specification.

Author: tatwaichong

mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp

@@ -1167,10 +1167,61 @@ bool checkErrorIfPad(Operation *op) {
   return true;
 }
 
+// Returns true if the operation takes no input operands, excluding attributes.
+static bool isNullaryOperation(Operation *op) {
+  if (isa<tosa::ConstOp>(op) || isa<tosa::ConstShapeOp>(op) ||
+      isa<tosa::YieldOp>(op) || isa<tosa::VariableOp>(op))
+    return true;
+  return false;
+}
+
+bool checkErrorIfCondIf(Operation *op) {
+  auto ifOp = dyn_cast<tosa::IfOp>(op);
+  if (!ifOp)
+    return true;
+
+  // Whether the types and shapes of operands between the input/output list and
+  // internal regions are validated by the operation verifier. However, with
+  // support for the simplified form - where redundant operand notations are
+  // omitted - is not conformant to the specification. According to the
+  // specification, all operands passed into an operation must be explicitly
+  // declared at each operation's structure. This code section verify that the
+  // operation's form complies with this requirement.
+
+  // Returns true if the region uses no external input operands.
+  auto isNullaryRegion = [](Region &region) -> bool {
+    bool noLiveInValue = true;
+    region.walk([&noLiveInValue](Operation *op) {
+      if (!isNullaryOperation(op)) {
+        noLiveInValue = false;
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    });
+    return noLiveInValue;
+  };
+
+  mlir::Region &thenGraph = ifOp.getThenGraph();
+  mlir::Region &elseGraph = ifOp.getElseGraph();
+  bool isThenGraphNullaryRegion = isNullaryRegion(thenGraph);
+  bool isElseGraphNullaryRegion = isNullaryRegion(elseGraph);
+  bool isInputListEmpty = ifOp.getInputList().size() == 0;
+
+  if ((isInputListEmpty != isThenGraphNullaryRegion) ||
+      (isInputListEmpty != isElseGraphNullaryRegion)) {
+    op->emitOpError()
+        << "the current simplified form is not strictly conformant to the "
+           "spec, please use the generic format\n";
+    return false;
+  }
+
+  return true;
+}
+
 LogicalResult TosaValidation::applyErrorIfCheck(Operation *op) {
   if (!checkErrorIfResize(op) || !checkErrorIfMul(op) ||
       !checkErrorIfTable(op) || !checkErrorIfRescale(op) ||
-      !checkErrorIfPad(op))
+      !checkErrorIfPad(op) || !checkErrorIfCondIf(op))
     return failure();
   return success();
 }

mlir/test/Dialect/Tosa/error_if_check.mlir

@@ -225,3 +225,17 @@ func.func @test_error_i32_unsigned_output(%arg0: tensor<1xi8>) -> tensor<1xi32>
   %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp {scale32 = false, rounding_mode = "SINGLE_ROUND", per_channel = false, input_unsigned = false, output_unsigned = true} : (tensor<1xi8>, tensor<1xi16>, tensor<1xi8>, tensor<1xi8>, tensor<1xi32>) -> tensor<1xi32>
   return %0 : tensor<1xi32>
 }
+
+// -----
+// CHECK-LABEL: cond_if_simplified_form
+func.func @test_cond_if_simplified_form(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> tensor<f32> {
+  // expected-error@+1 {{'tosa.cond_if' op the current simplified form is not strictly conformant to the spec, please use the generic format}}
+  %0 = tosa.cond_if %arg2 -> (tensor<f32>) {
+    %1 = tosa.add %arg0, %arg1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
+    tosa.yield %1 : tensor<f32>
+  } else {
+    %1 = tosa.sub %arg0, %arg1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
+    tosa.yield %1 : tensor<f32>
+  }
+  return %0 : tensor<f32>
+}