[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

111

Generally we avoid contractions in debug statements

isn't -> is not

114

ditto

120

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.

122

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

125

ditto

[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

120

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

 } // namespace tosa
 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::tosa;
 
 namespace {
 
+static LogicalResult checkConstantOperandPad(Operation *op) {
+  if (auto pad_op = dyn_cast<tosa::PadOp>(op)) {
+    DenseElementsAttr paddings;
+    if (!matchPattern(pad_op.getPadding(), m_Constant(&paddings)))
+      return op->emitOpError("padding of pad is not constant");
+
+    DenseElementsAttr pad_const;
+    // Assume this op is zero-padding if pad_const is not presented.
+    if (pad_op.getPadConst() &&
+        !matchPattern(pad_op.getPadConst(), m_Constant(&pad_const)))
+      return op->emitOpError("pad_const of pad is not constant");
+  }
+  return success();
+}
+
+static LogicalResult checkConstantOperandTranspose(Operation *op) {
+  if (auto transpose_op = dyn_cast<tosa::TransposeOp>(op)) {
+    DenseElementsAttr perms;
+    if (!matchPattern(transpose_op.getPerms(), m_Constant(&perms)))
+      return op->emitOpError("perms of transpose is not constant");
+  }
+  return success();
+}
+
+static LogicalResult checkConstantOperandFullyConnected(Operation *op) {
+  if (auto fc_op = dyn_cast<tosa::FullyConnectedOp>(op)) {
+    DenseElementsAttr weight;
+    if (!matchPattern(fc_op.getWeight(), m_Constant(&weight)))
+      return op->emitOpError("weight of fully_connected is not constant");
+
+    DenseElementsAttr bias;
+    if (!matchPattern(fc_op.getBias(), m_Constant(&bias)))
+      return op->emitOpError("bias of fully_connected is not constant");
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TOSA Validation Pass.
 //===----------------------------------------------------------------------===//
 
 struct TosaValidation : public tosa::impl::TosaValidationBase<TosaValidation> {
 public:
-  explicit TosaValidation() = default;
+  explicit TosaValidation() { populateConstantOperandChecks(); }
+  void runOnOperation() override;
+
+  LogicalResult applyConstantOperandCheck(Operation *op) {
+    for (auto &checker : const_checkers) {
+      if (failed(checker(op)))
+        return failure();
+    }
+    return success();
+  }
 
 private:
-  void runOnOperation() override;
+  void populateConstantOperandChecks() {
+    const_checkers.emplace_back(checkConstantOperandPad);
+    const_checkers.emplace_back(checkConstantOperandTranspose);
+    const_checkers.emplace_back(checkConstantOperandFullyConnected);
+  }
 
+  SmallVector<std::function<LogicalResult(Operation *)>> const_checkers;
   std::optional<TosaProfileEnum> profileType;
 };
 
 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();
       }
     }
+
+    // Some uses of TOSA rely on the constant operands of particular operations.
+    if (StrictOperationSpecAlignment && failed(applyConstantOperandCheck(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_concat(%arg0 : tensor<2x1xf32>, %arg1 : tensor<2x2xf32>) -> tensor<?x?xf32> {
   // expected-error@+1 {{Cannot concat tensors with different sizes on the non-axis dimension 1}}
   %0 = "tosa.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<2x1xf32>, tensor<2x2xf32>) -> tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
+
+// -----
+
+func.func @test_pad_non_const(%arg0: tensor<13x21x3xf32>, %arg1: tensor<3x2xi32>) -> tensor<13x21x3xf32> {
+  // expected-error@+1 {{'tosa.pad' op padding of pad is not 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 is not 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 is not constant}}
+  %0 = "tosa.transpose"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<3xi32>) -> tensor<3x13x21xf32>
+  return %0 : tensor<3x13x21xf32>
+}
+
+// -----
+
+func.func @test_fully_connected_non_const(%arg0: tensor<13x21x3xf32>, %arg1: tensor<2x3xf32>) -> tensor<273x2xf32> {
+  %0 = "tosa.const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32>
+  %1 = "tosa.reshape"(%arg0) {new_shape = array<i64: 273, 3>} : (tensor<13x21x3xf32>) -> tensor<273x3xf32>
+  // expected-error@+1 {{'tosa.fully_connected' op weight of fully_connected is not constant}}
+  %2 = "tosa.fully_connected"(%1, %arg1, %0) : (tensor<273x3xf32>, tensor<2x3xf32>, tensor<2xf32>) -> tensor<273x2xf32>
+  return %2 : tensor<273x2xf32>
+}
+
+// -----
+
+func.func @test_fully_connected_non_const(%arg0: tensor<13x21x3xf32>, %arg1: tensor<2xf32>) -> tensor<273x2xf32> {
+  %0 = "tosa.const"() {value = dense<[[-0.613216758, -0.63714242, -0.73500061], [0.180762768, 0.773053169, -0.933686495]]> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
+  %1 = "tosa.reshape"(%arg0) {new_shape = array<i64: 273, 3>} : (tensor<13x21x3xf32>) -> tensor<273x3xf32>
+  // expected-error@+1 {{'tosa.fully_connected' op bias of fully_connected is not constant}}
+  %2 = "tosa.fully_connected"(%1, %0, %arg1) : (tensor<273x3xf32>, tensor<2x3xf32>, tensor<2xf32>) -> tensor<273x2xf32>
+  return %2 : tensor<273x2xf32>
+}