Modify code to conform with style guidelines

Signed-off-by: Ruoqian Guo <[email protected]>

Author: ruoqianguo

core/conversion/converters/impl/expand.cpp

@@ -46,30 +46,30 @@ nvinfer1::ILayer* create_plugin(
 void addSliceInput(nvinfer1::Dims& dims, int idx, ConversionCtx* ctx, nvinfer1::ISliceLayer* slice) {
   int32_t rank = static_cast<int32_t>(dims.nbDims);
   int32_t* tmp = new int32_t[rank];
-  for(int i=0;i<rank;i++)
+  for (int i = 0; i < rank; i++)
     tmp[i] = dims.d[i];
   const nvinfer1::Dims d{1, {rank}};
   const nvinfer1::Weights w{nvinfer1::DataType::kINT32, tmp, rank};
   auto t = ctx->net->addConstant(d, w)->getOutput(0);
   slice->setInput(idx, *t);
 }
 
-nvinfer1::ITensor* vec2Tensor(int32_t *dim, int rank, ConversionCtx* ctx){
+nvinfer1::ITensor* vec2Tensor(int32_t* dim, int rank, ConversionCtx* ctx) {
   const nvinfer1::Dims d{1, {static_cast<int32_t>(rank)}};
   const nvinfer1::Weights w{nvinfer1::DataType::kINT32, dim, rank};
   return ctx->net->addConstant(d, w)->getOutput(0);
 }
 
-nvinfer1::ITensor * concat(int max_rank, int old_rank, ConversionCtx* ctx, nvinfer1::ITensor*tensor){
-  if(max_rank - old_rank > 0){
+nvinfer1::ITensor* concat(int max_rank, int old_rank, ConversionCtx* ctx, nvinfer1::ITensor* tensor) {
+  if (max_rank - old_rank > 0) {
     int32_t* tmp = new int32_t[max_rank - old_rank];
-    for(int i=0;i<(max_rank - old_rank);i++)
+    for (int i = 0; i < (max_rank - old_rank); i++)
       tmp[i] = 1;
     auto max_rank_tensor = vec2Tensor(tmp, max_rank - old_rank, ctx);
     auto in_shape_tensor = ctx->net->addShape(*tensor)->getOutput(0);
     nvinfer1::ITensor* const args[2] = {max_rank_tensor, in_shape_tensor};
     return ctx->net->addConcatenation(args, 2)->getOutput(0);
-  }else{ // max_rank - old_rank == 0
+  } else { // max_rank - old_rank == 0
     return ctx->net->addShape(*tensor)->getOutput(0);
   }
 }
@@ -86,7 +86,7 @@ bool add_expand(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor
     int64_t dim = input_dims.nbDims - 1 - offset;
     int64_t size = (dim >= 0) ? input_dims.d[dim] : 1;
     int64_t targetSize = expandedDims.d[i];
-    if(targetSize != -1){
+    if (targetSize != -1) {
       if (size != targetSize) {
         if (size != 1) {
           TRTORCH_THROW_ERROR(
@@ -95,12 +95,12 @@ bool add_expand(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor
                                               << " at dimension " << i);
         }
       }
-    }else{
-      if(dim < 0){
-        TRTORCH_THROW_ERROR("The expanded size of the tensor (" << \
-            targetSize << ") isn't allowed in a leading, non-existing dimension " << \
-            i);
-      }else{
+    } else {
+      if (dim < 0) {
+        TRTORCH_THROW_ERROR(
+            "The expanded size of the tensor (" << targetSize << ") isn't allowed in a leading, non-existing dimension "
+                                                << i);
+      } else {
         // in(3, 1), expand(3, -1, 4) -> expand(3, 3, 4)
         expandedDims.d[i] = input_dims.d[dim];
       }
@@ -146,18 +146,23 @@ bool add_expand(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor
   return true;
 }
 
-bool add_expand_dynamic(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* in, nvinfer1::ITensor* expandedDimsTensor){
-  auto input_shape_tensor =  ctx->net->addShape(*in)->getOutput(0);
+bool add_expand_dynamic(
+    ConversionCtx* ctx,
+    const torch::jit::Node* n,
+    nvinfer1::ITensor* in,
+    nvinfer1::ITensor* expandedDimsTensor) {
+  auto input_shape_tensor = ctx->net->addShape(*in)->getOutput(0);
   auto input_rank = in->getDimensions().nbDims;
   auto output_rank = expandedDimsTensor->getDimensions().d[0];
   TRTORCH_CHECK(
       input_rank <= output_rank,
       "Number of dimensions of the desired expansion must be greater than or equal to the number of input dimensions");
-  
+
   // add a plugin to check expandedDimsTensor whether match input_shape_tensor
-  auto expandShape_layer = create_plugin(ctx, n, input_shape_tensor, expandedDimsTensor, input_rank, output_rank, "expandShape");
+  auto expandShape_layer =
+      create_plugin(ctx, n, input_shape_tensor, expandedDimsTensor, input_rank, output_rank, "expandShape");
   auto _tensor = expandShape_layer->getOutput(0);
-  
+
   size_t max_rank = std::max(input_rank, output_rank);
 
   // Dimensions are right alignment
@@ -174,16 +179,19 @@ bool add_expand_dynamic(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1:
   nvinfer1::Dims starts_dim = util::toDims(c10::IntArrayRef(start_vec));
 
   // compute sizes = max(x,y).
-  auto sizes = ctx->net->addElementWise(*new_input_shape_tensor, *new_output_shape_tensor, nvinfer1::ElementWiseOperation::kMAX)->getOutput(0);
+  auto sizes =
+      ctx->net->addElementWise(*new_input_shape_tensor, *new_output_shape_tensor, nvinfer1::ElementWiseOperation::kMAX)
+          ->getOutput(0);
   nvinfer1::Dims sizes_dim{-1, {}};
   sizes_dim.nbDims = max_rank;
-  
+
   // Compute (x > 1 ? 1 : 0) for x in newDims, assuming positive x, using only TensorRT operations.
   // min(1, sub(input_shape, 1))
   int32_t* one_vector_tmp = new int32_t[1];
   one_vector_tmp[0] = 1;
   auto one_vector = vec2Tensor(one_vector_tmp, 1, ctx);
-  auto x_sub_one = ctx->net->addElementWise(*new_input_shape_tensor, *one_vector, nvinfer1::ElementWiseOperation::kSUB)->getOutput(0);
+  auto x_sub_one = ctx->net->addElementWise(*new_input_shape_tensor, *one_vector, nvinfer1::ElementWiseOperation::kSUB)
+                       ->getOutput(0);
   auto strides = ctx->net->addElementWise(*one_vector, *x_sub_one, nvinfer1::ElementWiseOperation::kMIN)->getOutput(0);
   nvinfer1::Dims strides_dim{-1, {}};
   strides_dim.nbDims = max_rank;
@@ -194,7 +202,7 @@ bool add_expand_dynamic(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1:
   slice->setInput(2, *sizes);
   slice->setInput(3, *strides);
 
-  auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], slice->getOutput(0)); 
+  auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], slice->getOutput(0));
 
   LOG_DEBUG("Expand layer output tensor shape: " << out_tensor->getDimensions());
 
@@ -203,94 +211,96 @@ bool add_expand_dynamic(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1:
 
 auto expand_registrations TRTORCH_UNUSED =
     RegisterNodeConversionPatterns()
-        .pattern({"aten::expand(Tensor(a) self, int[] size, *, bool implicit=False) -> (Tensor(a))",
-                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-                    auto in = args[0].ITensor();
-                    auto input_dims = in->getDimensions();
-                    auto expanded_size = args[1].unwrapToIntList();
-                    auto expandedDims = util::toDims(expanded_size);
-                    LOG_DEBUG("(expand layer) Expand input from " << input_dims << " to " << expandedDims);
-                    if(ctx->input_is_dynamic){
-                      int expanded_size_rank = static_cast<int>(expanded_size.size());
-                      int32_t* tmp = new int32_t[expanded_size_rank];
-                      for(int i=0;i<expanded_size_rank;i++)
-                        tmp[i] = expanded_size[i];
-                      auto expandedDimsTensor = vec2Tensor(tmp, expanded_size_rank, ctx);
-                      return add_expand_dynamic(ctx, n, in, expandedDimsTensor);
-                    }else{
-                      return add_expand(ctx, n, in, expandedDims);
-                    }
-                  }})
-        .pattern({"aten::expand_as(Tensor(a) self, Tensor other) -> (Tensor(a))",
-                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-                    auto in = args[0].ITensor();
-                    auto input_dims = in->getDimensions();
-                    auto targetTensor = args[1].ITensor();
-                    auto targetDims = targetTensor->getDimensions();
-                    LOG_DEBUG("(expand_as layer) Expand input from " << input_dims << " to " << targetDims);
-                    if(ctx->input_is_dynamic){
-                      return add_expand_dynamic(ctx, n, in, ctx->net->addShape(*targetTensor)->getOutput(0));
-                    }else{
-                      return add_expand(ctx, n, in, targetDims);
-                    }
-                    
-                  }})
-        .pattern({"aten::repeat(Tensor self, int[] repeats) -> (Tensor)",
-                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-                    auto in = args[0].ITensor();
-                    auto input_dims = in->getDimensions();
-                    auto repeats = args[1].unwrapToIntList().vec();
-                    int repeats_rank = repeats.size();
-                    TRTORCH_CHECK(
-                        repeats_rank >= input_dims.nbDims,
-                        "Number of repeat dimensions cannot be smaller than number of input dimensions");
-                    auto num_expand_dims = repeats_rank - input_dims.nbDims;
-
-                    if(ctx->input_is_dynamic){
-                      int input_rank = input_dims.nbDims;
-                      int output_rank= repeats_rank;
-                      auto new_input_shape_tensor = concat(output_rank, input_rank, ctx, in);
-
-                      // Add a reshape layer to expand dims
-                      auto shuffle = ctx->net->addShuffle(*in);
-                      shuffle->setInput(1, *new_input_shape_tensor);
-                      in = shuffle->getOutput(0);
-                    }else{
-                      if (num_expand_dims > 0) {
-                        nvinfer1::Dims reshape_dims;
-                        reshape_dims.nbDims = repeats.size();
-                        for (int i = 0; i < num_expand_dims; i++) {
-                          reshape_dims.d[i] = 1;
-                        }
-                        for (int i = 0; i < input_dims.nbDims; i++) {
-                          reshape_dims.d[num_expand_dims + i] = input_dims.d[i];
-                        }
-                        // Add a reshape layer to expand dims
-                        auto reshape_layer = ctx->net->addShuffle(*in);
-                        reshape_layer->setReshapeDimensions(reshape_dims);
-                        in = reshape_layer->getOutput(0);
-                        LOG_DEBUG("Input reshaped to : " << in->getDimensions() << " from " << input_dims);
-                      }
-                      LOG_DEBUG("Repeats: " << repeats);
-                    }
-
-                    // Concat across all repeat axes.
-                    // TODO: Implementation might not be performant. Explore other strategies to improve performance.
-                    for (int i = repeats.size() - 1; i >= 0; --i) {
-                      std::vector<nvinfer1::ITensor*> tensors_vec;
-                      for (int j = 0; j < repeats[i]; j++) {
-                        tensors_vec.push_back(in);
-                      }
-                      auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
-                      concat_layer->setAxis(i);
-                      in = concat_layer->getOutput(0);
-                    }
-
-                    auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
-
-                    LOG_DEBUG("Repeat layer output tensor shape: " << in->getDimensions());
-                    return true;
-                  }});
+        .pattern(
+            {"aten::expand(Tensor(a) self, int[] size, *, bool implicit=False) -> (Tensor(a))",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto in = args[0].ITensor();
+               auto input_dims = in->getDimensions();
+               auto expanded_size = args[1].unwrapToIntList();
+               auto expandedDims = util::toDims(expanded_size);
+               LOG_DEBUG("(expand layer) Expand input from " << input_dims << " to " << expandedDims);
+               if (ctx->input_is_dynamic) {
+                 int expanded_size_rank = static_cast<int>(expanded_size.size());
+                 int32_t* tmp = new int32_t[expanded_size_rank];
+                 for (int i = 0; i < expanded_size_rank; i++)
+                   tmp[i] = expanded_size[i];
+                 auto expandedDimsTensor = vec2Tensor(tmp, expanded_size_rank, ctx);
+                 return add_expand_dynamic(ctx, n, in, expandedDimsTensor);
+               } else {
+                 return add_expand(ctx, n, in, expandedDims);
+               }
+             }})
+        .pattern(
+            {"aten::expand_as(Tensor(a) self, Tensor other) -> (Tensor(a))",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto in = args[0].ITensor();
+               auto input_dims = in->getDimensions();
+               auto targetTensor = args[1].ITensor();
+               auto targetDims = targetTensor->getDimensions();
+               LOG_DEBUG("(expand_as layer) Expand input from " << input_dims << " to " << targetDims);
+               if (ctx->input_is_dynamic) {
+                 return add_expand_dynamic(ctx, n, in, ctx->net->addShape(*targetTensor)->getOutput(0));
+               } else {
+                 return add_expand(ctx, n, in, targetDims);
+               }
+             }})
+        .pattern(
+            {"aten::repeat(Tensor self, int[] repeats) -> (Tensor)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto in = args[0].ITensor();
+               auto input_dims = in->getDimensions();
+               auto repeats = args[1].unwrapToIntList().vec();
+               int repeats_rank = repeats.size();
+               TRTORCH_CHECK(
+                   repeats_rank >= input_dims.nbDims,
+                   "Number of repeat dimensions cannot be smaller than number of input dimensions");
+               auto num_expand_dims = repeats_rank - input_dims.nbDims;
+
+               if (ctx->input_is_dynamic) {
+                 int input_rank = input_dims.nbDims;
+                 int output_rank = repeats_rank;
+                 auto new_input_shape_tensor = concat(output_rank, input_rank, ctx, in);
+
+                 // Add a reshape layer to expand dims
+                 auto shuffle = ctx->net->addShuffle(*in);
+                 shuffle->setInput(1, *new_input_shape_tensor);
+                 in = shuffle->getOutput(0);
+               } else {
+                 if (num_expand_dims > 0) {
+                   nvinfer1::Dims reshape_dims;
+                   reshape_dims.nbDims = repeats.size();
+                   for (int i = 0; i < num_expand_dims; i++) {
+                     reshape_dims.d[i] = 1;
+                   }
+                   for (int i = 0; i < input_dims.nbDims; i++) {
+                     reshape_dims.d[num_expand_dims + i] = input_dims.d[i];
+                   }
+                   // Add a reshape layer to expand dims
+                   auto reshape_layer = ctx->net->addShuffle(*in);
+                   reshape_layer->setReshapeDimensions(reshape_dims);
+                   in = reshape_layer->getOutput(0);
+                   LOG_DEBUG("Input reshaped to : " << in->getDimensions() << " from " << input_dims);
+                 }
+                 LOG_DEBUG("Repeats: " << repeats);
+               }
+
+               // Concat across all repeat axes.
+               // TODO: Implementation might not be performant. Explore other strategies to improve performance.
+               for (int i = repeats.size() - 1; i >= 0; --i) {
+                 std::vector<nvinfer1::ITensor*> tensors_vec;
+                 for (int j = 0; j < repeats[i]; j++) {
+                   tensors_vec.push_back(in);
+                 }
+                 auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
+                 concat_layer->setAxis(i);
+                 in = concat_layer->getOutput(0);
+               }
+
+               auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
+
+               LOG_DEBUG("Repeat layer output tensor shape: " << in->getDimensions());
+               return true;
+             }});
 
 } // namespace
 } // namespace impl