Merge remote-tracking branch 'origin/feature/onnx-to-tosa' into matthias.adjust_for_llvm_bump_to_a58e774f

Author: mgehre-amd

docs/Dialects/onnx.md

@@ -1,4 +1,45 @@
 <!-- Autogenerated by mlir-tblgen; don't manually edit -->
+### `onnx.AMDQuarkBFPQuantizeDequantizeOp` (AMDQuarkBFPQuantizeDequantizeOp)
+
+_BFPQuantizeDequantize_
+
+Block Floating Point (BFP) groups numbers (e.g., tensors, arrays) into blocks, where each block shares a common exponent, and the values in the block are represented with individual mantissas (and the sign bit). This approach offers the performance and speed of 8-bit operations while bringing the precision closer to 16-bit operations.
+
+MicroeXponents (MX) extends the concept of BFP by introducing two levels of exponents: shared exponents for entire blocks and micro exponents for finer-grained sub-blocks. This two-level approach enables more precise scaling of individual elements within a block, reducing quantization error and improving the representational range. The paper https://arxiv.org/abs/2302.08007 introduces three specific formats: MX4, MX6 and MX9, which have different bits of mantissa.
+
+This operator converts floating-point values (typically 32-bit floating-point numbers) into BFP or MX values, then convert them back. It approximates the Quantize-Dequantize process and introduces quantization errors.
+
+Traits: `AlwaysSpeculatableImplTrait`, `OpVersionTrait<1>`, `SameOperandsAndResultElementType`
+
+Interfaces: `ConditionallySpeculatable`, `NoMemoryEffect (MemoryEffectOpInterface)`, `ShapeHelperOpInterface`, `ShapeInferenceOpInterface`
+
+Effects: `MemoryEffects::Effect{}`
+
+#### Attributes:
+
+<table>
+<tr><th>Attribute</th><th>MLIR Type</th><th>Description</th></tr>
+<tr><td><code>bfp_method</code></td><td>::mlir::StringAttr</td><td>string attribute</td></tr>
+<tr><td><code>axis</code></td><td>::mlir::IntegerAttr</td><td>64-bit signed integer attribute</td></tr>
+<tr><td><code>bit_width</code></td><td>::mlir::IntegerAttr</td><td>64-bit signed integer attribute</td></tr>
+<tr><td><code>block_size</code></td><td>::mlir::IntegerAttr</td><td>64-bit signed integer attribute</td></tr>
+<tr><td><code>rounding_mode</code></td><td>::mlir::IntegerAttr</td><td>64-bit signed integer attribute</td></tr>
+<tr><td><code>sub_block_size</code></td><td>::mlir::IntegerAttr</td><td>64-bit signed integer attribute</td></tr>
+<tr><td><code>sub_block_shift_bits</code></td><td>::mlir::IntegerAttr</td><td>64-bit signed integer attribute</td></tr>
+</table>
+
+#### Operands:
+
+| Operand | Description |
+| :-----: | ----------- |
+| `X` | tensor of 32-bit float values
+
+#### Results:
+
+| Result | Description |
+| :----: | ----------- |
+| `Y` | tensor of 32-bit float values
+
 ### `onnx.Abs` (ONNXAbsOp)
 
 _ONNX Abs operation_

docs/ImportONNXDefs.md

@@ -155,3 +155,9 @@ necessary.
 It is not always needed to keep the code for an older version, which may be rewritten into the new
 operation. Thus, we just need to have the dialect definition, but not the code for inference or
 lowering.
+
+# Adding Operations from not-builtin domains
+To add an operation from a not-builtin domain, it needs to be added to the `additional_op_version_dict` in gen_onnx_mlir.py. The key is the domain name and the value is the per-operation version dictionary.
+The new domain also needs to be added to the `domain_abrv_dict` in gen_onnx_mlir.py. The key is the domain name and the value is the abbreviation/prefix used in ONNX-MLIR for this domain.
+For operations from not-builtin domains, the operation definition specification needs to be manually provided.
+This can be done via custom TableGen records for the operations. See [/src/Dialect/ONNX/AdditionalONNXOps.td](../src/Dialect/ONNX/AdditionalONNXOps.td) for examples.

src/Builder/FrontendDialectTransformer.cpp

@@ -68,6 +68,13 @@ bool isDefaultDomain(std::string_view domain) {
   return domain.empty() || (domain == "ai.onnx");
 }
 
+std::string canonicalizeDomain(std::string_view domain) {
+  // Handle aliasing of "ai.onnx" and "". According to the onnx documentation,
+  // the default domain is "ai.onnx", but in practice it seems like the
+  // empty-string domain is used by default.
+  return isDefaultDomain(domain) ? "" : std::string(domain);
+}
+
 /// We consider opset < 6 is old. Users will see a warning if their model
 /// contains ops of old opset.
 constexpr int32_t MINIMUM_SUPPORTED_OPSET = 6;
@@ -86,7 +93,13 @@ template <class T>
 OpsetImportsMap GetOpsetImportsFromProto(const T &proto) {
   OpsetImportsMap opset_imports;
   for (const auto &opset_import : proto.opset_import()) {
-    opset_imports[opset_import.domain()] = opset_import.version();
+    const auto domain = canonicalizeDomain(opset_import.domain());
+    const auto [iter_, inserted] =
+        opset_imports.emplace(domain, opset_import.version());
+    if (!inserted) {
+      llvm::errs() << "Warning: Domain " << domain
+                   << " found multiple times in opset imports.\n";
+    }
   }
   return opset_imports;
 }
@@ -101,7 +114,8 @@ ModelLocalFunctionsMap GetModelLocalFunctions(const onnx::ModelProto &m) {
   for (const auto &function_proto : m.functions()) {
     model_local_functions_by_id.insert(
         {GetModelLocalFunctionsMapIdentifier(
-             function_proto.domain(), function_proto.name()),
+             canonicalizeDomain(function_proto.domain()),
+             function_proto.name()),
             &function_proto});
   }
   return model_local_functions_by_id;
@@ -197,12 +211,12 @@ class FrontendGenImpl {
   ModuleOp module_;
   OpBuilder builder_;
 
-  // onnxop: list of versions supported by onnx-mlir for dialect
-  std::map<std::string, std::vector<int>> op_dialect_version_map_;
-  // onnxop: list of versions for dialect
-  std::map<std::string, std::vector<int>> op_opsets_map_;
-  // onnxop: the top version in third_part/onnx
-  std::map<std::string, int> op_dialect_top_version_map_;
+  // onnxop: list of versions supported by onnx-mlir for dialect, op
+  std::map<std::string, std::map<std::string, std::vector<int>>>
+      dialect_op_version_map_;
+  // onnxop: list of versions for dialect, op
+  std::map<std::string, std::map<std::string, std::vector<int>>>
+      dialect_op_opsets_map_;
 
   // mapping between string name and symbol
   ValueSymbolMapping frontend_symbols_;
@@ -214,7 +228,8 @@ class FrontendGenImpl {
       onnx_mlir::detail::FrontendGenImpl::*)(
       const onnx::NodeProto &, std::string & /*errorMessage*/);
 
-  std::map<std::string, ImportHandlerType> import_handler_map_;
+  std::map<std::string, std::map<std::string, ImportHandlerType>>
+      import_handler_map_;
 
   // The total number of elements in all initializers. This value is a rough
   // counter of the number of parameters in a model.
@@ -848,23 +863,33 @@ class FrontendGenImpl {
     // Use the type map or types in input model to determine the data type of
     // output.
     std::vector<int> outputMap = T::getTypeMap();
+    const bool shouldTakeShapeFromModelForCustomOp =
+        isCustomOp &&
+        (options_.useOnnxModelTypesForCustomOps || givenOutputTypes.empty());
     for (unsigned int i = 0; i < (unsigned int)node.output().size(); i++) {
       // Optional outputs using empty string.
       if (node.output()[i].empty()) {
         outputTypes.emplace_back(builder_.getNoneType());
       } else {
-        if (options_.useOnnxModelTypes ||
-            (isCustomOp && options_.useOnnxModelTypesForCustomOps)) {
+        if (options_.useOnnxModelTypes || shouldTakeShapeFromModelForCustomOp) {
           auto onnxModelType = ConvertOnnxType(node.output(i), errorMessage);
           if (onnxModelType) {
             const auto ec = onnxModelType->getError();
             if (!ec) {
-              outputTypes.emplace_back(*onnxModelType.value());
+              Type outputType = *onnxModelType.value();
+              if (!options_.useOnnxModelTypesForCustomOps &&
+                  !options_.useOnnxModelTypes) {
+                if (auto shapedType = mlir::dyn_cast<ShapedType>(outputType)) {
+                  Type elementType = shapedType.getElementType();
+                  outputType = UnrankedTensorType::get(elementType);
+                }
+              }
+              outputTypes.emplace_back(outputType);
               continue;
             }
             if (!options_.allowMissingOutputTypes || ec != InvalidOnnxFormat) {
               errorMessage +=
-                  "Failed to get type for '" + node.output(i) + "\n";
+                  "Failed to get type for '" + node.output(i) + "'\n";
               return ec;
             }
             llvm::errs() << "Warning: "
@@ -874,13 +899,19 @@ class FrontendGenImpl {
         }
         unsigned int j = i;
         // Variadic output is a single ODS result.
-        if (variadicOut)
+        if (variadicOut) {
           j = 0;
+        }
         if (!givenOutputTypes.empty()) {
+          assert(givenOutputTypes.size() > i &&
+                 "givenOutputTypes size is less than number of outputs");
           outputTypes.emplace_back(
               UnrankedTensorType::get(givenOutputTypes[i]));
         } else if (j < outputMap.size() && outputMap[j] >= MAX_NUM_TYPES) {
           // Mapping gives a connection with an input.
+          assert(
+              outputMap[j] - MAX_NUM_TYPES < static_cast<int>(inputs.size()) &&
+              "output type mapping to input is out of range");
           Type inputType = inputs[outputMap[j] - MAX_NUM_TYPES].getType();
           if (mlir::isa<TensorType>(inputType)) {
             Type elementType =
@@ -1281,25 +1312,34 @@ class FrontendGenImpl {
     return onnx::OpSchemaRegistry::Schema(node.op_type(), version, domain);
   }
 
-  std::string GetImportVersionOfNode(const onnx::NodeProto &node) {
-    auto current_opset_it = opset_map_.find(node.domain());
+  std::string GetImportVersionOfNode(
+      const onnx::NodeProto &node, const std::string &domain) {
+    auto current_opset_it = opset_map_.find(domain);
     if (current_opset_it == opset_map_.end())
       return "";
 
     const int current_opset = current_opset_it->second;
 
+    const auto op_version_map = dialect_op_version_map_.find(domain);
+    if (op_version_map == dialect_op_version_map_.end())
+      return "";
+
+    const auto op_opsets_map = dialect_op_opsets_map_.find(domain);
+    if (op_opsets_map == dialect_op_opsets_map_.end())
+      return "";
+
     LLVM_DEBUG(llvm::dbgs() << DEBUG_TYPE << ": Importing ONNX"
                             << node.op_type() << " (" << node.name() << ")"
                             << ", Opset: " << current_opset << "\n");
 
     const auto supported_opset_list_it =
-        op_dialect_version_map_.find(node.op_type());
-    const auto opset_list_it = op_opsets_map_.find(node.op_type());
+        op_version_map->second.find(node.op_type());
+    const auto opset_list_it = op_opsets_map->second.find(node.op_type());
 
-    // Custom ops may not be present in op_dialect_version_map_. If no version
+    // Custom ops may not be present in op_version_map. If no version
     // info is found, treat as unversioned (no renaming).
-    if (supported_opset_list_it == op_dialect_version_map_.end() ||
-        opset_list_it == op_opsets_map_.end())
+    if (supported_opset_list_it == op_version_map->second.end() ||
+        opset_list_it == op_opsets_map->second.end())
       return "";
 
     // To determine the opset version for a node/op:
@@ -1338,7 +1378,7 @@ class FrontendGenImpl {
 
     // A new opset is added to onnx-mlir when it becomes incompatible.
     // All opset newest than the last opset should use the last opset(version)
-    if (isDefaultDomain(node.domain()) &&
+    if (isDefaultDomain(domain) &&
         upperRangeOfNewestValidOpsetVersion < supported_opset_list.back() &&
         upperRangeOfNewestValidOpsetVersion < MINIMUM_SUPPORTED_OPSET)
       llvm::errs() << "\nWarning: ONNX " << node.op_type()
@@ -1567,13 +1607,12 @@ class FrontendGenImpl {
     auto mlirAttr = builder_.getStringAttr(funcName);
     auto funcAttr = builder_.getNamedAttr("function_name", mlirAttr);
     attributes.push_back(funcAttr);
-    auto domainAttr = builder_.getNamedAttr(
-        "domain_name", builder_.getStringAttr(node.domain()));
+    auto domainAttr = builder_.getNamedAttr("domain_name",
+        builder_.getStringAttr(canonicalizeDomain(node.domain())));
     attributes.push_back(domainAttr);
-    int nIn = 0;
-    int nOut = 0;
+    const int nIn = ONNXCustomOp::getNumberOfOperands();
+    const int nOut = ONNXCustomOp::getNumberOfResults();
     getNodeInputs(node, inputs);
-    nOut = node.output().size();
     std::vector<Type> givenOutputTypes;
 
     // We lack a way of specifying import behavior for custom domains. For now
@@ -1611,13 +1650,16 @@ class FrontendGenImpl {
 
   [[nodiscard]] std::error_code ImportNode(
       const onnx::NodeProto &node, std::string &errorMessage) {
-    if (isDefaultDomain(node.domain()) || (node.domain() == "ai.onnx.ml") ||
-        (node.domain() == "ai.onnx.preview.training")) {
-      std::string opName = node.op_type() + GetImportVersionOfNode(node);
-      auto handler = import_handler_map_.find(opName);
+    const std::string domain = canonicalizeDomain(node.domain());
+
+    const std::string opName =
+        node.op_type() + GetImportVersionOfNode(node, domain);
+    auto domainIt = import_handler_map_.find(domain);
+    if (domainIt != import_handler_map_.end()) {
+      auto handler = domainIt->second.find(opName);
       std::vector<std::string> funcs = options_.functionsToDecompose;
-      if (!(std::find(funcs.begin(), funcs.end(), opName) != funcs.end())) {
-        if (handler != import_handler_map_.end()) {
+      if (!llvm::is_contained(funcs, opName)) {
+        if (handler != domainIt->second.end()) {
           // It's a regular op with a registered handler.
           return (this->*(handler->second))(node, errorMessage);
         }
@@ -1633,7 +1675,7 @@ class FrontendGenImpl {
     }
 
     auto model_function = in_model_functions_.find(
-        GetModelLocalFunctionsMapIdentifier(node.domain(), node.op_type()));
+        GetModelLocalFunctionsMapIdentifier(domain, node.op_type()));
     if (model_function != in_model_functions_.end()) {
       return ImportFunctionCallNode(
           node, /*schema=*/nullptr, model_function->second, errorMessage);