[ETHOSN] Transpose fully connected weights (apache#12970)

* [ETHOSN] Transpose fully connected weights

The NPU driver stack expects weights in IO (HWIO) format, however, Relay
uses an OI representation. Although the shape of the weight tensor was
correctly changed during codegen, the values in the weights tensor were
not being transposed. This lead to an output mismatch when the output
"units" was > 1. The tests didn't catch this due to using a weights
tensor of all 1's.

Change-Id: I51b2bcd14b677280ef3b6a6845d56b7dfacc7d6a

* Address comments

* Refactor use of weight transpose to common file between contrib
codegens.
* Make function areguments more explicit.
* Update network hashes.

Change-Id: Ib53bc7d2837b62908b92fd09062cbe9a8bb4ab30

* Fix lint

Change-Id: I6a1d9ffa8e3a747b7c77c9b27aa1b1c0d4c5cbff

* Fix cmsis-nn weights transpose

Change-Id: Ie89e429da222ffe17bc8faf831bf59217008a68a

* Address comments

Change-Id: Ie5ded2db3024b9e2c5095f01adea65798fc1da55

cmake/modules/contrib/CMSISNN.cmake

@@ -18,6 +18,8 @@
 if(USE_CMSISNN)
   add_definitions(-DTVM_USE_CMSISNN)
   message(STATUS "Build with CMSIS-NN support")
-  tvm_file_glob(GLOB RELAY_CONTRIB_CMSISNN_SRCS src/relay/backend/contrib/cmsisnn/*.cc)
+  tvm_file_glob(GLOB RELAY_CONTRIB_CMSISNN_SRCS
+    src/relay/backend/contrib/cmsisnn/*.cc
+    src/relay/backend/contrib/constant_transforms.cc)
   list(APPEND COMPILER_SRCS ${RELAY_CONTRIB_CMSISNN_SRCS})
 endif(USE_CMSISNN)

cmake/modules/contrib/EthosN.cmake

@@ -35,7 +35,8 @@ if(NOT USE_ETHOSN STREQUAL "OFF")
     list(APPEND RUNTIME_SRCS ${ETHOSN_RUNTIME_CONTRIB_SRC})
 
     tvm_file_glob(GLOB COMPILER_ETHOSN_SRCS
-                  src/relay/backend/contrib/ethosn/*)
+                  src/relay/backend/contrib/ethosn/*
+                  src/relay/backend/contrib/constant_transforms.cc)
     list(APPEND COMPILER_SRCS ${COMPILER_ETHOSN_SRCS})
 
     list(APPEND TVM_LINKER_LIBS ${ETHOSN_COMPILER_LIBRARY}

src/relay/backend/contrib/cmsisnn/generate_constants.cc

@@ -31,6 +31,7 @@
 #include "../../../op/make_op.h"
 #include "../../../qnn/utils.h"
 #include "../../../transforms/pattern_utils.h"
+#include "../constant_transforms.h"
 #include "convolutions.h"
 
 namespace tvm {
@@ -64,22 +65,9 @@ class GenerateConstantsMutator : public MixedModeMutator {
     attrs->out_dtype = std::move(conv2d_attrs->out_dtype);
     *new_attrs = tvm::Attrs{attrs};
 
-    std::string kernel_layout = conv2d_attrs->kernel_layout.c_str();
-    int pos_o = kernel_layout.find("O");
-    int pos_h = kernel_layout.find("H");
-    int pos_w = kernel_layout.find("W");
-    int pos_i = kernel_layout.find("I");
-
-    IRModule kernel_module;
-    auto func_body = MakeTranspose(
-        kernel_expr, {Integer(pos_o), Integer(pos_h), Integer(pos_w), Integer(pos_i)});
-    auto kernel_func =
-        Function(FreeVars(func_body), func_body, Type(), FreeTypeVars(func_body, kernel_module));
-    GlobalVar kernel_var("main");
-    kernel_module->Add(kernel_var, kernel_func);
-    kernel_module = relay::transform::FoldConstant()(kernel_module);
-    kernel_func = Downcast<Function>(kernel_module->Lookup("main"));
-    return kernel_func->body;
+    Constant conv2d_kernel = Downcast<Constant>(kernel_expr);
+    conv2d_kernel = TransposeWeights(conv2d_kernel, conv2d_attrs->kernel_layout, "OHWI");
+    return conv2d_kernel;
   }
 
   /*!  * \brief Performs weight transpose and substitutes existing constants in the composite

src/relay/backend/contrib/constant_transforms.cc

@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "constant_transforms.h"
+
+#include <string>
+
+#include "../../transforms/pattern_utils.h"
+#include "../../transforms/simplify_expr.h"
+
+/*!
+ * \file src/relay/backend/contrib/constant_transforms.cc
+ * \brief Transforms applied to constant operations during codegen for BYOC backends.
+ */
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+
+Expr FoldConstantExpr(const Expr& expr, bool fold_qnn) {
+  auto mod = IRModule::FromExpr(expr);
+  mod = transform::FoldConstant(fold_qnn)(mod);
+  auto entry_func = Downcast<Function>(mod->Lookup("main"));
+  return expr.as<FunctionNode>() == nullptr ? entry_func->body : entry_func;
+}
+
+Constant TransposeWeights(const Constant& data, const std::string& source_layout,
+                          const std::string& target_layout) {
+  Array<Integer> transpose_matrix;
+  for (const char& c : target_layout) {
+    int pos = source_layout.find(c);
+    transpose_matrix.push_back(pos);
+  }
+  Expr transpose = MakeTranspose(data, transpose_matrix);
+  transpose = InferType(FoldConstantExpr(transpose));
+  Constant transposed_data = Downcast<Constant>(transpose);
+  return transposed_data;
+}
+
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm

src/relay/backend/contrib/constant_transforms.h

@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file src/relay/backend/contrib/constant_transforms.h
+ * \brief Transforms applied to constant operations during codegen for BYOC backends.
+ */
+
+#ifndef TVM_RELAY_BACKEND_CONTRIB_CONSTANT_TRANSFORMS_H_
+#define TVM_RELAY_BACKEND_CONTRIB_CONSTANT_TRANSFORMS_H_
+
+#include <tvm/relay/expr.h>
+
+#include <string>
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+
+/*!
+ * \brief Apply constant folding on an expression.
+ *
+ * \param expr The expression to fold.
+ * \param fold_qnn Whether to fold constants for QNN operations.
+ * \returns The new folded expression.
+ */
+Expr FoldConstantExpr(const Expr& expr, bool fold_qnn = true);
+
+/*!
+ *\brief Transpose weights from `source_layout` to `target_layout`
+ *
+ * \param data The constant expression to transpose.
+ * \param source_layout The current layout of the constant e.g. "OHWI".
+ * \param target_layout The target layout of the constant e.g. "HWIO".
+ */
+Constant TransposeWeights(const Constant& data, const std::string& source_layout,
+                          const std::string& target_layout);
+
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_BACKEND_CONTRIB_CONSTANT_TRANSFORMS_H_

src/relay/backend/contrib/ethosn/codegen.cc

@@ -412,8 +412,8 @@ EthosnError ConstructNetworkVisitor::MakeFullyConnectedLayer(const Call& call,
     return err;
   }
 
-  auto weights = AddConstant(network_, params.weights_info, params.raw_weights).tensor;
-  auto bias = AddConstant(network_, params.bias_info, params.raw_bias).tensor;
+  auto weights = AddConstant(network_, params.weights_info, params.raw_weights->data).tensor;
+  auto bias = AddConstant(network_, params.bias_info, params.raw_bias->data).tensor;
   try {
     auto input =
         AddReshape(network_, *operand_table_[call->args[0]][0], params.input_info.m_Dimensions)

src/relay/backend/contrib/ethosn/convert_equivalent.cc

@@ -32,6 +32,7 @@
 #include "../../../qnn/utils.h"
 #include "../../../transforms/pattern_utils.h"
 #include "../../../transforms/simplify_expr.h"
+#include "../constant_transforms.h"
 #include "ethosn_api.h"
 
 namespace tvm {

src/relay/backend/contrib/ethosn/ethosn_api.cc

@@ -41,6 +41,7 @@
 #include "../../../op/make_op.h"
 #include "../../../transforms/pattern_utils.h"
 #include "../../../transforms/simplify_expr.h"
+#include "../constant_transforms.h"
 #include "ethosn_support_library/Support.hpp"
 #include "ethosn_support_library/SupportQueries.hpp"
 #include "tvm/relay/qnn/attrs.h"
@@ -197,7 +198,10 @@ EthosnError EthosnAPI::QnnFullyConnected(const Expr& expr, FullyConnectedParams*
   sl::QuantizationInfo output_q_info;
   err += Tvm2Npu(input_zero_point, input_scale, &data_q_info);
   err += Tvm2Npu(kernel_zero_point, kernel_scale, &weights_q_info);
-  err += Tvm2Npu(0, data_q_info.GetScale() * weights_q_info.GetScale(), &bias_q_info);
+  std::valarray<float> bias_scales = data_q_info.GetScale() * weights_q_info.GetScales();
+  const int bias_zero_point = 0;
+  const unsigned int bias_axis = 3;
+  err += Tvm2Npu(bias_zero_point, bias_scales, bias_axis, &bias_q_info);
   err += Tvm2Npu(output_zero_point, output_scale, &output_q_info);
 
   // Create fc info
@@ -213,27 +217,30 @@ EthosnError EthosnAPI::QnnFullyConnected(const Expr& expr, FullyConnectedParams*
                                       data_data_type, sl::DataFormat::NHWC, data_q_info);
 
   // Create weights info
-  const auto* weights_dtype = dense->args[1]->checked_type().as<TensorTypeNode>();
+  Constant weights_data = Downcast<Constant>(dense->args[1]);
+  weights_data = TransposeWeights(weights_data, "OI", "IO");
+  const auto* weights_ttype = weights_data->checked_type().as<TensorTypeNode>();
   sl::TensorShape weights_tensor_shape;
   sl::DataType weights_data_type;
   sl::DataFormat weights_data_format;
   // Ignore the error here because weights don't have a batch axis
-  Tvm2Npu(weights_dtype->shape, &weights_tensor_shape);
-  err += Tvm2Npu(weights_dtype->dtype, &weights_data_type);
+  Tvm2Npu(weights_ttype->shape, &weights_tensor_shape);
+  err += Tvm2Npu(weights_ttype->dtype, &weights_data_type);
   err += Tvm2Npu("HWIO", &weights_data_format);
-  params->weights_info = sl::TensorInfo({1, 1, weights_tensor_shape[1], weights_tensor_shape[0]},
+  // Weights tensor shape is 1, 1, I, O
+  params->weights_info = sl::TensorInfo({1, 1, weights_tensor_shape[0], weights_tensor_shape[1]},
                                         weights_data_type, weights_data_format, weights_q_info);
-  params->raw_weights = dense->args[1].as<ConstantNode>()->data->data;
+  params->raw_weights = weights_data->data;
 
   // Create bias info
   params->bias_info =
-      sl::TensorInfo({1, 1, 1, weights_tensor_shape[0]}, sl::DataType::INT32_QUANTIZED,
+      sl::TensorInfo({1, 1, 1, weights_tensor_shape[1]}, sl::DataType::INT32_QUANTIZED,
                      sl::DataFormat::NHWC, bias_q_info);
-  params->raw_bias = bias_add->args[1].as<ConstantNode>()->data->data;
+  params->raw_bias = bias_add->args[1].as<ConstantNode>()->data;
 
   sl::TensorInfo output_tensor_info;
   err += Tvm2Npu(requantize->checked_type(), &output_tensor_info);
-  output_tensor_info.m_Dimensions = {data_tensor_shape[0], 1, 1, weights_tensor_shape[0]};
+  output_tensor_info.m_Dimensions = {data_tensor_shape[0], 1, 1, weights_tensor_shape[1]};
   output_tensor_info.m_QuantizationInfo = output_q_info;
   params->output_info = output_tensor_info;
 
@@ -449,21 +456,6 @@ EthosnError EthosnAPI::Mean(const Expr& expr, MeanParams* params) {
   return err;
 }
 
-Constant TransposeWeights(const Constant& data, const std::string& input_layout) {
-  int pos_h = input_layout.find("H");
-  int pos_w = input_layout.find("W");
-  int pos_i = input_layout.find("I");
-  int pos_o = input_layout.find("O");
-
-  // Currently the expected target layout is HWIO only.
-  Array<Integer> target_shape = {pos_h, pos_w, pos_i, pos_o};
-
-  Expr transpose = MakeTranspose(data, target_shape);
-  transpose = InferType(FoldConstantExpr(transpose));
-  Constant transposed_data = Downcast<Constant>(transpose);
-  return transposed_data;
-}
-
 EthosnError EthosnAPI::QnnConv2dTranspose(const Expr& expr, QnnConv2dTransposeParams* params) {
   Call requantize = Downcast<Call>(expr);
   Call bias;
@@ -530,7 +522,7 @@ EthosnError EthosnAPI::QnnConv2dTranspose(const Expr& expr, QnnConv2dTransposePa
   // Create weights info
   Constant weights_data = Downcast<Constant>(conv2d_transpose->args[1]);
   if (conv_attr->kernel_layout != "HWIO") {
-    weights_data = TransposeWeights(weights_data, conv_attr->kernel_layout);
+    weights_data = TransposeWeights(weights_data, conv_attr->kernel_layout, "HWIO");
   }
   const auto* weights_ttype = weights_data->checked_type().as<TensorTypeNode>();
   sl::TensorShape weights_tensor_shape;
@@ -1080,13 +1072,6 @@ EthosnError EthosnAPI::AsConstant(const Expr& expr, T* out) {
   return EthosnError();
 }
 
-Expr FoldConstantExpr(const Expr& expr, bool fold_qnn) {
-  auto mod = IRModule::FromExpr(expr);
-  mod = transform::FoldConstant(fold_qnn)(mod);
-  auto entry_func = Downcast<Function>(mod->Lookup("main"));
-  return expr.as<FunctionNode>() == nullptr ? entry_func->body : entry_func;
-}
-
 }  // namespace ethosn
 }  // namespace contrib
 }  // namespace relay

src/relay/backend/contrib/ethosn/ethosn_api.h

@@ -66,8 +66,8 @@ struct FullyConnectedParams {
   sl::TensorInfo weights_info;
   sl::TensorInfo bias_info;
   sl::TensorInfo output_info;
-  void* raw_weights = nullptr;
-  void* raw_bias = nullptr;
+  runtime::NDArray raw_weights;
+  runtime::NDArray raw_bias;
 };
 
 struct MaxPool2DParams {
@@ -324,15 +324,6 @@ class EthosnAPI {
   static EthosnError AsConstant(const Expr& expr, std::valarray<float>* out);
 };
 
-/*!
- * \brief Apply constant folding on an expression.
- *
- * \param expr The expression to fold.
- * \param fold_qnn Whether to fold constants for QNN operations.
- * \returns The new folded expression.
- */
-Expr FoldConstantExpr(const Expr& expr, bool fold_qnn = true);
-
 }  // namespace ethosn
 }  // namespace contrib
 }  // namespace relay