[NNPack] Fix automatically cast fail on some platforms (#388)

apache · Aug 26, 2017 · 46f141f · 46f141f
1 parent 26b433b
commit 46f141f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/src/contrib/nnpack/convolution.cc b/src/contrib/nnpack/convolution.cc
@@ -20,8 +20,10 @@ TVM_REGISTER_GLOBAL("tvm.contrib.nnpack.convolution_inference")
     DLTensor* kernel = args[1];
     DLTensor* bias   = args[2];
     DLTensor* output = args[3];
-    nnp_padding input_padding{args[4], args[5], args[6], args[7]};
-    nnp_size stride_size{args[8], args[9]};
+    uint64_t pad_top = args[4], pad_right = args[5], pad_bottom = args[6], pad_left = args[7];
+    nnp_padding input_padding{pad_top, pad_right, pad_bottom, pad_left};
+    uint64_t stride_width = args[8], stride_height = args[9];
+    nnp_size stride_size{stride_width, stride_height};
 
     CHECK_EQ(input->ndim, 3);
     CHECK_EQ(kernel->ndim, 4);
@@ -76,7 +78,8 @@ TVM_REGISTER_GLOBAL("tvm.contrib.nnpack.convolution_output")
     DLTensor* kernel = args[1];
     DLTensor* bias   = args[2];
     DLTensor* output = args[3];
-    nnp_padding input_padding{args[4], args[5], args[6], args[7]};
+    uint64_t pad_top = args[4], pad_right = args[5], pad_bottom = args[6], pad_left = args[7];
+    nnp_padding input_padding{pad_top, pad_right, pad_bottom, pad_left};
 
     CHECK_EQ(input->ndim, 4);
     CHECK_EQ(kernel->ndim, 4);

diff --git a/src/contrib/nnpack/nnpack_utils.cc b/src/contrib/nnpack/nnpack_utils.cc
@@ -17,7 +17,7 @@ NNPackThreadLocalEntry* NNPackThreadLocalEntry::ThreadLocal() {
 TVM_REGISTER_GLOBAL("contrib.nnpack._Config")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     NNPackThreadLocalEntry *entry = NNPackThreadLocalEntry::ThreadLocal();
-    size_t nthreads = args[0];
+    size_t nthreads = args[0].operator uint64_t();
     if (entry->threadpool != NULL &&
         pthreadpool_get_threads_count(entry->threadpool) != nthreads) {
       pthreadpool_destroy(entry->threadpool);