Fix small issue that FP16 is not loaded mmap directly.

liuliu · Mar 13, 2024 · a27fef3 · a27fef3
1 parent b6add38
commit a27fef3
Showing 1 changed file with 23 additions and 1 deletion.
diff --git a/nnc/Store.swift b/nnc/Store.swift
@@ -3430,6 +3430,28 @@ private let q8pDecodeJitWithExternalEager:
     return 1
   }
 
+private let decodeWithExternalEager:
+  @convention(c) (
+    UnsafeRawPointer?, Int, Int32, UnsafePointer<Int32>?, Int32, UInt32, UnsafeMutableRawPointer?,
+    ccv_nnc_tensor_param_t, UnsafeMutablePointer<UnsafeMutablePointer<ccv_nnc_tensor_t>?>?,
+    UnsafeMutableRawPointer?, UnsafeMutablePointer<Int>?
+  ) -> Int32 = {
+    data, dataSize, dataType, dimensions, dimensionCount, identifier, context, params, tensorOut,
+    decoded, decodedSize
+    in
+    guard let data = data, dimensionCount > 0 else { return 0 }
+    guard tensorOut!.pointee == nil else {
+      return decodeWithExternalStore(
+        data, dataSize, dataType, dimensions, dimensionCount, identifier, context, params,
+        tensorOut, decoded, decodedSize)
+    }
+    let store = Unmanaged<DynamicGraph._Store>.fromOpaque(context!).takeUnretainedValue()
+    let offset = Int(data.load(as: UInt64.self))
+    tensorOut!.pointee = ccv_nnc_tensor_new_from_file(
+      params, store.externalStore, off_t(offset), Int32(CCV_NNC_TENSOR_MEMORY_MAP_EAGER))
+    return 1
+  }
+
 private let uDecodeJitWithExternalStore:
   @convention(c) (
     UnsafeRawPointer?, Int, Int32, UnsafePointer<Int32>?, Int32, UInt32, UnsafeMutableRawPointer?,
@@ -3478,7 +3500,7 @@ private let uDecodeJitWithExternalStore:
         data, dataSize, dataType, dimensions, dimensionCount, identifier, context, params,
         tensorOut, decoded, decodedSize)
     default:
-      return decodeWithExternalStore(
+      return decodeWithExternalEager(
         data, dataSize, dataType, dimensions, dimensionCount, identifier, context, params,
         tensorOut, decoded, decodedSize)
     }