clean transpose axis setup and sort. and add comments.

inference-engine/src/plaidml_plugin/ops/space_to_batch.cpp

@@ -26,7 +26,11 @@ static OpRegistration reg("SpaceToBatch", [](const Context& ctx) {
     auto crops_begin = get_coords_from_constant_operand(2, ctx.layer);
     auto crops_end = get_coords_from_constant_operand(3, ctx.layer);
 
-    // padding the original Tensor.
+    /** According to openvino op spec, follow below operate step:
+     *
+     *  Zero-pad the start and end of dimensions [D_0, ..., D_{N - 1}] of the input according to `pads_begin` and `pads_end`:
+     *  x = [batch + P_0, D_1 + P_1, D_2 + P_2, ..., D_{N - 1} + P_{N - 1}], where P_i = pads_begin[i] + pads_end[i]
+     **/
     std::vector<int> lo_pads;
     for (auto pad : crops_begin) {
         lo_pads.push_back(pad);
@@ -40,7 +44,10 @@ static OpRegistration reg("SpaceToBatch", [](const Context& ctx) {
     }
     edsl::Tensor padding_I = op::explicit_padding(I, lo_pads, hi_pads).padval(edsl::Constant(0));
 
-    // reshape input tensor dim.
+    /** reshape padding tensor
+     *  x' = reshape(x, [batch, (D_1 + P_1) / B_1, B_1, (D_2 + P_2) / B_2, B_2, ..., (D_{N - 1} + P_{N - 1}) / B_{N - 1},
+     *                  B_{N - 1}]), where B_i = block_shape[i]
+     */
     std::vector<edsl::TensorDim> I_dims(padding_I.rank());
     padding_I.bind_dims(I_dims);
 
@@ -52,31 +59,28 @@ static OpRegistration reg("SpaceToBatch", [](const Context& ctx) {
     }
     auto reshape_I = edsl::reshape(padding_I, temp_dims);
 
-    /** transpose padding tensor, According to openvino op spec.
-     x'' = transpose(x',  [2, 4, ..., (N - 1) + (N - 1), 0, 1, 3, ..., N + (N - 1)])**/
-
-    // setup the dims
+    /** transpose reshape tensor.
+     *  x'' = transpose(x',  [2, 4, ..., (N - 1) + (N - 1), 0, 1, 3, ..., N + (N - 1)])
+     **/
+    // setup the sort dims
     std::vector<size_t> reshape_dims(reshape_I.rank());
     auto dims_size = reshape_dims.size();
-    for (size_t i = 0; i < dims_size; i++) {
-        reshape_dims[i] = i;
-    }
-    // sort the dims.
-    for (int i = dims_size - 2; i > 0; i = i - 2) {
-        for (int j = 0; j < (dims_size - i) / 2; j++) {
-            std::swap(reshape_dims[i + j], reshape_dims[i + j + 1]);
-        }
-    }
-    for (size_t i = 0; i < dims_size / 2; i++) {
-        std::swap(reshape_dims[i], reshape_dims[i + 1]);
+    auto mid_dim = dims_size / 2;
+    reshape_dims[mid_dim] = 0;
+    for (size_t i = 0; i < mid_dim; i++) {
+        reshape_dims[i] = 2 * (i + 1);
+        reshape_dims[i + mid_dim + 1] = 2 * i + 1;
     }
     std::vector<edsl::Value> dims_wrapper;
     for (auto dim : reshape_dims) {
         dims_wrapper.emplace_back(dim);
     }
     auto transpose_I = op::transpose(reshape_I, edsl::Value(dims_wrapper));
 
-    // final reshape.
+    /** reshape to the final tensor.
+     *  y = reshape(x'', [batch * B_1 * ... * B_{N - 1}, (D_1 + P_1) / B_1, (D_2 + P_2) / B_2, ... ,
+     *              (D_{N - 1} + P_{N - 1}) / B_{N - 1}])
+     */
     size_t total_block_size = 1;
     for (size_t i = 0; i < block_shape.size(); i++) {
         total_block_size *= block_shape[i];