Merge onnx/onnx-mlir dce84b5 into zosdev/onnx-mlir metis (onnx#192)

* 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]>
cjvolzka · Nov 14, 2023 · fa9a71c · fa9a71c
1 parent c3c267b
commit fa9a71c
Show file tree

Hide file tree

Showing 55 changed files with 847 additions and 253 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -11,6 +11,7 @@ option(ONNX_MLIR_ENABLE_STABLEHLO "Enable StableHLO support." ON)
 option(ONNX_MLIR_DECOMP_ONNX_CONVTRANSPOSE "Enable ONNXConvTransposeOp decomposition." ON)
 option(ONNX_MLIR_ENABLE_WERROR "Enable warnings as errors." OFF)
 option(ONNX_MLIR_SUPPRESS_THIRD_PARTY_WARNINGS "Suppress warning in third_party code." ON)
+option(ONNX_MLIR_ENABLE_JAVA "Set to ON for building the Java runtime, tools, and tests" ON)
 
 set(CMAKE_CXX_STANDARD 17)
 
@@ -87,37 +88,42 @@ endif()
 
 find_package(Python3 ${LLVM_MINIMUM_PYTHON_VERSION} REQUIRED COMPONENTS Interpreter Development)
 
-# On MacOS, find_package for Java and JNI doesn't work if JAVA_INCLUDE_PATH
-# and JAVA_INCLUDE_PATH2 are not set. Setting them requires JAVA_HOME to be
-# known. If the JAVA_HOME env var isn't set, we try to figure it out by running
-# java and looking for the java.home property.
-if (APPLE)
-  set(JAVA_HOME $ENV{JAVA_HOME})
-  if ("${JAVA_HOME}" STREQUAL "")
-    execute_process(
-      COMMAND bash -c "java -XshowSettings:properties -version 2>&1 | \
-                       sed -n 's/[ \t]*java.home = //p'"
-      OUTPUT_STRIP_TRAILING_WHITESPACE
-      OUTPUT_VARIABLE JAVA_HOME)
-  endif()
-  if (NOT "${JAVA_HOME}" STREQUAL "")
-    SET(JAVA_INCLUDE_PATH  "${JAVA_HOME}/include")
-    SET(JAVA_INCLUDE_PATH2 "${JAVA_HOME}/include/darwin")
+if(ONNX_MLIR_ENABLE_JAVA)
+  # On MacOS, find_package for Java and JNI doesn't work if JAVA_INCLUDE_PATH
+  # and JAVA_INCLUDE_PATH2 are not set. Setting them requires JAVA_HOME to be
+  # known. If the JAVA_HOME env var isn't set, we try to figure it out by running
+  # java and looking for the java.home property.
+  if (APPLE)
+    set(JAVA_HOME $ENV{JAVA_HOME})
+    if ("${JAVA_HOME}" STREQUAL "")
+      execute_process(
+        COMMAND bash -c "java -XshowSettings:properties -version 2>&1 | \
+                         sed -n 's/[ \t]*java.home = //p'"
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        OUTPUT_VARIABLE JAVA_HOME)
+    endif()
+    if (NOT "${JAVA_HOME}" STREQUAL "")
+      SET(JAVA_INCLUDE_PATH  "${JAVA_HOME}/include")
+      SET(JAVA_INCLUDE_PATH2 "${JAVA_HOME}/include/darwin")
+    endif()
   endif()
-endif()
 
-# Don't require AWT to allow headless JDK to reduce docker image size
-set(JAVA_AWT_INCLUDE_PATH headless)
-set(JAVA_AWT_LIBRARY headless)
-find_package(Java COMPONENTS Development)
-find_package(JNI)
+  # Don't require AWT to allow headless JDK to reduce docker image size
+  set(JAVA_AWT_INCLUDE_PATH headless)
+  set(JAVA_AWT_LIBRARY headless)
+  find_package(Java COMPONENTS Development)
+  find_package(JNI)
 
-if (Java_Development_FOUND AND JNI_FOUND)
-  set(ONNX_MLIR_ENABLE_JNI TRUE)
-  include(UseJava)
+  if (Java_Development_FOUND AND JNI_FOUND)
+    set(ONNX_MLIR_ENABLE_JNI TRUE)
+    include(UseJava)
+  else()
+    set(ONNX_MLIR_ENABLE_JNI FALSE)
+    message(WARNING "Java Development component or JNI not found, JNI targets will not work")
+  endif()
 else()
   set(ONNX_MLIR_ENABLE_JNI FALSE)
-  message(WARNING "Java Development component or JNI not found, JNI targets will not work")
+  message(WARNING "Java explicitly disabled with the ONNX_MLIR_ENABLE_JAVA option, JNI targets will not work")
 endif()
 
 if (CMAKE_CXX_COMPILER_ID MATCHES "MSVC")

diff --git a/docs/LocationInfo.md b/docs/LocationInfo.md
@@ -72,9 +72,9 @@ Then location info can be found in the output of test_add.onnx.onnx.mlir
 The test_add.onnx.mlir content:
 
 ```
-  1 module attributes {llvm.data_layout = "e-m:o-p270:32:32-p271:32:32-p    272:64:64-i64:64-f80:128-n8:16:32:64-S128", llvm.target_triple = "x8    6_64-apple-darwin22.3.0", "onnx-mlir.symbol-postfix" = "test_add"} {
-  2   func.func @main_graph(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x4x    5xf32>) -> tensor<3x4x5xf32> attributes {input_names = ["x", "y"], o    utput_names = ["sum"]} {
-  3     %0 = "onnx.Add"(%arg0, %arg1) : (tensor<3x4x5xf32>, tensor<3x4x5    xf32>) -> tensor<3x4x5xf32>
+  1 module attributes {llvm.data_layout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-apple-darwin22.3.0", "onnx-mlir.symbol-postfix" = "test_add"} {
+  2   func.func @main_graph(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x4x5xf32>) -> tensor<3x4x5xf32> {
+  3     %0 = "onnx.Add"(%arg0, %arg1) : (tensor<3x4x5xf32>, tensor<3x4x5xf32>) -> tensor<3x4x5xf32>
   4     onnx.Return %0 : tensor<3x4x5xf32>
   5   }
   6   "onnx.EntryPoint"() {func = @main_graph} : () -> ()

diff --git a/docs/Testing.md b/docs/Testing.md
@@ -316,7 +316,7 @@ $gdb Debug/bin/run-onnx-lib
 (gdb) run ./test_add.so
 (gdb) list
 1	builtin.module  {
-2	  builtin.func @main_graph(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x4x5xf32>) -> tensor<3x4x5xf32> attributes {input_names = ["x", "y"], output_names = ["sum"]} {
+2	  builtin.func @main_graph(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x4x5xf32>) -> tensor<3x4x5xf32> {
 3	    %0 = "onnx.Add"(%arg0, %arg1) : (tensor<3x4x5xf32>, tensor<3x4x5xf32>) -> tensor<3x4x5xf32>
 4	    return %0 : tensor<3x4x5xf32>
 5	  }

diff --git a/src/Builder/FrontendDialectTransformer.cpp b/src/Builder/FrontendDialectTransformer.cpp
@@ -190,6 +190,7 @@ class FrontendGenImpl {
   // mapping between string name and symbol
   ValueSymbolMapping frontend_symbols_;
 
+  // Keep shape information set by users.
   ModelInputShaper modelInputShaper_;
 
   using ImportHandlerType = void (onnx_mlir::detail::FrontendGenImpl::*)(
@@ -286,8 +287,10 @@ class FrontendGenImpl {
   /*!
    * Import an onnx tensor type by determining and returning its type
    * @param type_proto onnx tensor TypeProto.
+   * @param dim_params a comma-separated string of dimIndex:dimParam.
    */
-  Type ImportTensorType(const onnx::TypeProto &type_proto) {
+  Type ImportTensorType(
+      const onnx::TypeProto &type_proto, std::string *dim_params = nullptr) {
     assert(type_proto.value_case() == onnx::TypeProto::kTensorType &&
            "expect tensor type");
     std::vector<int64_t> dims;
@@ -300,6 +303,7 @@ class FrontendGenImpl {
     auto shape_proto = tensor_type.shape();
     for (int i = 0; i < shape_proto.dim_size(); i++) {
       if (shape_proto.dim()[i].dim_value()) {
+        // Dim is a constant value.
         int dim_numeric_size = shape_proto.dim()[i].dim_value();
         assert(dim_numeric_size != 0 &&
                "Parsed an tensor with a dimension size of zero");
@@ -309,6 +313,14 @@ class FrontendGenImpl {
           // If dim_value < 0, then dim is parametric.
           dims.push_back(ShapedType::kDynamic);
         }
+      } else if (dim_params && shape_proto.dim()[i].has_dim_param()) {
+        // Dim is unknown but assigned a string ID that can be used to check
+        // equality between unknown dimensions.
+        if (!dim_params->empty())
+          *dim_params += ",";
+        *dim_params +=
+            std::to_string(i) + ":" + shape_proto.dim()[i].dim_param();
+        dims.push_back(ShapedType::kDynamic);
       } else {
         dims.push_back(ShapedType::kDynamic);
       }
@@ -318,13 +330,14 @@ class FrontendGenImpl {
     return RankedTensorType::get(tensor_dims, elementType);
   }
 
-  Type ImportSequenceType(const onnx::TypeProto &type_proto) {
+  Type ImportSequenceType(
+      const onnx::TypeProto &type_proto, std::string *dim_params = nullptr) {
     auto input_seq_type = type_proto.sequence_type();
     if (input_seq_type.has_elem_type()) {
       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);
+      Type mlir_elem_type = ImportTensorType(elem_type, dim_params);
       if (!mlir_elem_type.isa<ShapedType>())
         llvm_unreachable("Seq type is incorrect");
       Type seq_type = mlir::SeqType::get(mlir_elem_type.cast<ShapedType>(), -1);
@@ -343,13 +356,14 @@ class FrontendGenImpl {
     llvm_unreachable("unexpected type");
   }
 
-  Type ImportType(const onnx::TypeProto &type_proto) {
+  Type ImportType(
+      const onnx::TypeProto &type_proto, std::string *dim_params = nullptr) {
     switch (type_proto.value_case()) {
     case onnx::TypeProto::kTensorType:
-      return ImportTensorType(type_proto);
+      return ImportTensorType(type_proto, dim_params);
       break;
     case onnx::TypeProto::kSequenceType:
-      return ImportSequenceType(type_proto);
+      return ImportSequenceType(type_proto, dim_params);
       break;
     case onnx::TypeProto::kOptionalType:
       return ImportOptionalType(type_proto);
@@ -468,17 +482,26 @@ class FrontendGenImpl {
     //  * maintain a list of the defined graph
     llvm::SmallVector<Type, 4> argTypes;
 
-    llvm::SmallVector<llvm::StringRef, 4> inputNames;
-    llvm::SmallVector<llvm::StringRef, 4> outputNames;
+    llvm::SmallVector<llvm::StringRef, 4> inputNames, outputNames;
+    // Keep dim_param for each dynamic dimension of each input tensor.
+    // In ONNX specification, two dynamic dimensions with the same dim_param
+    // string would be the same at runtime.
+    //
+    // See https://github.com/onnx/onnx/blob/main/docs/IR.md for more
+    // information about dim_param.
+    llvm::SmallVector<std::string, 4> inputDimParams, outputDimParams;
 
     // Import the input tensor types that are not constant and not initialized.
     int inputIndex = 0;
     for (const auto &input : graph.input()) {
       AddValueInfo(input);
       if (initializerNames.count(input.name()) == 0) {
         inputNames.push_back(input.name());
-        Type argTy = ImportType(input.type());
+        std::string dimParams = "";
+        Type argTy = ImportType(input.type(), &dimParams);
         argTy = modelInputShaper_.reshape(inputIndex, argTy);
+        if (!dimParams.empty())
+          inputDimParams.emplace_back(dimParams);
 
         argTypes.emplace_back(argTy);
 
@@ -524,7 +547,10 @@ class FrontendGenImpl {
     llvm::SmallVector<Value, 4> retVals;
     // Import the output tensors
     for (const auto &output : graph.output()) {
-      ImportOutputTensor(output, retTys, retVals);
+      std::string dimParams = "";
+      ImportOutputTensor(output, retTys, retVals, &dimParams);
+      if (!dimParams.empty())
+        outputDimParams.emplace_back(dimParams);
     }
 
     if (useReturn)
@@ -533,8 +559,21 @@ class FrontendGenImpl {
       // Create a return operation to return all ONNX output tensors.
       builder_.create<ONNXYieldOp>(UnknownLoc(), retVals);
 
-    op->setAttr("input_names", builder_.getStrArrayAttr(inputNames));
-    op->setAttr("output_names", builder_.getStrArrayAttr(outputNames));
+    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]));
+    if (!inputNames.empty())
+      op->setAttr("input_names", builder_.getStrArrayAttr(inputNames));
+    if (!outputNames.empty())
+      op->setAttr("output_names", builder_.getStrArrayAttr(outputNames));
+    if (!inputDimParamsRefs.empty())
+      op->setAttr(
+          "input_dim_params", builder_.getStrArrayAttr(inputDimParamsRefs));
+    if (!outputDimParamsRefs.empty())
+      op->setAttr(
+          "output_dim_params", builder_.getStrArrayAttr(outputDimParamsRefs));
 
     frontend_symbols_.popScope(graph.name());
     onnx_type_map.popScope(graph.name());
@@ -1328,23 +1367,73 @@ class FrontendGenImpl {
    * @param output onnx output ValueInfoProto.
    * @param ret_types a vector of types representing graph's output types.
    * @param ret_vals a vector of mlir Value representing graph's output.
+   * @param dim_params a comma-separated string of dimIndex:dimParam.
    */
   void ImportOutputTensor(const onnx::ValueInfoProto &output,
       llvm::SmallVectorImpl<Type> &ret_types,
-      llvm::SmallVectorImpl<Value> &ret_vals) {
+      llvm::SmallVectorImpl<Value> &ret_vals,
+      std::string *dim_params = nullptr) {
     const Value *valPtr = frontend_symbols_.GetByOnnxName(output.name());
     Value val = *valPtr;
     if (output.type().value_case() == onnx::TypeProto::kTensorType) {
       if (output.type().tensor_type().has_shape()) {
-        val.setType(ImportType(output.type()));
+        val.setType(ImportType(output.type(), dim_params));
       }
       ret_types.emplace_back(val.getType());
     } else {
-      ret_types.emplace_back(ImportType(output.type()));
+      ret_types.emplace_back(ImportType(output.type(), dim_params));
     }
     ret_vals.push_back(val);
   }
 
+  // Move function attributes for argument/result names and dim_params into
+  // argument/result attributes.
+  void moveFuncAttrsToArgAttrs(func::FuncOp funcOp,
+      ArrayRef<std::string> funcAttrNames, ArrayRef<std::string> argAttrNames,
+      bool isArg) {
+    assert(funcAttrNames.size() == argAttrNames.size() &&
+           "The number of attributes to move mismatched");
+    Operation *op = funcOp.getOperation();
+    size_t numOfArgs =
+        (isArg) ? funcOp.getNumArguments() : funcOp.getNumResults();
+
+    // Only move attributes that exists.
+    SmallVector<ArrayAttr, 2> funcAttrsToMove;
+    SmallVector<std::string, 2> targetArgAttrNames;
+    for (size_t i = 0; i < funcAttrNames.size(); ++i) {
+      ArrayAttr attr = op->getAttrOfType<ArrayAttr>(funcAttrNames[i]);
+      if (!attr)
+        continue;
+      funcAttrsToMove.emplace_back(attr);
+      targetArgAttrNames.emplace_back(argAttrNames[i]);
+    }
+
+    // Move function attributes to argument/result attributes.
+    for (size_t i = 0; i < numOfArgs; ++i) {
+      SmallVector<NamedAttribute, 2> argAttrs;
+      for (size_t k = 0; k < funcAttrsToMove.size(); ++k) {
+        if (i < funcAttrsToMove[k].size()) {
+          auto name = (funcAttrsToMove[k].getValue()[i]).cast<StringAttr>();
+          if (name) {
+            NamedAttribute namedAttr =
+                builder_.getNamedAttr(argAttrNames[k], name);
+            argAttrs.emplace_back(namedAttr);
+          }
+        }
+      }
+      if (!argAttrs.empty()) {
+        if (isArg)
+          funcOp.setArgAttrs(i, argAttrs);
+        else
+          funcOp.setResultAttrs(i, argAttrs);
+      }
+    }
+
+    // Clean up the function attributes.
+    for (std::string s : funcAttrNames)
+      op->removeAttr(s);
+  }
+
   /*!
    * Import ONNX main computation graph.
    * @param graph onnx graph proto.
@@ -1363,6 +1452,13 @@ class FrontendGenImpl {
         /*op=*/mainFunc.getOperation(), /*useReturn=*/true);
     mainFunc.setType(funcType);
 
+    // Move function attributes for argument/result names and dim_params into
+    // argument/result attributes.
+    moveFuncAttrsToArgAttrs(mainFunc, {"input_names", "input_dim_params"},
+        {"onnx.name", "onnx.dim_params"}, /*isArg=*/true);
+    moveFuncAttrsToArgAttrs(mainFunc, {"output_names", "output_dim_params"},
+        {"onnx.name", "onnx.dim_params"}, /*isArg=*/false);
+
     // Emit entry point op describing inference function signature.
     auto entryPoint = ONNXEntryPointOp::create(UnknownLoc(), mainFunc);
     module_.push_back(entryPoint);