[BUG FIX] reformat based on comments

openvinotoolkit · Aug 5, 2024 · 79dfad7 · 79dfad7
1 parent e93b04e
commit 79dfad7
Showing 1 changed file with 13 additions and 13 deletions.
diff --git a/source/thead_rvv/fp32/prelu.c b/source/thead_rvv/fp32/prelu.c
@@ -74,18 +74,17 @@ int shl_rvv_prelu_fp32(struct csinn_tensor *input, struct csinn_tensor *alpha,
             }
         } else {
             // broadcast alpha
-            int new_alpha_size = csinn_tensor_size(input);
-            float *alpha_data_b = shl_mem_alloc(new_alpha_size * sizeof(float));
+            int input_size = csinn_tensor_size(input);
+            float *alpha_data_b = shl_mem_alloc(input_size * sizeof(float));
             struct csinn_tensor *alpha_ = csinn_alloc_tensor(NULL);
             csinn_tensor_copy(alpha_, input);
             alpha_->data = alpha_data_b;
             shl_ref_broadcast_to_shape_f32(alpha, alpha_, alpha_->dim, alpha_->dim_count);
             alpha_data = (float *)alpha_->data;
 
             // calculation
-            int ele_size = input->dim[0] * input->dim[1] * input->dim[2] * input->dim[3];
-            while (ele_size > 0) {
-                int vl = vsetvl_e32m2(ele_size);
+            while (input_size > 0) {
+                int vl = vsetvl_e32m2(input_size);
                 vfloat32m2_t _input = vle32_v_f32m2(input_data, vl);
                 vfloat32m2_t _a = vle32_v_f32m2(alpha_data, vl);
                 vbool16_t _mask = vmflt_vf_f32m2_b16(_input, 0.0f, vl);
@@ -94,20 +93,21 @@ int shl_rvv_prelu_fp32(struct csinn_tensor *input, struct csinn_tensor *alpha,
                 input_data += vl;
                 alpha_data += vl;
                 output_data += vl;
-                ele_size -= vl;
-            }
-            if (output->layout == CSINN_LAYOUT_NC1HWC0) {
-                const int packn = csrr_vlenb() / sizeof(float);
-                output->dim[1] *= packn;
-                output->dim[4] = 0;
-                output->dim_count = 4;
-                output->layout = CSINN_LAYOUT_NCHW;
+                input_size -= vl;
             }
 
             // free memory and tensor
             shl_mem_free(alpha_data_b);
             csinn_free_tensor(alpha_);
         }
+
+        if (output->layout == CSINN_LAYOUT_NC1HWC0) {
+            const int packn = csrr_vlenb() / sizeof(float);
+            output->dim[1] *= packn;
+            output->dim[4] = 0;
+            output->dim_count = 4;
+            output->layout = CSINN_LAYOUT_NCHW;
+        }
     } else {
         shl_debug_error("prelu unsupported layout: %d\n", input->layout);
     }