fix alignment of pointers

apache · Aug 15, 2018 · 87b2275 · 87b2275
1 parent 5524e7b
commit 87b2275
Showing 1 changed file with 3 additions and 10 deletions.
diff --git a/src/operator/contrib/bounding_box-inl.h b/src/operator/contrib/bounding_box-inl.h
@@ -423,15 +423,8 @@ void BoxNMSForward(const nnvm::NodeAttrs& attrs,
       dtype_size += buffer_shape.Size();
     }
     // ceil up when sizeof(DType) is larger than sizeof(DType)
-    index_t workspace_size = (int32_size * sizeof(int32_t) - 1) / sizeof(DType) + 1 + dtype_size;
-    // workspace_size += workspace_size * 2;  // all_sorted_index, batch_id
-    // workspace_size += 2 * sort_index_shape.Size();  // scores, batch_id, areas
-    //
-    // // batch_start
-    // workspace_size += ((batch_start_shape.Size() * sizeof(int32_t) - 1) / sizeof(DType)) + 1;
-    // if (req[0] == kWriteInplace) {
-    //   workspace_size += buffer_shape.Size();
-    // }
+    index_t int32_offset = (int32_size * sizeof(int32_t) - 1) / sizeof(DType) + 1;
+    index_t workspace_size = int32_offset + dtype_size;
     Tensor<xpu, 1, DType> workspace = ctx.requested[box_nms_enum::kTempSpace]
       .get_space_typed<xpu, 1, DType>(Shape1(workspace_size), s);
     Tensor<xpu, 1, int32_t> sorted_index(
@@ -441,7 +434,7 @@ void BoxNMSForward(const nnvm::NodeAttrs& attrs,
     Tensor<xpu, 1, int32_t> batch_id(
       all_sorted_index.dptr_ + all_sorted_index.MSize(), sort_index_shape, s);
     Tensor<xpu, 1, int32_t> batch_start(batch_id.dptr_ + batch_id.MSize(), batch_start_shape, s);
-    Tensor<xpu, 1, DType> scores(reinterpret_cast<DType*>(batch_start.dptr_ + batch_start.MSize()),
+    Tensor<xpu, 1, DType> scores(workspace.dptr_ + int32_offset,
       sort_index_shape, s);
     Tensor<xpu, 1, DType> areas(scores.dptr_ + scores.MSize(), sort_index_shape, s);
     Tensor<xpu, 3, DType> buffer = data;