Apply code review comments

Author: echuraev

include/tvm/relay/transform.h

@@ -120,8 +120,8 @@ TVM_DLL Pass FoldConstant(bool fold_qnn = false);
 /*!
  * \brief Split function with huge number of arguments to smaller pieces.
  *
- * \param max_function_args Maximum number of function arguments. If it is 0 then SplitArgs won't
- *                          split function.
+ * \param max_function_args Maximum number of function arguments. If it equals 0 then SplitArgs
+ *                          shouldn't split the function.
  *
  * \return The pass.
  */

python/tvm/relay/transform/transform.py

@@ -1379,7 +1379,8 @@ def SplitArgs(max_function_args):
     Parameters
     ----------
     max_function_args: int
-      Maximum number of function arguments. If it is 0 then SplitArgs won't split function.
+      Maximum number of function arguments. If it equals 0 then SplitArgs
+      shouldn't split the function.
 
 
     Returns

src/relay/analysis/graph_partitioner.cc

@@ -243,33 +243,33 @@ size_t GraphPartitioner::CountArgs_(IndexedForwardGraph::Node* src,
   ICHECK(gnode != nullptr);
   auto sum = gnode->args_num;
   visited_groups.insert(gnode->FindRoot());
-  auto calcArgs = [this, src, &graph, &visited_groups,
-                   update_postpone](const relay::Expr& arg) -> size_t {
+  auto calc_args_number = [this, src, &graph, &visited_groups,
+                           update_postpone](const relay::Expr& arg) -> size_t {
     if (arg.as<VarNode>()) return 0;
     auto* node = graph.node_map.at(arg.get());
     Group* prev_group = groups_[node->index]->FindRoot();
     if (visited_groups.count(prev_group) == 0) {
       visited_groups.insert(prev_group);
       if (prev_group->args_num > 0) {
-        // Get number of arguments from group
+        // Get the number of arguments from the group
         return prev_group->args_num;
       } else if (update_postpone) {
-        // Update pointer to node which should be postponed for deferred fusing
+        // Update pointer to the node which should be postponed for deferred fusing
         postpone_node_ = src;
       } else {
-        // Calculate number of arguments for the node which wasn't processed before
+        // Calculate the number of arguments for the node which wasn't processed before
         return CountArgs_(node, graph, update_postpone);
       }
     }
     return 0;
   };
   if (auto call_node = GetRef<ObjectRef>(src->ref).as<CallNode>()) {
     for (auto& it : call_node->args) {
-      sum += calcArgs(it);
+      sum += calc_args_number(it);
     }
   } else if (auto tuple_node = GetRef<ObjectRef>(src->ref).as<TupleNode>()) {
     for (auto& it : tuple_node->fields) {
-      sum += calcArgs(it);
+      sum += calc_args_number(it);
     }
   }
   return sum;
@@ -357,11 +357,11 @@ void GraphPartitioner::RunFuse(const IndexedForwardGraph& graph,    //
     Group* group_node = groups_[nid];
     ICHECK(group_node != nullptr);
     postpone_node_ = nullptr;
-    // Check if fusing of some inputs was postponed
+    // Check if the fusing of some inputs was postponed
     if (postponed_fusing_map_.count(graph_node)) {
       auto range = postponed_fusing_map_.equal_range(graph_node);
       for (auto it = range.first; it != range.second; ++it) {
-        // If number of arguments is less than limit then the input can be fused
+        // If the number of arguments is less than the limit then the input can be fused
         if (CountArgs_(graph_node, graph, false) <= CountArgsLimit_(graph_node)) {
           auto* src = it->second;
           auto* snode = post_dom_tree.nodes[src->index]->parent->gnode;
@@ -381,7 +381,7 @@ void GraphPartitioner::RunFuse(const IndexedForwardGraph& graph,    //
     // refuse the fusion if too many ops are going to be fused together
     if (CountFusedNodesWithNewChild(graph_node, dom_node->parent->gnode) > max_fuse_depth_)
       continue;
-    // refuse the fusion if too many arguments are going to be in fused function
+    // Refuse the fusion if too many arguments are going to be in the fused function
     if (max_function_args_ > 0) {
       auto limit = CountArgsLimit_(graph_node);
       if (limit > 0) {

src/relay/analysis/graph_partitioner.h

@@ -224,8 +224,8 @@ class GraphPartitioner {
       postponed_fusing_map_;
   /*!
    * \brief Fusing of this node should be postponed till all child nodes will be evaluated.
-   *        It is used to calculate number of arguments which will be passed to this node in
-   * generated function.
+   *        It is used to calculate the number of arguments which will be passed to this node in
+   *        the generated function.
    */
   const IndexedForwardGraph::Node* postpone_node_{nullptr};
   // Internal implementation of CheckPath
@@ -266,21 +266,21 @@ class GraphPartitioner {
   void CommitFuse(IndexedForwardGraph::Node* src, IndexedForwardGraph::Node* sink);
 
   size_t CountNodesUptoSink_(IndexedForwardGraph::Node* src, IndexedForwardGraph::Node* sink);
-  // Count the number of additional arguments. In case of dynamic shape,
-  // generated function takes several additional arguments, such as size of
-  // dynamic dimension and strides.
-  // This function calculates number of such additional arguments.
+  // Count the number of additional arguments. In the case of dynamic shape,
+  // generated function takes several additional arguments, such as the sizes of
+  // the dynamic dimensions and strides.
+  // This function calculates the number of such additional arguments.
   size_t CountAdditionalArgs_(const TensorTypeNode* ttype, bool with_strides = true);
   // Calculate the number of arguments for the node.
   size_t CountArgs_(IndexedForwardGraph::Node* src, const IndexedForwardGraph& graph,
                     bool update_postpone = true);
-  // Count actual limit of arguments for a generated function.
+  // Count the actual limit of arguments for a generated function.
   // max_function_args_ specifies the number of maximum function arguments. But
-  // usually, output tensors also passed to the function as arguments.
-  // Additionally, in case of dynamic shape, it is necessary to take into
-  // account the number of parameters which specifies the size of dynamic
-  // dimension.
-  // This function computes limit of arguments by the following formula:
+  // usually, output tensors are also passed to the function as arguments.
+  // Additionally, in the case of dynamic shape, it is necessary to take into
+  // account the number of parameters which specifies the sizes of the dynamic
+  // dimensions.
+  // This function computes the limit of arguments by the following formula:
   // limit = max_function_args_ - output_args_count
   size_t CountArgsLimit_(const IndexedForwardGraph::Node* child);
 
@@ -292,7 +292,9 @@ class GraphPartitioner {
   // is important for correct calculation.
   size_t CountFusedNodesWithNewChild(IndexedForwardGraph::Node* child,
                                      IndexedForwardGraph::Node* dom_parent);
-  // Count the number of arguments in a fused subgraph if the output of child is additionally fused.
+  // Count the number of arguments in a fused subgraph. This function also takes into account the
+  // number of the child's output node argument. It helps to stop fusing before the node when the
+  // limit will be exceeded.
   size_t CountFusedArgs(const IndexedForwardGraph& graph, IndexedForwardGraph::Node* child);
 
   // Initialize the groups.

src/relay/transforms/split_args.cc

@@ -61,16 +61,17 @@ class ArgumentSplitter : public ExprRewriter {
     return lastExpr;
   }
 
-  // In case of dynamic shape in tensor, size of any_dims and strides are passed as function args
+  // In the case of dynamic shape in tensor, the sizes of any_dims and strides are passed as
+  // function args
   size_t CalculateNumberOfAdditionalArgs_(const TensorTypeNode* arg, bool isOutput = false) {
     size_t num = 0;
     for (const auto& dim : arg->shape) {
       if (dim.as<AnyNode>()) {
         num++;
       }
     }
-    // In case of dynamic shape also strides will be passed to function
-    // as arguments. Number of strides equals to the rank of the tensor.
+    // In the case of dynamic shape, strides are also passed to a function as arguments. The number
+    // of strides equals the rank of the tensor.
     if (num > 0 && isOutput)
       return arg->shape.size();
     else if (num > 0)

src/target/target_kind.cc

@@ -365,8 +365,8 @@ TVM_REGISTER_TARGET_KIND("opencl", kDLOpenCL)
     .add_attr_option<Integer>("max_num_threads", Integer(256))
     .add_attr_option<Integer>("thread_warp_size", Integer(1))
     .add_attr_option<Integer>("texture_spatial_limit", Integer(16384))
-    // Faced that Qualcomm OpenCL runtime was crashed without any error message in
-    // case when the number of kernel arguments was pretty big. OpenCL doesn't
+    // Faced that Qualcomm OpenCL runtime crashed without any error message in
+    // the case when the number of kernel arguments was pretty big. OpenCL doesn't
     // specify any limitations on the number of kernel arguments. max_function_args
     // equals to 128 looks like a reasonable number of kernel arguments.
     .add_attr_option<Integer>("max_function_args", Integer(128))

tests/python/unittest/test_target_codegen_opencl.py

@@ -357,10 +357,11 @@ def ref_impl(inputs):
 @tvm.testing.requires_opencl
 def test_fuse_concat_max_num_args(executor_type, shape_type):
     """
-    In this test before concat we have an operation with 3 inputs. In the
-    SplitArgs we cannot calculate these inputs as inputs to concat, because
-    they will be added to the concat after fusing operation. So FuseOps pass
-    should handle this case and stop fusing before concat.
+    In this test, we have an operation with 3 inputs before concat. In the
+    SplitArgs we cannot calculate these inputs as inputs to the concat layer,
+    because they will be added to the concat after the fusing operation. So
+    FuseOps pass should handle this case and stop fusing before the concat
+    layer.
 
     The example:
        x     y     z                  x     y     z