Add StatefulOps to TOSA Dialect

This patch adds tosa.variable, tosa.variable.read and
tosa.variable.write operators and tests.

Signed-off-by: Jerry Ge <[email protected]>
Change-Id: I647e2e5c3762d7890b03f6aa7c09a29198b7d355

Author: Jerry-Ge

mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h

@@ -35,8 +35,8 @@ std::unique_ptr<Pass> createTosaToLinalgNamed();
 void addTosaToLinalgPasses(
     OpPassManager &pm, const TosaToLinalgOptions &options,
     // Note: Default to 'none' level unless otherwise specified.
-    tosa::ValidationOptions const &validationOptions =
-        tosa::ValidationOptions().setLevel(tosa::TosaLevelEnum::None));
+    tosa::TosaValidationOptions const &validationOptions = {
+        tosa::TosaProfileEnum::Undefined, false, tosa::TosaLevelEnum::None});
 
 /// Populates conversion passes from TOSA dialect to Linalg dialect.
 void populateTosaToLinalgConversionPatterns(RewritePatternSet *patterns);

mlir/include/mlir/Dialect/Tosa/IR/TosaUtilOps.td

@@ -114,7 +114,7 @@ def Tosa_VariableWriteOp : Tosa_Op<"variable.write", []> {
   }];
 
   let arguments = (ins
-    FlatSymbolRefAttr:$name,
+    SymbolNameAttr:$name,
     AnyType:$value
   );
 
@@ -134,7 +134,7 @@ def Tosa_VariableReadOp : Tosa_Op<"variable.read", []> {
   }];
 
   let arguments = (ins
-    FlatSymbolRefAttr:$name
+    SymbolNameAttr:$name
   );
 
   let results = (outs

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

@@ -68,9 +68,6 @@ struct ValidationOptions {
   }
 };
 
-std::unique_ptr<OperationPass<ModuleOp>> createTosaValidationPass(
-    ValidationOptions const &options = ValidationOptions());
-
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Tosa/Transforms/Passes.h.inc"
 

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

@@ -95,7 +95,6 @@ def TosaValidation : Pass<"tosa-validate", "mlir::ModuleOp"> {
     This pass validates if input TOSA operations match the specification for given
     criteria, e.g. TOSA profile.
   }];
-  let constructor = "createTosaValidationPass()";
 
   let options = [
       Option<"profile", "profile", "mlir::tosa::TosaProfileEnum",

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgPass.cpp

@@ -76,7 +76,7 @@ std::unique_ptr<Pass> mlir::tosa::createTosaToLinalg() {
 
 void mlir::tosa::addTosaToLinalgPasses(
     OpPassManager &pm, const TosaToLinalgOptions &options,
-    tosa::ValidationOptions const &validationOptions) {
+    tosa::TosaValidationOptions const &validationOptions) {
   // Optional decompositions are designed to benefit linalg.
   if (!options.disableTosaDecompositions)
     pm.addNestedPass<func::FuncOp>(tosa::createTosaOptionalDecompositions());
@@ -90,7 +90,6 @@ void mlir::tosa::addTosaToLinalgPasses(
   pm.addNestedPass<func::FuncOp>(tosa::createTosaLayerwiseConstantFoldPass(
       {options.aggressiveReduceConstant}));
   pm.addNestedPass<func::FuncOp>(tosa::createTosaMakeBroadcastablePass());
-  pm.addNestedPass<func::FuncOp>(
-      tosa::createTosaValidationPass(validationOptions));
+  pm.addNestedPass<func::FuncOp>(tosa::createTosaValidation(validationOptions));
   pm.addNestedPass<func::FuncOp>(tosa::createTosaToLinalg());
 }

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

@@ -99,9 +99,10 @@ static constexpr tosa_level_t TOSA_LEVEL_NONE = {0, 0, 0, 0};
 struct TosaValidation : public tosa::impl::TosaValidationBase<TosaValidation> {
 public:
   explicit TosaValidation() { populateConstantOperandChecks(); }
-  explicit TosaValidation(const ValidationOptions &options) : TosaValidation() {
+  explicit TosaValidation(const TosaValidationOptions &options)
+      : TosaValidation() {
     this->profile = options.profile;
-    this->StrictOperationSpecAlignment = options.strictOperationSpecAlignment;
+    this->StrictOperationSpecAlignment = options.StrictOperationSpecAlignment;
     this->level = options.level;
   }
   void runOnOperation() final;
@@ -409,7 +410,7 @@ struct TosaValidation : public tosa::impl::TosaValidationBase<TosaValidation> {
 
   SmallVector<std::function<LogicalResult(Operation *)>> const_checkers;
   tosa_level_t tosa_level;
-  std::unordered_map<std::string, mlir::Type> variables_map;
+  DenseMap<const mlir::StringAttr *, mlir::Type> variables_map;
 };
 
 LogicalResult TosaValidation::applyLevelCheck(Operation *op) {
@@ -445,26 +446,17 @@ inline bool CompatibleTypes(const mlir::Type &type,
 
 bool TosaValidation::CheckVariable(Operation *op) {
   if (isa<mlir::tosa::VariableOp>(op)) {
-    auto name_attr = dyn_cast<mlir::StringAttr>(op->getAttr("name"));
-    if (!name_attr) {
-      op->emitOpError() << "Name attribute is not StringAttr";
-      return false;
-    }
-    std::string name = name_attr.getValue().str();
+    auto name_attr = cast<mlir::StringAttr>(op->getAttr("name"));
 
-    if (variables_map.count(name)) {
+    if (variables_map.count(&name_attr)) {
       op->emitOpError() << "name has already been declared";
       return false;
     }
 
-    auto type_attr = dyn_cast<mlir::TypeAttr>(op->getAttr("type"));
-    if (!type_attr) {
-      op->emitOpError() << "type attribute is not TypeAttr";
-      return false;
-    }
+    auto type_attr = cast<mlir::TypeAttr>(op->getAttr("type"));
     mlir::Type type = type_attr.getValue();
 
-    variables_map[name] = type;
+    variables_map[&name_attr] = type;
   }
 
   return true;
@@ -473,19 +465,14 @@ bool TosaValidation::CheckVariable(Operation *op) {
 bool TosaValidation::CheckVariableReadOrWrite(Operation *op) {
   if (isa<mlir::tosa::VariableReadOp>(op) ||
       isa<mlir::tosa::VariableWriteOp>(op)) {
-    auto name_attr = dyn_cast<mlir::FlatSymbolRefAttr>(op->getAttr("name"));
-    if (!name_attr) {
-      op->emitOpError() << "name attribute is not FlatSymbolRefAttr";
-      return false;
-    }
-    std::string name = name_attr.getValue().str();
+    auto name_attr = cast<mlir::StringAttr>(op->getAttr("name"));
 
-    if (!variables_map.count(name)) {
+    if (!variables_map.count(&name_attr)) {
       op->emitOpError() << "name has not been declared";
       return false;
     }
 
-    auto var_type = variables_map[name];
+    auto var_type = variables_map[&name_attr];
 
     for (auto v : op->getOperands()) {
       auto type = v.getType();
@@ -542,8 +529,3 @@ void TosaValidation::runOnOperation() {
   });
 }
 } // namespace
-
-std::unique_ptr<OperationPass<ModuleOp>>
-mlir::tosa::createTosaValidationPass(ValidationOptions const &options) {
-  return std::make_unique<TosaValidation>(options);
-}