prevent workspace from being bigger than anticipated

TuKJet · Nov 14, 2015 · 0e680e0 · 0e680e0
1 parent 168a713
commit 0e680e0
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/src/operator/convolution-inl.h b/src/operator/convolution-inl.h
@@ -240,10 +240,15 @@ class ConvolutionOp : public Operator {
                                      param_.num_filter / param_.num_group,
                                      oshape[2] * oshape[3]);
     const uint64_t workspace_size = param_.workspace;  // In elements of sizeof(real_t)
-    nstep_ = std::max(std::min(static_cast<index_t>(workspace_size / shape_colunit_.Size()),
-                               ishape[0]), 1U);
-    int nop = (ishape[0] + nstep_ - 1) / nstep_;
+    index_t nstep = std::max(std::min(static_cast<index_t>(workspace_size / shape_colunit_.Size()),
+                                      ishape[0]), 1U);
+    index_t nop = (ishape[0] + nstep - 1) / nstep;
     nstep_ = (ishape[0] + nop - 1) / nop;
+
+    if (nstep_ == nstep) {
+      nstep_ = std::max(nstep_ - 1, 1U);
+    }
+
     mshadow::Shape<2> scol = mshadow::Shape2(shape_colunit_[0],
                                              shape_colunit_[1] * nstep_);
     mshadow::Shape<3> sdst = mshadow::Shape3(shape_dstunit_[0],