[mlir][tosa] Add check if the operand of the operations is constant.

[tatwaichong]

Summary

Some uses of TOSA rely on the constant operands of particular operations,
e.g. paddings and pad_const in pad op. Add a verification pattern in the
validation pass, and this is optionally enabled.

Change-Id: I1628c0840a27ab06ef91150eee56ad4f5ac9543d

[tatwaichong]

Define these constant operand checking functions as function template. That might make
code cleaner as it is likely more operations will be added.

[rsuderman]

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

64

Generally we avoid contractions in debug statements

isn't -> is not

67

ditto

78

ditto

101

Not necessary in this CR but longterm we should probably have a table lookup instead of a series of if/else operations. It is excessive right now as we have only two but useful to keep in mind.

103

We should probably just invoke signalPassFailure() within the checkConstantsOperand helper. This way we do not need to repeat the check statements.

[tatwaichong]

Instead of checking a series of if/else operations, adopt operation type look-up method
by implementing polymorphism. I think that may easier to extend to involve more operations
or other checking patterns. And address other review comments.

[tatwaichong]

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

101

Great suggestion. I intend to get rid of if-else-checking manner in this PR as there is a new - fully_connected op check added, and likely more operand checks or other patterns will be added in the future. hopefully, this change can make the code easier to extend.

[rsuderman]

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

40

Typically you don't want to define a default implementation for a virtual function in this case as it is expected to be overrideen.

112

Using the class ConstantOperandCheck for the vector is unnecessary as you are just using it as a wrapper around a function reference. Instead you can using SmallVector<std::function<LogicalResult(Operation*)>>, then emplace the statically defined checkers above. This avoids having an extra virtual function lookup and removes all of the class boilerplate for each check.

[tatwaichong]

Populate and call checkers (function) directly to avoid virtual function lookup and class
overhead.

[rsuderman]

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

40

You don't need a class at all. Just move the population and validation to the Pass itself, and make the individual checks standalone functions.

[tatwaichong]

Rebase onto top-of-tree to include the fix that isn't relevant to this patch.

[tatwaichong]

As suggested, minimize the change by adding the necessary checkers only.

[tatwaichong]

As suggested, minimize the change by adding the necessary checkers only.

[rsuderman]

Awesome. This is so much cleaner.

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td

def TosaValidation : Pass<"tosa-validate", "func::FuncOp"> {
   let description = [{
     This pass validates if input TOSA operations match the specification for given
     criteria, e.g. TOSA profile.
   }];
   let constructor = "createTosaValidationPass()";
 
   let options = [
       Option<"profileName", "profile", "std::string",
              /*default=*/"\"undefined\"",
-      "Validation if ops match for given profile">];
+             "Validate if operations match for the given profile">,
+      Option<"StrictOperationSpecAlignment", "strict-op-spec-alignment", "bool",
+             /*default=*/"false",
+             "Verify if the properties of certain operations align the spec requirement">,
+   ];
 }
 
 #endif // MLIR_DIALECT_TOSA_TRANSFORMS_PASSES

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

 
 using namespace mlir;
 using namespace mlir::tosa;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // TOSA Validation Pass.
 //===----------------------------------------------------------------------===//
 
 struct TosaValidation : public tosa::impl::TosaValidationBase<TosaValidation> {
 public:
   explicit TosaValidation() = default;
+  void runOnOperation() override;
 
 private:
-  void runOnOperation() override;
+  template <typename TosaOp>
+  LogicalResult checkConstantOperand(TosaOp op) {
+    return success();
+  }
 
   std::optional<TosaProfileEnum> profileType;
 };
 
+template <>
+LogicalResult TosaValidation::checkConstantOperand<PadOp>(PadOp op) {
+  DenseElementsAttr paddings;
+  if (!matchPattern(op.getPadding(), m_Constant(&paddings))) {
+    return op.emitOpError("padding of pad isn't constant");
+  }
+
+  DenseElementsAttr pad_const;
+  // Assume this op is zero-padding if pad_const isn't presented.
+  if (op.getPadConst() &&
+      !matchPattern(op.getPadConst(), m_Constant(&pad_const))) {
+    return op.emitOpError("pad_const of pad isn't constant");
+  }
+
+  return success();
+}
+
+template <>
+LogicalResult
+TosaValidation::checkConstantOperand<TransposeOp>(TransposeOp op) {
+  DenseElementsAttr perms;
+  if (!matchPattern(op.getPerms(), m_Constant(&perms))) {
+    return op.emitOpError("perms of transpose isn't constant");
+  }
+
+  return success();
+}
+
 void TosaValidation::runOnOperation() {
   profileType = symbolizeEnum<TosaProfileEnum>(profileName);
 
   getOperation().walk([&](Operation *op) {
     for (Value operand : op->getOperands()) {
       if ((profileType == TosaProfileEnum::BaseInference) &&
           getElementTypeOrSelf(operand).isa<FloatType>()) {
         return signalPassFailure();
       }
       if (getElementTypeOrSelf(operand).isF64()) {
         return signalPassFailure();
       }
     }
+
+    if (StrictOperationSpecAlignment) {
+      // Some uses of TOSA rely on the constant operands of particular
+      // operations.
+      if (auto pad_op = dyn_cast<PadOp>(op))
+        if (failed(checkConstantOperand<PadOp>(pad_op)))
+          signalPassFailure();
+
+      if (auto transpose_op = dyn_cast<TransposeOp>(op))
+        if (failed(checkConstantOperand<TransposeOp>(transpose_op)))
+          signalPassFailure();
+    }
   });
 }
 } // namespace
 
 std::unique_ptr<Pass> mlir::tosa::createTosaValidationPass() {
   return std::make_unique<TosaValidation>();
 }

mlir/test/Dialect/Tosa/invalid.mlir

-// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics --tosa-validate=strict-op-spec-alignment
 
 
 func.func @test_conv2d(%arg0: tensor<1x29x29x4xf32>, %arg1: tensor<16x3x3x4xi8>, %arg2: tensor<16xi8>) -> tensor<1x27x27x16xi8> {
   // expected-error@+1 {{expect both input and weight to be float or not together, got 'f32' and 'i8'}}
   %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>}
            : (tensor<1x29x29x4xf32>, tensor<16x3x3x4xi8>, tensor<16xi8>) -> tensor<1x27x27x16xi8>
   return %0 : tensor<1x27x27x16xi8>
 }

 func.func @test_conv2d(%arg0: tensor<1x29x29x4xi8>, %arg1: tensor<16x3x3x4xi8>, %arg2: tensor<16xi8>) -> tensor<1x27x27x16xi8> {
   // expected-error@+1 {{'tosa.conv2d' op quantizationattr is required for quantized type, and not allowed for float type}}
   %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>}
            : (tensor<1x29x29x4xi8>, tensor<16x3x3x4xi8>, tensor<16xi8>) -> tensor<1x27x27x16xi8>
   return %0 : tensor<1x27x27x16xi8>
 }
 
 
+// -----
+
+func.func @test_pad_non_const(%arg0: tensor<13x21x3xf32>, %arg1: tensor<3x2xi32>) -> tensor<13x21x3xf32> {
+  // expected-error@+1 {{'tosa.pad' op padding of pad isn't constant}}
+  %0 = "tosa.pad"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<3x2xi32>) -> tensor<13x21x3xf32>
+  return %0 : tensor<13x21x3xf32>
+}
+
+// -----
+
+func.func @test_pad_non_const(%arg0: tensor<13x21x3xi8>, %arg1: tensor<i8>) -> tensor<13x21x3xi8> {
+  %0 = "tosa.const"() {value = dense<[[0, 0], [0, 1], [0, 1]]> : tensor<3x2xi32>} : () -> tensor<3x2xi32>
+  // expected-error@+1 {{'tosa.pad' op pad_const of pad isn't constant}}
+  %1 = "tosa.pad"(%arg0, %0, %arg1) : (tensor<13x21x3xi8>, tensor<3x2xi32>, tensor<i8>) -> tensor<13x21x3xi8>
+  return %1 : tensor<13x21x3xi8>
+}
+
+// -----
+
+func.func @test_transpose_non_const(%arg0: tensor<13x21x3xf32>, %arg1: tensor<3xi32>) -> tensor<3x13x21xf32> {
+  // expected-error@+1 {{'tosa.transpose' op perms of transpose isn't constant}}
+  %0 = "tosa.transpose"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<3xi32>) -> tensor<3x13x21xf32>
+  return %0 : tensor<3x13x21xf32>
+}
+