Merge "Fix Conv2D crash in 64bit binary." into pi-dev

5 files changed, 22 insertions, 40 deletions

diff --git a/nn/common/operations/Conv2D.cpp b/nn/common/operations/Conv2D.cpp
index 3787c0502..344ab8feb 100644
--- a/nn/common/operations/Conv2D.cpp
+++ b/nn/common/operations/Conv2D.cpp
@@ -55,8 +55,8 @@ static char static_scratch_buffer[kStaticBufferSize];
     for (int i=0; i<4; i++) {                                                   \
         im2colByteSize *= im2colDim.sizes[i];                                   \
     }                                                                           \
-                                                                                \
-    if (sizeof(size_t) == 4 && (im2colByteSize / sizeof(Type)) > 0xFFFFFFFF)  { \
+    /* http://b/77982879, tflite::optimized_ops::Conv uses int for offsets */   \
+    if (im2colByteSize >= 0x7fffffff)  {                                        \
         LOG(ERROR) << "Conv size is too large, not enough memory";              \
         return false;                                                           \
     }                                                                           \
diff --git a/nn/runtime/test/Android.bp b/nn/runtime/test/Android.bp
index 0f652124d..6de26e320 100644
--- a/nn/runtime/test/Android.bp
+++ b/nn/runtime/test/Android.bp
@@ -58,7 +58,10 @@ cc_test {
     name: "NeuralNetworksTest_shared_partial",
     defaults: ["NeuralNetworksTest_defaults"],
     srcs: [
-        "TestMainPartial.cpp",
+        "TestMain.cpp",
+    ],
+    cflags: [
+        "-DNNTEST_ONLY_PUBLIC_API"
     ],
     shared_libs: [
         "libneuralnetworks",
diff --git a/nn/runtime/test/TestMain.cpp b/nn/runtime/test/TestMain.cpp
index 358c5b4c3..7409c9376 100644
--- a/nn/runtime/test/TestMain.cpp
+++ b/nn/runtime/test/TestMain.cpp
@@ -28,13 +28,14 @@ int main(int argc, char** argv) {
     android::nn::initVLogMask();
     // Test with the installed drivers.
     int n1 = RUN_ALL_TESTS();
-
-#ifdef TESTMAINPARTIAL
+#ifdef NNTEST_ONLY_PUBLIC_API
+    // Can't use non-public functionality, because we're linking against
+    // the shared library version of the runtime.
     return n1;
 #else
     // Test with the CPU driver only.
     android::nn::DeviceManager::get()->setUseCpuOnly(true);
     int n2 = RUN_ALL_TESTS();
     return n1 | n2;
-#endif
+#endif // NNTEST_ONLY_PUBLIC_API
 }
diff --git a/nn/runtime/test/TestMainPartial.cpp b/nn/runtime/test/TestMainPartial.cpp
deleted file mode 100644
index a15f71fd6..000000000
--- a/nn/runtime/test/TestMainPartial.cpp
+++ /dev/null
@@ -1,21 +0,0 @@
-/*
- * Copyright (C) 2018 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-// Can't use non-public functionality, because we're linking against
-// the shared library version of the runtime.
-#define TESTMAINPARTIAL
-
-#include "TestMain.cpp"
diff --git a/nn/runtime/test/TestMemoryInternal.cpp b/nn/runtime/test/TestMemoryInternal.cpp
index 766829159..5368ac251 100644
--- a/nn/runtime/test/TestMemoryInternal.cpp
+++ b/nn/runtime/test/TestMemoryInternal.cpp
@@ -21,6 +21,7 @@
 #include "TestMemory.h"
 
 #include "NeuralNetworksWrapper.h"
+#include "Manager.h"
 #include "Memory.h"
 
 #include <android/sharedmem.h>
@@ -199,16 +200,19 @@ TEST_F(MemoryLeakTest, IterativelyInstantiate) {
     }
 }
 
+#ifndef NNTEST_ONLY_PUBLIC_API
 // Regression test for http://b/73663843, conv_2d trying to allocate too much memory.
 TEST_F(MemoryLeakTest, convTooLarge) {
+    android::nn::DeviceManager::get()->setUseCpuOnly(true);
     WrapperModel model;
 
-    // This kernel/input size will make convQuant8 allocate 40 * 13 * 13 * 128 * 92 * 92 ~= 7G
-    // This will alloways fail on 32 bit binaries and probably fail on 64bit devices.
+    // This kernel/input size will make convQuant8 allocate 12 * 13 * 13 * 128 * 92 * 92, which is
+    // just outside of signed int range (0x82F56000) - this will fail due to CPU implementation
+    // limitations
     WrapperOperandType type3(WrapperType::INT32, {});
     WrapperOperandType type2(WrapperType::TENSOR_INT32, {128}, 0.25, 0);
-    WrapperOperandType type0(WrapperType::TENSOR_QUANT8_ASYMM, {40, 104, 104, 128}, 0.5, 0);
-    WrapperOperandType type4(WrapperType::TENSOR_QUANT8_ASYMM, {40, 92, 92, 128}, 1.0, 0);
+    WrapperOperandType type0(WrapperType::TENSOR_QUANT8_ASYMM, {12, 104, 104, 128}, 0.5, 0);
+    WrapperOperandType type4(WrapperType::TENSOR_QUANT8_ASYMM, {12, 92, 92, 128}, 1.0, 0);
     WrapperOperandType type1(WrapperType::TENSOR_QUANT8_ASYMM, {128, 13, 13, 128}, 0.5, 0);
 
     // Operands
@@ -244,21 +248,16 @@ TEST_F(MemoryLeakTest, convTooLarge) {
     WrapperExecution execution(&compilation);
 
     // Set input and outputs
-    static uint8_t input[40 * 104 * 104 * 128] = {};
+    static uint8_t input[12 * 104 * 104 * 128] = {};
     ASSERT_EQ(WrapperResult::NO_ERROR, execution.setInput(0, input, sizeof(input)));
-    static uint8_t output[40 * 92 * 92 * 128] = {};
+    static uint8_t output[12 * 92 * 92 * 128] = {};
     ASSERT_EQ(WrapperResult::NO_ERROR, execution.setOutput(0, output, sizeof(output)));
 
     // This shouldn't segfault
     WrapperResult r = execution.compute();
 
-    if (sizeof(size_t) == 4) {
-      // 32 bit binary
-      ASSERT_EQ(WrapperResult::OP_FAILED, r);
-    } else {
-      // 64 bit binary
-      ASSERT_TRUE((WrapperResult::OP_FAILED == r) || (WrapperResult::NO_ERROR == r));
-    }
+    ASSERT_EQ(WrapperResult::OP_FAILED, r);
 }
+#endif // NNTEST_ONLY_PUBLIC_API
 
 }  // end namespace