Import dim_param for model inputs and outputs #2616

tungld · 2023-11-10T13:35:12Z

In ONNX specification, a dimension can be dim_value (static) or dim_param (a named dynamic dimension).
If two dynamic dimensions has the same dim_param, they must have the same value at runtime.

This patch will import dim_param of the inputs and outputs of a model, and embed such dim_param in argument and result attributes. This patch also moves input_names and output_names into argument and result attributes onnx.name. for example

func.func @main_graph(
         %arg0: tensor<?x?xf32> {onnx.dim_params = "0:batch_size,1:sequence_len", onnx.name = "X"}, 
         %arg1: tensor<1x?xf32> {onnx.dim_params = "1:sequence_len", onnx.name = "Y"}) 
    -> (tensor<?x?xf32> {onnx.dim_params = "0:batch_size,1:sequence_len", onnx.name = "Z"})

onnx.dim_params is a list of "dim_index:dim_param"

In a practical model, e.g. t5-decoder with KV cache, dim_param is like the following, and we know that all the first dimensions (batch_size) must be the same.

These dim_param information will be used later by the dynamic dimension analysis to make the analysis stronger. Also we can check if users' inputs to the model are consistent or not.

Signed-off-by: Tung D. Le <[email protected]>

AlexandreEichenberger

LGTM, I had wanted something like this for a long time, it's great that it is part of the ONNX protocol so that we don't have to reinvent the wheel. With our efficient dimAnalysis, this should really help out disambiguating dynamic dimensions that are really identical.

sorenlassen

really nice PR, I had always hoped that your dynamic dimensions framework would apply to onnx's dim_params

sorenlassen · 2023-11-10T15:14:57Z

src/Builder/FrontendDialectTransformer.cpp

@@ -324,7 +334,8 @@ class FrontendGenImpl {
      onnx::TypeProto elem_type = input_seq_type.elem_type();
      assert(elem_type.value_case() == onnx::TypeProto::kTensorType &&
             "expect tensor inside sequence type");
-      Type mlir_elem_type = ImportTensorType(elem_type);
+      std::string s;


a little bit easier to read if you rename s to dimParam

or consider adding convenience methods ImportTensorType and ImportType without the dimParam parameter, or make dimParam a pointer that defaults to nullptr

Done. Changed to use optional arguments.

sorenlassen · 2023-11-10T15:15:43Z

src/Builder/FrontendDialectTransformer.cpp

@@ -287,7 +288,8 @@ class FrontendGenImpl {
   * Import an onnx tensor type by determining and returning its type
   * @param type_proto onnx tensor TypeProto.


please add a @param description of dimParam

Done. Thanks!

sorenlassen · 2023-11-10T15:29:22Z

src/Builder/FrontendDialectTransformer.cpp

+    SmallVector<llvm::StringRef> inputDimParamsRefs, outputDimParamsRefs;
+    for (uint64_t i = 0; i < inputDimParams.size(); ++i)
+      inputDimParamsRefs.emplace_back(llvm::StringRef(inputDimParams[i]));
+    for (uint64_t i = 0; i < outputDimParams.size(); ++i)
+      outputDimParamsRefs.emplace_back(llvm::StringRef(outputDimParams[i]));
    op->setAttr("input_names", builder_.getStrArrayAttr(inputNames));
    op->setAttr("output_names", builder_.getStrArrayAttr(outputNames));
+    op->setAttr(
+        "input_dim_params", builder_.getStrArrayAttr(inputDimParamsRefs));
+    op->setAttr(
+        "output_dim_params", builder_.getStrArrayAttr(outputDimParamsRefs));


I think it would be nicer if we made all of these argument and result attributes
https://github.com/sorenlassen/llvm-project/blob/main/mlir/include/mlir/Dialect/Func/IR/FuncOps.td#L241-L245

func.func @main_graph( %arg0: tensor<?x?xf32> {onnx.name = "X", onnx.dims = "0:batch_size,1:sequence_len"}, %arg1: tensor<1x?xi64> {onnx.name = "Y", onnx.dims = "1:sequence_len"} ) -> ( tensor<?x?xf32> {onnx.name = "Z", onnx.dims = "0:batch_size,1:sequence_len"} )

and then you can omit the onnx.dims attribute in the cases where it's empty

Great suggestion! I moved to use argument/result attributes instead of function attributes. Thanks!

sorenlassen · 2023-11-10T15:40:38Z

utils/testing/dim_param.py

+model_def = helper.make_model(graph_def, producer_name="onnx-mlir")
+
+onnx.checker.check_model(model_def)
+print(MessageToJson(model_def))


if you like the more concise textual representation of the .onnxtext lit tests you can change the last line to print(onnx.printer.to_text(model_def)) - see utils/onnx2text.py

It's quite simple so let me keep the current one.

Signed-off-by: Tung D. Le <[email protected]>

jenkins-droid · 2023-11-13T10:26:10Z

Jenkins Linux ppc64le Build #12382 [push] Import dim_param for mod... started at 05:34

jenkins-droid · 2023-11-13T10:26:10Z

Jenkins Linux amd64 Build #13362 [push] Import dim_param for mod... started at 04:26

jenkins-droid · 2023-11-13T10:26:11Z

Jenkins Linux s390x Build #13388 [push] Import dim_param for mod... started at 05:26

jenkins-droid · 2023-11-13T11:40:00Z

Jenkins Linux amd64 Build #13362 [push] Import dim_param for mod... failed after 1 hr 13 min

jenkins-droid · 2023-11-13T11:53:27Z

Jenkins Linux s390x Build #13388 [push] Import dim_param for mod... passed after 1 hr 27 min

jenkins-droid · 2023-11-13T12:13:40Z

Jenkins Linux ppc64le Build #12382 [push] Import dim_param for mod... passed after 1 hr 47 min

* detect LayerNorm in presence of reciprocal and div of 1 (onnx#2609) Signed-off-by: Alexandre Eichenberger <[email protected]> * [NNPA] Use F16 as element type for zTensor (onnx#2611) * Use f16 as element type for zTensor Signed-off-by: Tung D. Le <[email protected]> --------- Signed-off-by: Tung D. Le <[email protected]> * Layernorm: convert instance norm and group norm to layer norm. (onnx#2595) Signed-off-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Tung D. Le <[email protected]> * Parse and set --mcpu in onnx-mlir-opt command (onnx#2614) Signed-off-by: Tung D. Le <[email protected]> * Import dim_param for model inputs and outputs (onnx#2616) * Import dim_param for model inputs and outputs * use argument attributes Signed-off-by: Tung D. Le <[email protected]> --------- Signed-off-by: Tung D. Le <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> * [DialectBuilder] add builder funcrions for ONNXSumOp and ONNXConvOp (onnx#2572) The DialectBuilder class seems to be missing the function create the ONNXSumOp and ONNXConOp nodes and check their shape. This patch adds the necessary functions. Signed-off-by: Ashay Rane <[email protected]> Signed-off-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> * [StableHLO] Lowers PadOp (constant mode) & GatherElements Op to StableHLO (onnx#2602) * [Stablehlo] Pad constant mode & GatherElements to Stablehlo Signed-off-by: chongsong.chen <[email protected]> Signed-off-by: Yan Xu <[email protected]> Co-authored-by: chongsong.chen <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> * [build] Add cmake option to enable/disable Java components build (onnx#2613) * Add ONNX_MLIR_ENABLE_JAVA cmake option (default TRUE) Signed-off-by: Boyana Norris <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Tung D. Le <[email protected]> Co-authored-by: Ashay Rane <[email protected]> Co-authored-by: Yan Xu <[email protected]> Co-authored-by: chongsong.chen <[email protected]> Co-authored-by: Boyana Norris <[email protected]>

* 'main' of github.ibm.com:zosdev/onnx-mlir: Use dim_params in dynamic dimension analysis (onnx#2620) Update rapidcheck to include the fix for missing <cstdint> include (onnx#2623) Initial changes for llvm uplift (onnx#2568) [build] Add cmake option to enable/disable Java components build (onnx#2613) [StableHLO] Lowers PadOp (constant mode) & GatherElements Op to StableHLO (onnx#2602) [DialectBuilder] add builder funcrions for ONNXSumOp and ONNXConvOp (onnx#2572) Import dim_param for model inputs and outputs (onnx#2616) Parse and set --mcpu in onnx-mlir-opt command (onnx#2614) Layernorm: convert instance norm and group norm to layer norm. (onnx#2595) [NNPA] Use F16 as element type for zTensor (onnx#2611) detect LayerNorm in presence of reciprocal and div of 1 (onnx#2609) # Conflicts: # test/mlir/conversion/onnx_to_krnl/NN/Normalization_O3_SIMD_canonicalize.mlir

tungld added 3 commits November 10, 2023 22:19

Import dim_param for model inputs and outputs

1aed7d0

Signed-off-by: Tung D. Le <[email protected]>

add files

5310a49

Signed-off-by: Tung D. Le <[email protected]>

format

5dffe4e

Signed-off-by: Tung D. Le <[email protected]>

tungld changed the title ~~Import dim params~~ Import dim_param for a model inputs and outputs Nov 10, 2023

tungld requested review from AlexandreEichenberger, sorenlassen, chentong319 and philass November 10, 2023 13:38

remove debug statements

cb9c919

Signed-off-by: Tung D. Le <[email protected]>

tungld changed the title ~~Import dim_param for a model inputs and outputs~~ Import dim_param for model inputs and outputs Nov 10, 2023

AlexandreEichenberger approved these changes Nov 10, 2023

View reviewed changes

Merge branch 'main' into import_dim_params

6e0e413

sorenlassen approved these changes Nov 10, 2023

View reviewed changes

tungld added 6 commits November 13, 2023 11:50

change dim_params to an optional argument

8c27f6f

Signed-off-by: Tung D. Le <[email protected]>

correct element types in the lit tests

8686764

Signed-off-by: Tung D. Le <[email protected]>

use argument attributes

dda5b99

Signed-off-by: Tung D. Le <[email protected]>

Update the lowering of EntryPoint to use argument attributes

24f5785

Signed-off-by: Tung D. Le <[email protected]>

clean up input_names and output_names attributes

0e9ebc2

Signed-off-by: Tung D. Le <[email protected]>

change onnx.names to onnx.name

da6d7e3

Signed-off-by: Tung D. Le <[email protected]>

tungld merged commit 04e26e7 into onnx:main Nov 13, 2023
5 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Import dim_param for model inputs and outputs #2616

Import dim_param for model inputs and outputs #2616

tungld commented Nov 10, 2023 •

edited

Loading

AlexandreEichenberger left a comment

sorenlassen left a comment

sorenlassen Nov 10, 2023

tungld Nov 13, 2023

sorenlassen Nov 10, 2023

tungld Nov 13, 2023

sorenlassen Nov 10, 2023 •

edited

Loading

tungld Nov 13, 2023

sorenlassen Nov 10, 2023

tungld Nov 13, 2023

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

		@@ -287,7 +288,8 @@ class FrontendGenImpl {
		* Import an onnx tensor type by determining and returning its type
		* @param type_proto onnx tensor TypeProto.

Import dim_param for model inputs and outputs #2616

Import dim_param for model inputs and outputs #2616

Conversation

tungld commented Nov 10, 2023 • edited Loading

AlexandreEichenberger left a comment

Choose a reason for hiding this comment

sorenlassen left a comment

Choose a reason for hiding this comment

sorenlassen Nov 10, 2023

Choose a reason for hiding this comment

tungld Nov 13, 2023

Choose a reason for hiding this comment

sorenlassen Nov 10, 2023

Choose a reason for hiding this comment

tungld Nov 13, 2023

Choose a reason for hiding this comment

sorenlassen Nov 10, 2023 • edited Loading

Choose a reason for hiding this comment

tungld Nov 13, 2023

Choose a reason for hiding this comment

sorenlassen Nov 10, 2023

Choose a reason for hiding this comment

tungld Nov 13, 2023

Choose a reason for hiding this comment

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

jenkins-droid commented Nov 13, 2023

tungld commented Nov 10, 2023 •

edited

Loading

sorenlassen Nov 10, 2023 •

edited

Loading