Merge branch 'main' into mem_reduction_stickified

Author: imaihal

docs/BuildOnLinuxOSX.md

@@ -15,7 +15,7 @@ Firstly, install MLIR (as a part of LLVM-Project):
 ``` bash
 git clone -n https://github.com/llvm/llvm-project.git
 # Check out a specific branch that is known to work with ONNX-MLIR.
-cd llvm-project && git checkout f142f8afe21bceb00fb495468aa0b5043e98c419 && cd ..
+cd llvm-project && git checkout eaa95a1c2bd38332c1a4e634595f29d22b28ffea && cd ..
 ```
 
 [same-as-file]: <> (utils/build-mlir.sh)

docs/BuildOnWindows.md

@@ -52,7 +52,7 @@ Install MLIR (as a part of LLVM-Project):
 ```shell
 git clone -n https://github.com/llvm/llvm-project.git
 # Check out a specific branch that is known to work with ONNX-MLIR.
-cd llvm-project && git checkout f142f8afe21bceb00fb495468aa0b5043e98c419 && cd ..
+cd llvm-project && git checkout eaa95a1c2bd38332c1a4e634595f29d22b28ffea && cd ..
 ```
 
 [same-as-file]: <> (utils/build-mlir.cmd)

docs/Dialects/onnx.md

@@ -3589,6 +3589,63 @@ where the mean and variance are computed per instance per group of channels, and
 groups `num_groups` should be divisible by the number of channels so that there are
 an equal number of channels per group.
 
+The overall computation has two stages: the first stage normalizes the elements to
+have zero mean and unit variance for each instance in each group, and the second
+stage scales and shifts the results of the first stage. The floating-point precision
+used in the first stage is determined by the `stash_type` attribute. For example,
+if `stash_type` is 1, the operator casts all input variables to 32-bit float,
+performs the computation, and finally casts the normalized results back to the
+original type of `X`. The second stage does not depend on `stash_type`.
+
+When the number of groups is the same as the number of channels, this operator is
+equivalent to InstanceNormalization. When there is only one group, this operator
+is equivalent to LayerNormalization.
+
+Traits: `AlwaysSpeculatableImplTrait`
+
+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>epsilon</code></td><td>::mlir::FloatAttr</td><td>32-bit float attribute</td></tr>
+<tr><td><code>num_groups</code></td><td>::mlir::IntegerAttr</td><td>64-bit signed integer attribute</td></tr>
+<tr><td><code>stash_type</code></td><td>::mlir::IntegerAttr</td><td>64-bit signed integer attribute</td></tr>
+</table>
+
+#### Operands:
+
+| Operand | Description |
+| :-----: | ----------- |
+| `X` | tensor of bfloat16 type values or tensor of 16-bit float values or tensor of 32-bit float values or tensor of 64-bit float values
+| `scale` | tensor of bfloat16 type values or tensor of 16-bit float values or tensor of 32-bit float values or tensor of 64-bit float values
+| `bias` | tensor of bfloat16 type values or tensor of 16-bit float values or tensor of 32-bit float values or tensor of 64-bit float values
+
+#### Results:
+
+| Result | Description |
+| :----: | ----------- |
+| `Y` | tensor of bfloat16 type values or tensor of 16-bit float values or tensor of 32-bit float values or tensor of 64-bit float values
+
+### `onnx.GroupNormalizationV18` (ONNXGroupNormalizationV18Op)
+
+_ONNX GroupNormalization operation_
+
+A GroupNormalization function. Carries out group normalization as described in
+the paper https://arxiv.org/abs/1803.08494
+
+This operator transforms input according to
+```
+y = scale * (x - mean) / sqrt(variance + epsilon) + bias,
+```
+where the mean and variance are computed per instance per group of channels, and
+`scale` and `bias` should be specified for each group of channels. The number of
+groups `num_groups` should be divisible by the number of channels so that there are
+an equal number of channels per group.
+
 When the number of groups is the same as the number of channels, this operator is
 equivalent to InstanceNormalization. When there is only one group, this operator
 is equivalent to LayerNormalization.

docs/SupportedONNXOps-NNPA.md

@@ -3,11 +3,11 @@
 
 # Supported ONNX Operation for Target *NNPA*.
 
-Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitations are listed when applicable. This documentation highlights the minimum and maximum opset versions that are fully supported by onnx-mlir and not the version changes.
+Onnx-mlir currently supports ONNX operations targeting up to opset 21. Limitations are listed when applicable. This documentation highlights the minimum and maximum opset versions that are fully supported by onnx-mlir and not the version changes.
 
 * Operations are defined by the [ONNX Standard](https://github.com/onnx/onnx/blob/main/docs/Operators.md).
 * **Supported Opsets** indicates the lowest and highest opset a model may have for onnx-mlir to support compiling a model with the operator.
-   * A * indicates onnx-mlir is compatible with the latest version of that operator available as of opset 20.
+   * A * indicates onnx-mlir is compatible with the latest version of that operator available as of opset 21.
 
 
 NNPA has hardware limitations in dimension index size and tensor size, which are described in [NNPALimit.hpp](../src/Accelerators/NNPA/Support/NNPALimit.hpp). They are large enough for normal use cases, but if your model exceeds the limitations, CPU is used instead of NNPA.

docs/SupportedONNXOps-cpu.md

@@ -3,11 +3,11 @@
 
 # Supported ONNX Operation for Target *cpu*.
 
-Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitations are listed when applicable. This documentation highlights the minimum and maximum opset versions that are fully supported by onnx-mlir and not the version changes.
+Onnx-mlir currently supports ONNX operations targeting up to opset 21. Limitations are listed when applicable. This documentation highlights the minimum and maximum opset versions that are fully supported by onnx-mlir and not the version changes.
 
 * Operations are defined by the [ONNX Standard](https://github.com/onnx/onnx/blob/main/docs/Operators.md).
 * **Supported Opsets** indicates the lowest and highest opset a model may have for onnx-mlir to support compiling a model with the operator.
-   * A * indicates onnx-mlir is compatible with the latest version of that operator available as of opset 20.
+   * A * indicates onnx-mlir is compatible with the latest version of that operator available as of opset 21.
 
 
 | Op |Supported Opsets (inclusive) |Limitations |Notes |
@@ -36,8 +36,8 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **BitwiseOr** |18 - * | | |
 | **BitwiseXor** |18 - * | | |
 | **BlackmanWindow** |none | | | |
-| **Cast** |6 - * |Cast only between float and double types. Only ppc64le and MacOS platforms support float16. | |
-| **CastLike** |19 - * |CastLike only between float and double types. Only ppc64le and MacOS platforms support float16. | |
+| **Cast** |6 - * |Cast only between float and double types. Only ppc64le and MacOS platforms support float16. Does not support int4 and uint4. | |
+| **CastLike** |19 - * |CastLike only between float and double types. Only ppc64le and MacOS platforms support float16. Does not support int4 and uint4. | |
 | **CastMap** |none | | | |
 | **CategoryMapper** |none | | | |
 | **Ceil** |6 - * | | |
@@ -48,8 +48,8 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **Compress** |9 - * | | |
 | **Concat** |6 - * | | |
 | **ConcatFromSequence** |none | | | |
-| **Constant** |6 - * | | |
-| **ConstantOfShape** |9 - * | | |
+| **Constant** |6 - * |Does not support int4 and uint4. | |
+| **ConstantOfShape** |9 - * |Does not support int4 and uint4. | |
 | **Conv** |6 - * | | |
 | **ConvInteger** |none | | | |
 | **ConvTranspose** |6 - * |Spatial dimensions (H and W in input `X`, and kH and kW in input `W`) must be static dimension. | |
@@ -59,7 +59,7 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **DFT** |17 - * | | |
 | **DeformConv** |none | | | |
 | **DepthToSpace** |13 - * | | |
-| **DequantizeLinear** |10 - * |Only support for per-tensor or layer dequantization. No support for per-axis dequantization. | |
+| **DequantizeLinear** |10 - * |Only support for per-tensor or layer dequantization. No support for per-axis dequantization. Does not support int4 and uint4. | |
 | **Det** |none | | | |
 | **DictVectorizer** |none | | | |
 | **Div** |6 - * |No support for short integers. | |
@@ -73,7 +73,7 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **Expand** |8 - * |Input `shape` must have static shape. | |
 | **EyeLike** |none | | | |
 | **FeatureVectorizer** |none | | | |
-| **Flatten** |6 - * | | |
+| **Flatten** |6 - * |Does not support int4 and uint4. | |
 | **Floor** |6 - * | | |
 | **GRU** |7 - * |W, B and R must be constants. | |
 | **Gather** |6 - * | | |
@@ -94,8 +94,8 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **HardSigmoid** |6 - * | | |
 | **HardSwish** |none | | | |
 | **Hardmax** |6 - * | | |
-| **Identity** |16 - * |Sequence identity not supported. | |
-| **If** |16 - * |Sequence and Optional outputs are not supported. | |
+| **Identity** |16 - * |Sequence identity not supported. Does not support int4 and uint4. | |
+| **If** |16 - * |Sequence and Optional outputs are not supported. Does not support int4 and uint4. | |
 | **Imputer** |none | | | |
 | **InstanceNormalization** |6 - * | | |
 | **IsInf** |20 - * |Currently no support for float16 infinity value. Only for float32 and float64. | |
@@ -111,7 +111,7 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **LinearRegressor** |none | | | |
 | **Log** |6 - * | | |
 | **LogSoftmax** |13 - * |Axis 0, 1, and default currently disabled due to changes in ONNX 1.8.1/Opset 13. |Temporally removed due to changes in onnx 1.8.1. |
-| **Loop** |6 - * |Input must have static shape. | |
+| **Loop** |6 - * |Input must have static shape. Does not support int4 and uint4. | |
 | **LpNormalization** |none | | | |
 | **LpPool** |none | | | |
 | **MatMul** |6 - * | | |
@@ -142,11 +142,11 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **OptionalHasElement** |none | | | |
 | **Or** |7 - * | | |
 | **PRelu** |6 - * | | |
-| **Pad** |6 - * |axes input not supported. | |
+| **Pad** |6 - * |axes input not supported. Does not support int4 and uint4. | |
 | **Pow** |7 - * |No support for power with integer types. | |
 | **QLinearConv** |none | | | |
 | **QLinearMatMul** |none | | | |
-| **QuantizeLinear** |10 - * |Do not support per-axis and i8 quantization. | |
+| **QuantizeLinear** |10 - * |Does not support per-axis and i8 quantization. Does not support int4 and uint4. | |
 | **RNN** |7 - * |W, B and R must be constants. | |
 | **RandomNormal** |none | | | |
 | **RandomNormalLike** |none | | | |
@@ -165,7 +165,7 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **ReduceSum** |6 - * |Default axis and do_not_keep_dim not supported. |Default axis and do_not_keep_dim temporarily removed due to changes in onnx 1.8.1. |
 | **ReduceSumSquare** |13 - * |Default axis and do_not_keep_dim not supported. | |
 | **Relu** |6 - * | | |
-| **Reshape** |6 - * |allowzero not supported. Input `shape` must have static dimension. | |
+| **Reshape** |6 - * |allowzero not supported. Input `shape` must have static dimension. Does not support int4 and uint4. | |
 | **Resize** |10 - * |Missing support for linear, cubic, crop, pytorch_half_pixel, and floor. Attributes antialias, axes and keep_aspect_ratio_policy are not supported. `scales` and `sizes` must have static dimension. | |
 | **ReverseSequence** |10 - * | | |
 | **RoiAlign** |none | | | |
@@ -174,7 +174,7 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **SVMClassifier** |none | | | |
 | **SVMRegressor** |none | | | |
 | **Scaler** |none | | | |
-| **Scan** |8 - * |Does not support dynamic shapes. |Precision issue with newer opset, maybe just unsupported. Dynamic shape?. |
+| **Scan** |8 - * |Does not support dynamic shapes. Does not support int4 and uint4. |Precision issue with newer opset, maybe just unsupported. Dynamic shape?. |
 | **Scatter** |none | | | |
 | **ScatterElements** |11 - * |Does not support duplicate indices. | |
 | **ScatterND** |11 - * |Does not support scatternd add/multiply. | |
@@ -186,13 +186,13 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **SequenceInsert** |11 - * |Does not support unranked sequence element. | |
 | **SequenceLength** |none | | | |
 | **SequenceMap** |none | | | |
-| **Shape** |15 - * |Does not support start and end attributes. | |
+| **Shape** |15 - * |Does not support start and end attributes. Does not support int4 and uint4. | |
 | **Shrink** |none | | | |
 | **Sigmoid** |6 - * | | |
 | **Sign** |9 - * | | |
 | **Sin** |7 - * | | |
 | **Sinh** |9 - * | | |
-| **Size** |13 - * | | |
+| **Size** |13 - * |Does not support int4 and uint4. | |
 | **Slice** |13 - * |Axis must be a constant argument. |Add tests to slices, currently have none. |
 | **Softmax** |6 - * | | |
 | **SoftmaxCrossEntropyLoss** |none | | | |
@@ -202,7 +202,7 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **Split** |6 - * |Does not support static and dynamic shape, zero size splits. |Temporally removed due to changes in onnx 1.8.1. |
 | **SplitToSequence** |none | | | |
 | **Sqrt** |6 - * | | |
-| **Squeeze** |6 - * |Does not support static and dynamic shape. |Temporally removed due to changes in onnx 1.8.1. |
+| **Squeeze** |6 - * |Does not support static and dynamic shape. Does not support int4 and uint4. |Temporally removed due to changes in onnx 1.8.1. |
 | **StringNormalizer** |none | | | |
 | **Sub** |6 - * |Does not support short integers. | |
 | **Sum** |6 - * | | |
@@ -212,12 +212,12 @@ Onnx-mlir currently supports ONNX operations targeting up to opset 20. Limitatio
 | **ThresholdedRelu** |none | | | |
 | **Tile** |6 - * | | |
 | **TopK** |10 - * |`K`, the number of top elements to retrieve, must have static shape. | |
-| **Transpose** |6 - * | | |
+| **Transpose** |6 - * |Does not support int4 and uint4. | |
 | **TreeEnsembleClassifier** |none | | | |
 | **TreeEnsembleRegressor** |none | | | |
 | **Trilu** |14 - * | | |
 | **Unique** |11 - * | | |
-| **Unsqueeze** |6 - * |Does not support static and dynamic shape. |Temporally removed due to changes in onnx 1.8.1. |
+| **Unsqueeze** |6 - * |Does not support static and dynamic shape. Does not support int4 and uint4. |Temporally removed due to changes in onnx 1.8.1. |
 | **Upsample** |7 - * |Input `X` and `Y` must have static shape. | |
 | **Where** |9 - * | | |
 | **Xor** |7 - * | | |

requirements.txt

@@ -1,7 +1,8 @@
 lit~=15.0
 # numpy 1.24 deprecates np.object, np.bool, np.float, np.complex, np.str,
 # and np.int which are used heavily in onnx-mlir.
-numpy~=1.22.2, <=1.23.5
+numpy==2.0.1
+onnx==1.16.2
 protobuf==4.21.12
-pytest~=7.2
-pytest-xdist~=3.0
+pytest==8.3.2
+pytest-xdist==3.6.1

src/Builder/OpBuildTable.inc

@@ -80,7 +80,7 @@ op_dialect_version_map_["Gradient"] = {1};
 op_dialect_version_map_["Greater"] = {13};
 op_dialect_version_map_["GreaterOrEqual"] = {16};
 op_dialect_version_map_["GridSample"] = {16};
-op_dialect_version_map_["GroupNormalization"] = {18};
+op_dialect_version_map_["GroupNormalization"] = {21, 18};
 op_dialect_version_map_["HammingWindow"] = {17};
 op_dialect_version_map_["HannWindow"] = {17};
 op_dialect_version_map_["HardSigmoid"] = {6};
@@ -358,6 +358,8 @@ import_handler_map_["GridSample"] =
    &onnx_mlir::detail::FrontendGenImpl::buildOperation<mlir::ONNXGridSampleOp>;
 import_handler_map_["GroupNormalization"] = 
    &onnx_mlir::detail::FrontendGenImpl::buildOperation<mlir::ONNXGroupNormalizationOp>;
+import_handler_map_["GroupNormalizationV18"] = 
+   &onnx_mlir::detail::FrontendGenImpl::buildOperation<mlir::ONNXGroupNormalizationV18Op>;
 import_handler_map_["HammingWindow"] = 
    &onnx_mlir::detail::FrontendGenImpl::buildOperation<mlir::ONNXHammingWindowOp>;
 import_handler_map_["HannWindow"] = 

src/Compiler/CompilerDialects.cpp

@@ -46,8 +46,9 @@ DialectRegistry registerDialects(ArrayRef<accel::Accelerator::Kind> accels) {
   for (auto *accel : accel::Accelerator::getAccelerators())
     accel->registerDialects(registry);
 
-  if (useOldBufferization)
-    memref::registerAllocationOpInterfaceExternalModels(registry);
+  // Register interface needed by both old and new buffer deallocation pass.
+  memref::registerAllocationOpInterfaceExternalModels(registry);
+  arith::registerBufferDeallocationOpInterfaceExternalModels(registry);
 
   return registry;
 }

src/Compiler/CompilerOptions.cpp

@@ -42,6 +42,8 @@ bool enableONNXHybridPass;                             // common for both
 std::vector<std::string> functionsToDecompose;         // common for both
 std::string opsForCall;                                // common for both
 bool disableKrnlOpFusion;                              // common for both
+bool enableKrnlBufferReuse;                            // common for both
+bool disableMemRefPrefetch;                            // common for both
 EmissionTargetType emissionTarget;                     // onnx-mlir only
 bool invokeOnnxVersionConverter;                       // onnx-mlir only
 bool preserveLocations;                                // onnx-mlir only
@@ -211,6 +213,21 @@ static llvm::cl::opt<bool, true> disableKrnlOpFusionOpt(
     llvm::cl::location(disableKrnlOpFusion), llvm::cl::init(false),
     llvm::cl::cat(OnnxMlirCommonOptions));
 
+static llvm::cl::opt<bool, true> enableKrnlBufferReuseOpt(
+    "enable-krnl-buffer-reuse",
+    llvm::cl::desc("enable buffer reuse within an op in onnx-to-krnl pass"
+                   "(default=false)\n"
+                   "Set to 'true' if you want to enable buffer reuse."),
+    llvm::cl::location(enableKrnlBufferReuse), llvm::cl::init(false),
+    llvm::cl::cat(OnnxMlirCommonOptions));
+
+static llvm::cl::opt<bool, true> disableMemRefPrefetchOpt(
+    "disable-memref-prefetch",
+    llvm::cl::desc("disable generation of memref.prefetch (default=false)\n"
+                   "Set to 'true' if you want to disable prefetch."),
+    llvm::cl::location(disableMemRefPrefetch), llvm::cl::init(false),
+    llvm::cl::cat(OnnxMlirCommonOptions));
+
 static llvm::cl::opt<bool, true> disableRecomposeOptionOpt("disable-recompose",
     llvm::cl::desc("Disable recomposition of ONNX operations."),
     llvm::cl::location(disableRecomposeOption), llvm::cl::init(false),

src/Compiler/CompilerOptions.hpp

@@ -87,6 +87,8 @@ extern bool enableONNXHybridPass;                             // common for both
 extern std::vector<std::string> functionsToDecompose;         // common for both
 extern std::string opsForCall;                                // common for both
 extern bool disableKrnlOpFusion;                              // common for both
+extern bool enableKrnlBufferReuse;                            // common for both
+extern bool disableMemRefPrefetch;                            // common for both
 extern EmissionTargetType emissionTarget;                     // onnx-mlir only
 extern bool invokeOnnxVersionConverter;                       // onnx-mlir only
 extern bool preserveLocations;                                // onnx-mlir only