Merge pull request #475 from Xilinx/jrickert.mx6

Allow the import of not-builtin domains and add support for AMD Quarks BFPQuantizeDequantizeOp

Author: jorickert

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

@@ -197,12 +197,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 +214,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.
@@ -1304,18 +1305,26 @@ class FrontendGenImpl {
 
     const int current_opset = current_opset_it->second;
 
+    const auto op_version_map = dialect_op_version_map_.find(node.domain());
+    if (op_version_map == dialect_op_version_map_.end())
+      return "";
+
+    const auto op_opsets_map = dialect_op_opsets_map_.find(node.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:
@@ -1626,13 +1635,14 @@ 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);
+
+    std::string opName = node.op_type() + GetImportVersionOfNode(node);
+    auto domainIt = import_handler_map_.find(node.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);
         }