diff --git a/cpp/external/katagocoreml/src/Converter.cpp b/cpp/external/katagocoreml/src/Converter.cpp index e0a01a348..377ecaff7 100644 --- a/cpp/external/katagocoreml/src/Converter.cpp +++ b/cpp/external/katagocoreml/src/Converter.cpp @@ -29,9 +29,12 @@ void KataGoConverter::convert(const std::string& input_path, throw std::invalid_argument("max_batch_size must be >= min_batch_size or <= 0 for unlimited"); } - // Parse KataGo model - KataGoParser parser(input_path); - KataGoModelDesc model = parser.parse(); + // Parse KataGo model (parser + its decompressed buffer freed at end of scope) + KataGoModelDesc model; + { + KataGoParser parser(input_path); + model = parser.parse(); + } // Determine if using FP16 precision bool use_fp16 = (options.compute_precision == "FLOAT16"); @@ -52,9 +55,8 @@ void KataGoConverter::convert(const std::string& input_path, options.use_fp16_io); auto program = builder.build(); - // Get weights from builder - auto weights = builder.getWeights(); - std::vector weights_copy(weights.begin(), weights.end()); + // Serialize directly from the builder's weight views (no copy). + std::vector& weights = builder.getWeightsMutable(); // Update options with model metadata for serialization ConversionOptions final_options = options; @@ -90,7 +92,7 @@ void KataGoConverter::convert(const std::string& input_path, // Serialize to .mlpackage CoreMLSerializer serializer(final_options.specification_version); - serializer.serialize(program.get(), weights_copy, output_path, final_options); + serializer.serialize(program.get(), weights, output_path, final_options); } bool KataGoConverter::wouldBuildFullyFp32(int trunk_num_channels, bool has_transformer_blocks) { diff --git a/cpp/external/katagocoreml/src/builder/MILBuilder.cpp b/cpp/external/katagocoreml/src/builder/MILBuilder.cpp index b2e2f6bb7..9540e5494 100644 --- a/cpp/external/katagocoreml/src/builder/MILBuilder.cpp +++ b/cpp/external/katagocoreml/src/builder/MILBuilder.cpp @@ -280,11 +280,30 @@ void MILBuilder::addConstOp(CoreML::Specification::MILSpec::Block* block, const std::string& name, const std::vector& data, const std::vector& shape) { - // Register weight for blob storage. Mark FP32 storage when this const is declared FP32 (e.g. - // inside an FP32 sub-region of an otherwise-FP16 model) so storage matches the declared type. + // Register weight for blob storage (non-owning view into the model). Mark FP32 storage when this + // const is declared FP32 (e.g. inside an FP32 sub-region of an otherwise-FP16 model) so storage + // matches the declared type. m_ops.registerWeight(name, data, shape, m_weight_dtype == CoreML::Specification::MILSpec::DataType::FLOAT32); + emitConstOp(block, name, shape); +} + +void MILBuilder::addOwnedConstOp(CoreML::Specification::MILSpec::Block* block, + const std::string& name, + std::vector&& data, + const std::vector& shape) { + // Register derived/owned weight. Mirror addConstOp's per-weight FP32 marking: emitConstOp + // declares this const's dtype as m_weight_dtype, so the stored bytes must follow the same flag + // or BNNS rejects the model ("Metadata data type does not match requested type") when a derived + // const lands in an FP32 sub-region of an FP16 model. + const bool is_fp32 = (m_weight_dtype == CoreML::Specification::MILSpec::DataType::FLOAT32); + m_ops.registerOwnedWeight(name, std::move(data), shape, is_fp32); + emitConstOp(block, name, shape); +} +void MILBuilder::emitConstOp(CoreML::Specification::MILSpec::Block* block, + const std::string& name, + const std::vector& shape) { // Add const operation auto* op = block->add_operations(); op->set_type("const"); @@ -1254,7 +1273,7 @@ void MILBuilder::addLinearOp(CoreML::Specification::MILSpec::Block* block, // Add transposed weight constant with shape [out_channels, in_channels] std::vector transposed_shape = {static_cast(out_ch), static_cast(in_ch)}; - addConstOp(block, weight_name, transposed_weights, transposed_shape); + addOwnedConstOp(block, weight_name, std::move(transposed_weights), transposed_shape); // Add bias constant std::vector bias_shape = {static_cast(bias.num_channels)}; @@ -2291,10 +2310,13 @@ std::string MILBuilder::buildTransformerAttentionBlock(CoreML::Specification::MI std::string cosName = prefix + "_" + tag + "_cos"; std::string sinName = prefix + "_" + tag + "_sin"; std::string rName = prefix + "_" + tag + "_R"; - addConstOp(block, cosName, cosFull, {1, nh, seq, qHeadDim}); - addConstOp(block, sinName, sinFull, {1, nh, seq, qHeadDim}); + // cosFull/sinFull/R are locals computed here, so register them as OWNED consts: the + // WeightEntry holds a non-owning FloatView and serialization runs after this lambda + // returns, so a non-owning addConstOp would dangle. + addOwnedConstOp(block, cosName, std::move(cosFull), {1, nh, seq, qHeadDim}); + addOwnedConstOp(block, sinName, std::move(sinFull), {1, nh, seq, qHeadDim}); // Rank-4 [1,1,qd,qd] so matmul batch dims broadcast cleanly against [B,nh,seq,qd]. - addConstOp(block, rName, R, {1, 1, qHeadDim, qHeadDim}); + addOwnedConstOp(block, rName, std::move(R), {1, 1, qHeadDim, qHeadDim}); std::string rotated = genVarName(prefix + "_" + tag + "_rot"); matmul(x, rName, rotated, {-1, nh, seq, qHeadDim}, false, false); std::string xc = genVarName(prefix + "_" + tag + "_xc"); @@ -2450,7 +2472,9 @@ std::string MILBuilder::buildTransformerAttentionBlock(CoreML::Specification::MI for (int d = 0; d < vHeadDim; d++) for (int c = 0; c < outC; c++) whData[d * outC + c] = desc.out_proj.weights[static_cast(h * vHeadDim + d) * outC + c]; - addConstOp(block, wh, whData, {vHeadDim, outC}); + // whData is a per-head local slice; register OWNED so its FloatView stays valid until + // serialization (a non-owning addConstOp would dangle after this loop iteration). + addOwnedConstOp(block, wh, std::move(whData), {vHeadDim, outC}); std::string contrib = genVarName(prefix + "_contrib"); matmul(aoh2d, wh, contrib, {-1, outC}, false, false); if (h == 0) { diff --git a/cpp/external/katagocoreml/src/builder/MILBuilder.hpp b/cpp/external/katagocoreml/src/builder/MILBuilder.hpp index 72cda533e..992f393a0 100644 --- a/cpp/external/katagocoreml/src/builder/MILBuilder.hpp +++ b/cpp/external/katagocoreml/src/builder/MILBuilder.hpp @@ -29,8 +29,8 @@ class MILBuilder { /// @return Unique pointer to MIL Program protobuf std::unique_ptr build(); - /// Get weight entries for blob serialization - const std::vector& getWeights() const { return m_ops.getWeights(); } + /// Get weight entries for blob serialization (mutable; serialization sets blob_offset) + std::vector& getWeightsMutable() { return m_ops.getWeightsMutable(); } /// Get board dimensions int getBoardXSize() const { return m_board_x_size; } @@ -116,6 +116,24 @@ class MILBuilder { const std::vector& data, const std::vector& shape); + // addConstOp registers a NON-OWNING view into `data` (see WeightEntry), so the + // backing storage must outlive serialization. Binding a temporary here would + // dangle. Deleted so such calls fail to compile; use addOwnedConstOp for + // derived/temporary tensors that KataGoOps should own instead. + void addConstOp(CoreML::Specification::MILSpec::Block* block, + const std::string& name, + std::vector&& data, + const std::vector& shape) = delete; + + void addOwnedConstOp(CoreML::Specification::MILSpec::Block* block, + const std::string& name, + std::vector&& data, + const std::vector& shape); + + void emitConstOp(CoreML::Specification::MILSpec::Block* block, + const std::string& name, + const std::vector& shape); + void addIntArrayConstOp(CoreML::Specification::MILSpec::Block* block, const std::string& name, const std::vector& values); diff --git a/cpp/external/katagocoreml/src/builder/Operations.cpp b/cpp/external/katagocoreml/src/builder/Operations.cpp index 1c625acdd..e86364943 100644 --- a/cpp/external/katagocoreml/src/builder/Operations.cpp +++ b/cpp/external/katagocoreml/src/builder/Operations.cpp @@ -18,7 +18,7 @@ std::string KataGoOps::registerWeight(const std::string& name, bool is_fp32) { WeightEntry entry; entry.name = name; - entry.data = data; + entry.data = FloatView{data.data(), data.size()}; entry.shape = shape; entry.blob_offset = 0; // Will be set during serialization entry.is_fp32 = is_fp32; @@ -26,6 +26,22 @@ std::string KataGoOps::registerWeight(const std::string& name, return name; } +std::string KataGoOps::registerOwnedWeight(const std::string& name, + std::vector&& data, + const std::vector& shape, + bool is_fp32) { + m_owned.push_back(std::move(data)); + const std::vector& stored = m_owned.back(); + WeightEntry entry; + entry.name = name; + entry.data = FloatView{stored.data(), stored.size()}; + entry.shape = shape; + entry.blob_offset = 0; + entry.is_fp32 = is_fp32; + m_weights.push_back(std::move(entry)); + return name; +} + std::string KataGoOps::genOpName(const std::string& prefix) { return prefix + "_" + std::to_string(m_op_counter++); } diff --git a/cpp/external/katagocoreml/src/builder/Operations.hpp b/cpp/external/katagocoreml/src/builder/Operations.hpp index a9d2a1466..385648d19 100644 --- a/cpp/external/katagocoreml/src/builder/Operations.hpp +++ b/cpp/external/katagocoreml/src/builder/Operations.hpp @@ -5,15 +5,29 @@ #include "../types/KataGoTypes.hpp" #include +#include #include #include namespace katagocoreml { -/// Weight entry for blob file storage +/// Minimal non-owning view over a contiguous float buffer. KataGo-local on +/// purpose: keeps the MILBlob dependency out of this header (conversion to +/// MILBlob::Util::Span happens only at the serializer boundary). +struct FloatView { + const float* ptr = nullptr; + size_t len = 0; + const float* data() const { return ptr; } + size_t size() const { return len; } + bool empty() const { return len == 0; } + float operator[](size_t i) const { return ptr[i]; } +}; + +/// Weight entry for blob file storage. `data` is a NON-OWNING view into the live +/// KataGoModelDesc (or into KataGoOps::m_owned for derived tensors). struct WeightEntry { std::string name; - std::vector data; + FloatView data; // non-owning view (replaces raw ptr + count) std::vector shape; uint64_t blob_offset = 0; // Set during serialization bool is_fp32 = false; // Store as FP32 (set when the const was declared FP32, e.g. inside an @@ -53,17 +67,33 @@ class KataGoOps { /// Get precomputed mask constants const MaskConstants& getMaskConstants() const { return m_mask_constants; } - /// Register a weight tensor and return its reference name. is_fp32 marks it for FP32 storage. + /// Register a weight that lives in the model (stored as a non-owning view). + /// is_fp32 marks it for FP32 storage. std::string registerWeight(const std::string& name, const std::vector& data, const std::vector& shape, bool is_fp32 = false); - /// Get all registered weights - const std::vector& getWeights() const { return m_weights; } + /// The stored WeightEntry is a non-owning view into `data`, so a temporary + /// would leave it dangling. Deleted to reject such calls at compile time; + /// use registerOwnedWeight for tensors KataGoOps should own. + std::string registerWeight(const std::string& name, + std::vector&& data, + const std::vector& shape) = delete; + + /// Register a derived/temporary weight; KataGoOps takes ownership so the + /// view stays valid through serialization. is_fp32 marks it for FP32 storage + /// (mirrors registerWeight) so the stored dtype matches the declared const dtype. + std::string registerOwnedWeight(const std::string& name, + std::vector&& data, + const std::vector& shape, + bool is_fp32 = false); + + /// Get all registered weights (mutable; serialization sets blob_offset) + std::vector& getWeightsMutable() { return m_weights; } - /// Clear all registered weights - void clearWeights() { m_weights.clear(); } + /// Clear all registered weights (and their owned backing buffers) + void clearWeights() { m_weights.clear(); m_owned.clear(); } /// Generate unique operation name std::string genOpName(const std::string& prefix); @@ -74,6 +104,7 @@ class KataGoOps { bool m_optimize_identity_mask; MaskConstants m_mask_constants; std::vector m_weights; + std::deque> m_owned; int m_op_counter = 0; }; diff --git a/cpp/external/katagocoreml/src/parser/KataGoParser.cpp b/cpp/external/katagocoreml/src/parser/KataGoParser.cpp index 921074e0c..43b09a0ff 100644 --- a/cpp/external/katagocoreml/src/parser/KataGoParser.cpp +++ b/cpp/external/katagocoreml/src/parser/KataGoParser.cpp @@ -5,7 +5,6 @@ #include #include #include -#include #include #include #include @@ -42,54 +41,41 @@ bool KataGoParser::isVersionSupported(int version) { } // ============================================================================ -// File Loading +// Stream Primitives // ============================================================================ -void KataGoParser::loadFile() { - // Check if gzip compressed - bool is_gzip = false; - if (m_model_path.size() >= 3) { - std::string ext = m_model_path.substr(m_model_path.size() - 3); - is_gzip = (ext == ".gz"); - } - - if (is_gzip) { - // Read gzipped file - gzFile gz = gzopen(m_model_path.c_str(), "rb"); - if (!gz) { - throw std::runtime_error("Cannot open gzip file: " + m_model_path); - } - - // Read in chunks - m_buffer.clear(); - std::vector chunk(1024 * 1024); // 1MB chunks - int bytes_read; - while ((bytes_read = gzread(gz, chunk.data(), static_cast(chunk.size()))) > 0) { - m_buffer.insert(m_buffer.end(), chunk.begin(), chunk.begin() + bytes_read); - } - - if (bytes_read < 0) { - int errnum; - const char* errmsg = gzerror(gz, &errnum); - gzclose(gz); - throw std::runtime_error("Error reading gzip file: " + std::string(errmsg)); - } - - gzclose(gz); - } else { - // Read regular file - std::ifstream file(m_model_path, std::ios::binary | std::ios::ate); - if (!file) { - throw std::runtime_error("Cannot open file: " + m_model_path); - } +bool KataGoParser::refill() { + if(!m_gz) return false; + int n = gzread(m_gz.get(), m_refill.data(), (unsigned)m_refill.size()); + if(n < 0) { + int errnum; + const char* errmsg = gzerror(m_gz.get(), &errnum); + throw std::runtime_error("Error reading gzip stream: " + std::string(errmsg)); + } + m_refillPos = 0; + m_refillLen = (size_t)n; + return n > 0; +} - std::streamsize size = file.tellg(); - file.seekg(0, std::ios::beg); +int KataGoParser::peekByte() { + if(m_refillPos >= m_refillLen) { + if(!refill()) return -1; + } + return (int)m_refill[m_refillPos]; +} - m_buffer.resize(static_cast(size)); - if (!file.read(reinterpret_cast(m_buffer.data()), size)) { - throw std::runtime_error("Error reading file: " + m_model_path); +void KataGoParser::readExact(uint8_t* dst, size_t n, const std::string& name) { + size_t got = 0; + while(got < n) { + if(m_refillPos >= m_refillLen) { + if(!refill()) + throw std::runtime_error(name + ": unexpected EOF in binary block"); } + size_t avail = m_refillLen - m_refillPos; + size_t take = std::min(avail, n - got); + std::memcpy(dst + got, m_refill.data() + m_refillPos, take); + m_refillPos += take; + got += take; } } @@ -98,15 +84,16 @@ void KataGoParser::loadFile() { // ============================================================================ KataGoModelDesc KataGoParser::parse() { - loadFile(); - m_pos = 0; - - // Detect if binary format (check for @BIN@ marker) - const std::string bin_marker = "@BIN@"; - auto it = std::search(m_buffer.begin(), m_buffer.end(), - bin_marker.begin(), bin_marker.end()); - m_binary_floats = (it != m_buffer.end()); - + // Allocate the refill buffer first; if this throws, no handle has been opened. + m_refill.resize(1024 * 1024); + m_gz.reset(gzopen(m_model_path.c_str(), "rb")); + if(!m_gz) + throw std::runtime_error("Cannot open file: " + m_model_path); + m_refillPos = 0; + m_refillLen = 0; + m_formatDetected = false; // decided at first readFloats + m_binary_floats = true; + // ~GzHandle closes the file on normal return OR exception — no try/catch needed. return parseModel(); } @@ -115,24 +102,20 @@ KataGoModelDesc KataGoParser::parse() { // ============================================================================ void KataGoParser::skipWhitespace() { - while (m_pos < m_buffer.size()) { - char c = static_cast(m_buffer[m_pos]); - if (c != ' ' && c != '\t' && c != '\n' && c != '\r') { - break; - } - m_pos++; + int c; + while((c = peekByte()) >= 0) { + if(c != ' ' && c != '\t' && c != '\n' && c != '\r') break; + m_refillPos++; } } void KataGoParser::readUntilWhitespace(std::string& out) { out.clear(); - while (m_pos < m_buffer.size()) { - char c = static_cast(m_buffer[m_pos]); - if (c == ' ' || c == '\t' || c == '\n' || c == '\r') { - break; - } - out += c; - m_pos++; + int c; + while((c = peekByte()) >= 0) { + if(c == ' ' || c == '\t' || c == '\n' || c == '\r') break; + out += (char)c; + m_refillPos++; } } @@ -170,55 +153,29 @@ bool KataGoParser::readBool() { } std::vector KataGoParser::readFloats(size_t count, const std::string& name) { - // Bound count by the remaining file size BEFORE allocating: text floats need at least - // 2 bytes each and binary exactly 4, so a count exceeding the remaining byte count is - // impossible. Without this, a crafted file (gzip compresses multi-GB of zeros to a few - // MB) could trigger an enormous zero-initializing allocation before any read fails. - if (count > m_buffer.size() - m_pos) { - throw std::runtime_error(name + ": not enough bytes for " + std::to_string(count) + " floats"); - } std::vector floats(count); + skipWhitespace(); + + // KataGo model files are uniformly text OR uniformly binary, so detecting the + // format once at the first weight block (binary blocks start with '@BIN@') + // is valid for all subsequent blocks. + if(!m_formatDetected) { + m_binary_floats = (peekByte() == '@'); + m_formatDetected = true; + } - if (!m_binary_floats) { + if(!m_binary_floats) { // Text format - for (size_t i = 0; i < count; i++) { + for(size_t i = 0; i < count; i++) floats[i] = readFloat(); - } } else { - // Binary format - find @BIN@ marker. Bound the scan like desc.cpp (which allows at - // most 100 chars before the marker): an unbounded scan would silently skip arbitrary - // junk between weight blocks, masking corruption or resyncing past a real error. - size_t scan_start = m_pos; - while (m_pos < m_buffer.size()) { - if (m_buffer[m_pos] == '@') { - break; - } - if (m_pos - scan_start >= 100) { - throw std::runtime_error(name + ": could not find @BIN@ marker near expected position." - " Invalid model - perhaps a .txt model with a binary header, or corrupted data?"); - } - m_pos++; - } - - // Check for @BIN@ header - if (m_pos + 5 > m_buffer.size() || - std::memcmp(&m_buffer[m_pos], "@BIN@", 5) != 0) { + // Binary: consume the "@BIN@" marker, then read count*4 raw bytes. + char marker[5]; + readExact(reinterpret_cast(marker), 5, name); + if(std::memcmp(marker, "@BIN@", 5) != 0) throw std::runtime_error(name + ": expected @BIN@ marker for binary float block"); - } - m_pos += 5; - - // Read binary floats (little-endian). Compare count against the remaining bytes - // directly so a huge count cannot overflow the num_bytes computation and slip - // past this bounds check. - size_t remaining = m_buffer.size() - m_pos; - if (count > remaining / 4) { - throw std::runtime_error(name + ": not enough bytes for " + std::to_string(count) + " floats"); - } - size_t num_bytes = count * 4; - // Copy as little-endian float32 - std::memcpy(floats.data(), &m_buffer[m_pos], num_bytes); - m_pos += num_bytes; + readExact(reinterpret_cast(floats.data()), count * 4, name); } // Reject NaN/Inf weights: corrupted or otherwise invalid models would diff --git a/cpp/external/katagocoreml/src/parser/KataGoParser.hpp b/cpp/external/katagocoreml/src/parser/KataGoParser.hpp index 404ca4862..7ed423fac 100644 --- a/cpp/external/katagocoreml/src/parser/KataGoParser.hpp +++ b/cpp/external/katagocoreml/src/parser/KataGoParser.hpp @@ -5,8 +5,11 @@ #include "../types/KataGoTypes.hpp" #include +#include #include +#include #include +#include namespace katagocoreml { @@ -35,9 +38,23 @@ class KataGoParser { private: std::string m_model_path; - std::vector m_buffer; - size_t m_pos = 0; + // Custom-deleter unique_ptr owns the gzFile so it closes on every exit path + // (normal return, exception, or bad_alloc) without manual try/catch. + struct GzCloser { + void operator()(gzFile f) const noexcept { if(f) gzclose(f); } + }; + using GzHandle = std::unique_ptr, GzCloser>; + GzHandle m_gz; + std::vector m_refill; // bounded refill buffer (~1 MB) + size_t m_refillPos = 0; // read cursor within m_refill + size_t m_refillLen = 0; // valid bytes in m_refill bool m_binary_floats = true; + bool m_formatDetected = false; + + // Stream primitives + bool refill(); // returns false at EOF + int peekByte(); // -1 at EOF + void readExact(uint8_t* dst, size_t n, const std::string& name); // Low-level reading functions void readUntilWhitespace(std::string& out); @@ -73,9 +90,6 @@ class KataGoParser { // Main model parsing KataGoModelDesc parseModel(); - - // Helper to load file (handles gzip) - void loadFile(); }; } // namespace katagocoreml diff --git a/cpp/external/katagocoreml/src/serializer/CoreMLSerializer.cpp b/cpp/external/katagocoreml/src/serializer/CoreMLSerializer.cpp index f271f5526..50df8003f 100644 --- a/cpp/external/katagocoreml/src/serializer/CoreMLSerializer.cpp +++ b/cpp/external/katagocoreml/src/serializer/CoreMLSerializer.cpp @@ -12,6 +12,8 @@ #include #include #include +#include +#include namespace katagocoreml { @@ -230,8 +232,13 @@ void CoreMLSerializer::createPackage(const std::string& output_path, if (!out) { throw std::runtime_error("Failed to create temp model file"); } - if (!model->SerializeToOstream(&out)) { - throw std::runtime_error("Failed to serialize model spec"); + { + google::protobuf::io::OstreamOutputStream zos(&out); + google::protobuf::io::CodedOutputStream cos(&zos); + cos.SetSerializationDeterministic(true); + if (!model->SerializeToCodedStream(&cos)) { + throw std::runtime_error("Failed to serialize model spec"); + } } } diff --git a/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp b/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp index e8fe861c8..69d590609 100644 --- a/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp +++ b/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp @@ -15,24 +15,25 @@ size_t WeightSerializer::serialize(std::vector& weights, size_t total_bytes = 0; for (auto& entry : weights) { + const size_t count = entry.data.size(); // Per-weight precision: store FP16 only when the global mode is FP16 AND this weight was not // declared FP32 (entry.is_fp32 marks consts inside an FP32 sub-region of an FP16 model), so // stored bytes stay consistent with each const's declared dtype. const bool store_fp16 = use_fp16 && !entry.is_fp32; if (store_fp16) { // Convert FP32 weights to FP16 - std::vector fp16_data(entry.data.size()); - for (size_t i = 0; i < entry.data.size(); ++i) { + std::vector fp16_data(count); + for (size_t i = 0; i < count; ++i) { fp16_data[i] = MILBlob::Fp16::FromFloat(entry.data[i]); } MILBlob::Util::Span span(fp16_data.data(), fp16_data.size()); entry.blob_offset = writer.WriteData(span); - total_bytes += entry.data.size() * sizeof(MILBlob::Fp16); + total_bytes += count * sizeof(MILBlob::Fp16); } else { - // Write FP32 weights - MILBlob::Util::Span span(entry.data.data(), entry.data.size()); + // Write FP32 weights — convert the KataGo-local view to a MILBlob span here. + MILBlob::Util::Span span(entry.data.data(), count); entry.blob_offset = writer.WriteData(span); - total_bytes += entry.data.size() * sizeof(float); + total_bytes += count * sizeof(float); } } diff --git a/cpp/neuralnet/desc.cpp b/cpp/neuralnet/desc.cpp index 252663397..2fa166427 100644 --- a/cpp/neuralnet/desc.cpp +++ b/cpp/neuralnet/desc.cpp @@ -197,6 +197,10 @@ void ConvLayerDesc::scaleOutputChannels(const std::vector& scaling) { } } +void ConvLayerDesc::releaseWeights() { + std::vector().swap(weights); +} + //----------------------------------------------------------------------------- BatchNormLayerDesc::BatchNormLayerDesc() : numChannels(0), epsilon(0.001f), hasScale(false), hasBias(false) {} @@ -389,6 +393,15 @@ ActivationLayerDesc::ActivationLayerDesc(istream& in, int modelVersion) { } } +void BatchNormLayerDesc::releaseWeights() { + std::vector().swap(mean); + std::vector().swap(variance); + std::vector().swap(scale); + std::vector().swap(bias); + std::vector().swap(mergedScale); + std::vector().swap(mergedBias); +} + ActivationLayerDesc::ActivationLayerDesc(ActivationLayerDesc&& other) { *this = std::move(other); } @@ -517,6 +530,10 @@ MatBiasLayerDesc::MatBiasLayerDesc(istream& in, bool binaryFloats) { throw StringError(name + ": matbiaslayer failed to parse expected number of matbias weights"); } +void MatMulLayerDesc::releaseWeights() { + std::vector().swap(weights); +} + MatBiasLayerDesc::MatBiasLayerDesc(MatBiasLayerDesc&& other) { *this = std::move(other); } @@ -538,6 +555,10 @@ void MatBiasLayerDesc::applyScale8ToReduceActivations() { } } +void MatBiasLayerDesc::releaseWeights() { + std::vector().swap(weights); +} + //----------------------------------------------------------------------------- ResidualBlockDesc::ResidualBlockDesc() {} @@ -617,6 +638,13 @@ void ResidualBlockDesc::applyScale8ToReduceActivations() { midActivation.applyScale8ToReduceActivations(); } +void ResidualBlockDesc::releaseWeights() { + preBN.releaseWeights(); + regularConv.releaseWeights(); + midBN.releaseWeights(); + finalConv.releaseWeights(); +} + //----------------------------------------------------------------------------- GlobalPoolingResidualBlockDesc::GlobalPoolingResidualBlockDesc() {} @@ -738,6 +766,16 @@ void GlobalPoolingResidualBlockDesc::applyScale8ToReduceActivations() { midActivation.applyScale8ToReduceActivations(); } +void GlobalPoolingResidualBlockDesc::releaseWeights() { + preBN.releaseWeights(); + regularConv.releaseWeights(); + gpoolConv.releaseWeights(); + gpoolBN.releaseWeights(); + gpoolToBiasMul.releaseWeights(); + midBN.releaseWeights(); + finalConv.releaseWeights(); +} + //----------------------------------------------------------------------------- NestedBottleneckResidualBlockDesc::NestedBottleneckResidualBlockDesc() {} @@ -992,6 +1030,38 @@ void NestedBottleneckResidualBlockDesc::applyScale8ToReduceActivations() { postActivation.applyScale8ToReduceActivations(); } +void NestedBottleneckResidualBlockDesc::releaseWeights() { + preBN.releaseWeights(); + preConv.releaseWeights(); + for(int i = 0; i < blocks.size(); i++) { + if(blocks[i].first == ORDINARY_BLOCK_KIND) { + ResidualBlockDesc* desc = (ResidualBlockDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else if(blocks[i].first == GLOBAL_POOLING_BLOCK_KIND) { + GlobalPoolingResidualBlockDesc* desc = (GlobalPoolingResidualBlockDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else if(blocks[i].first == NESTED_BOTTLENECK_BLOCK_KIND) { + NestedBottleneckResidualBlockDesc* desc = (NestedBottleneckResidualBlockDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else if(blocks[i].first == TRANSFORMER_ATTENTION_BLOCK_KIND) { + TransformerAttentionDesc* desc = (TransformerAttentionDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else if(blocks[i].first == TRANSFORMER_FFN_BLOCK_KIND) { + TransformerFFNDesc* desc = (TransformerFFNDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else { + ASSERT_UNREACHABLE; + } + } + postBN.releaseWeights(); + postConv.releaseWeights(); +} + //----------------------------------------------------------------------------- RMSNormLayerDesc::RMSNormLayerDesc() : numChannels(0), epsilon(0), spatial(false), cgroupSize(0) {} @@ -1043,6 +1113,11 @@ int64_t RMSNormLayerDesc::getNumParameters() const { return (int64_t)gamma.size() + (int64_t)beta.size(); } +void RMSNormLayerDesc::releaseWeights() { + std::vector().swap(gamma); + std::vector().swap(beta); +} + //----------------------------------------------------------------------------- TransformerRMSNormDesc::TransformerRMSNormDesc() : numChannels(0), epsilon(0) {} @@ -1083,6 +1158,10 @@ int64_t TransformerRMSNormDesc::getNumParameters() const { return (int64_t)weight.size(); } +void TransformerRMSNormDesc::releaseWeights() { + std::vector().swap(weight); +} + //----------------------------------------------------------------------------- TransformerAttentionDesc::TransformerAttentionDesc() @@ -1209,6 +1288,15 @@ int64_t TransformerAttentionDesc::getNumParameters() const { (int64_t)ropeFreqs.size(); // learnable RoPE frequencies, empty for fixed/no RoPE } +void TransformerAttentionDesc::releaseWeights() { + preLN.releaseWeights(); + qProj.releaseWeights(); + kProj.releaseWeights(); + vProj.releaseWeights(); + outProj.releaseWeights(); + std::vector().swap(ropeFreqs); +} + void TransformerAttentionDesc::computeRopeCosSin(int nnXLen, int nnYLen, int paddedNNXYLen, std::vector& cosTable, std::vector& sinTable) const { if(!useRope) throw StringError("TransformerAttentionDesc::computeRopeCosSin called but useRope is false"); @@ -1344,6 +1432,13 @@ int64_t TransformerFFNDesc::getNumParameters() const { linear2.getNumParameters(); } +void TransformerFFNDesc::releaseWeights() { + preLN.releaseWeights(); + linear1.releaseWeights(); + linearGate.releaseWeights(); + linear2.releaseWeights(); +} + //----------------------------------------------------------------------------- static void parseResidualBlockStack( @@ -1550,6 +1645,14 @@ int64_t SGFMetadataEncoderDesc::getNumParameters() const { mul3.getNumParameters(); } +void SGFMetadataEncoderDesc::releaseWeights() { + mul1.releaseWeights(); + bias1.releaseWeights(); + mul2.releaseWeights(); + bias2.releaseWeights(); + mul3.releaseWeights(); +} + //----------------------------------------------------------------------------- TrunkDesc::TrunkDesc() @@ -1906,6 +2009,40 @@ void TrunkDesc::applyScale8ToReduceActivations() { } } +void TrunkDesc::releaseWeights() { + initialConv.releaseWeights(); + initialMatMul.releaseWeights(); + if(metaEncoderVersion > 0) + sgfMetadataEncoder.releaseWeights(); + for(int i = 0; i < blocks.size(); i++) { + if(blocks[i].first == ORDINARY_BLOCK_KIND) { + ResidualBlockDesc* desc = (ResidualBlockDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else if(blocks[i].first == GLOBAL_POOLING_BLOCK_KIND) { + GlobalPoolingResidualBlockDesc* desc = (GlobalPoolingResidualBlockDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else if(blocks[i].first == NESTED_BOTTLENECK_BLOCK_KIND) { + NestedBottleneckResidualBlockDesc* desc = (NestedBottleneckResidualBlockDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else if(blocks[i].first == TRANSFORMER_ATTENTION_BLOCK_KIND) { + TransformerAttentionDesc* desc = (TransformerAttentionDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else if(blocks[i].first == TRANSFORMER_FFN_BLOCK_KIND) { + TransformerFFNDesc* desc = (TransformerFFNDesc*)blocks[i].second.get(); + desc->releaseWeights(); + } + else { + ASSERT_UNREACHABLE; + } + } + // Whichever trunk tip norm is unused has empty parameter vectors, so releasing both is safe. + trunkTipBN.releaseWeights(); + trunkTipRMSNorm.releaseWeights(); +} //----------------------------------------------------------------------------- @@ -2086,6 +2223,18 @@ void PolicyHeadDesc::applyScale8ToReduceActivations() { passActivation.applyScale8ToReduceActivations(); } +void PolicyHeadDesc::releaseWeights() { + p1Conv.releaseWeights(); + g1Conv.releaseWeights(); + g1BN.releaseWeights(); + gpoolToBiasMul.releaseWeights(); + p1BN.releaseWeights(); + p2Conv.releaseWeights(); + gpoolToPassMul.releaseWeights(); + gpoolToPassBias.releaseWeights(); + gpoolToPassMul2.releaseWeights(); +} + //----------------------------------------------------------------------------- ValueHeadDesc::ValueHeadDesc() : modelVersion(-1) {} @@ -2246,6 +2395,17 @@ void ValueHeadDesc::applyScale8ToReduceActivations() { sv3Bias.applyScale8ToReduceActivations(); } +void ValueHeadDesc::releaseWeights() { + v1Conv.releaseWeights(); + v1BN.releaseWeights(); + v2Mul.releaseWeights(); + v2Bias.releaseWeights(); + v3Mul.releaseWeights(); + v3Bias.releaseWeights(); + sv3Mul.releaseWeights(); + sv3Bias.releaseWeights(); + vOwnershipConv.releaseWeights(); +} //----------------------------------------------------------------------------- @@ -2562,6 +2722,12 @@ void ModelDesc::applyScale8ToReduceActivations() { postProcessParams.outputScaleMultiplier *= 8.0f; } +void ModelDesc::releaseWeights() { + trunk.releaseWeights(); + policyHead.releaseWeights(); + valueHead.releaseWeights(); +} + struct NonCopyingStreamBuf : public std::streambuf { NonCopyingStreamBuf(string& str) { diff --git a/cpp/neuralnet/desc.h b/cpp/neuralnet/desc.h index 36c5a11d8..ef41dfca6 100644 --- a/cpp/neuralnet/desc.h +++ b/cpp/neuralnet/desc.h @@ -36,6 +36,8 @@ struct ConvLayerDesc { int64_t getNumParameters() const; void scaleOutputChannels(const std::vector& scaling); + + void releaseWeights(); }; struct BatchNormLayerDesc { @@ -68,6 +70,8 @@ struct BatchNormLayerDesc { void extractChannelFactorsAbsLtOne(std::vector& channelFactors); void extractChannelFactorsAbsLtOneWithInverses(std::vector& channelFactors, std::vector& invChannelFactors); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct ActivationLayerDesc { @@ -105,6 +109,8 @@ struct MatMulLayerDesc { int64_t getNumParameters() const; void scaleOutputChannels(const std::vector& scaling); + + void releaseWeights(); }; struct MatBiasLayerDesc { @@ -124,6 +130,8 @@ struct MatBiasLayerDesc { int64_t getNumParameters() const; void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct ResidualBlockDesc { @@ -150,6 +158,8 @@ struct ResidualBlockDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct GlobalPoolingResidualBlockDesc { @@ -181,6 +191,8 @@ struct GlobalPoolingResidualBlockDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct NestedBottleneckResidualBlockDesc { @@ -215,6 +227,8 @@ struct NestedBottleneckResidualBlockDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; // Trunk final normalization kind (stored in trunk header) @@ -240,6 +254,7 @@ struct RMSNormLayerDesc { RMSNormLayerDesc& operator=(RMSNormLayerDesc&& other); int64_t getNumParameters() const; + void releaseWeights(); }; // Lightweight RMSNorm used inside transformer blocks (weight only, no bias, no spatial modes) @@ -259,6 +274,7 @@ struct TransformerRMSNormDesc { TransformerRMSNormDesc& operator=(TransformerRMSNormDesc&& other); int64_t getNumParameters() const; + void releaseWeights(); }; struct TransformerAttentionDesc { @@ -294,6 +310,7 @@ struct TransformerAttentionDesc { TransformerAttentionDesc& operator=(TransformerAttentionDesc&& other); int64_t getNumParameters() const; + void releaseWeights(); // Compute cos/sin tables for RoPE given board dimensions. // Output tables are indexed as: @@ -324,6 +341,7 @@ struct TransformerFFNDesc { TransformerFFNDesc& operator=(TransformerFFNDesc&& other); int64_t getNumParameters() const; + void releaseWeights(); }; struct SGFMetadataEncoderDesc { @@ -349,6 +367,7 @@ struct SGFMetadataEncoderDesc { SGFMetadataEncoderDesc& operator=(SGFMetadataEncoderDesc&& other); int64_t getNumParameters() const; + void releaseWeights(); }; @@ -397,6 +416,8 @@ struct TrunkDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct PolicyHeadDesc { @@ -431,6 +452,8 @@ struct PolicyHeadDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct ValueHeadDesc { @@ -463,6 +486,8 @@ struct ValueHeadDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct ModelPostProcessParams { @@ -534,6 +559,11 @@ struct ModelDesc { //Fills supported with true if desiredRules itself was exactly supported, false if some modifications had to be made. Rules getSupportedRules(const Rules& desiredRules, bool& supported) const; + // Frees all weight arrays (conv/matmul/bias/batchnorm), keeping scalar shape + // metadata intact. Safe once weights are no longer needed (e.g. CoreML/ANE + // inference, which reads weights from the compiled .mlmodelc). + void releaseWeights(); + }; #endif // #ifndef DESC_H diff --git a/cpp/neuralnet/metalbackend.cpp b/cpp/neuralnet/metalbackend.cpp index e46d6720e..09a43f294 100644 --- a/cpp/neuralnet/metalbackend.cpp +++ b/cpp/neuralnet/metalbackend.cpp @@ -508,13 +508,27 @@ ComputeContext* NeuralNet::createComputeContext( const LoadedModel* loadedModel, ConfigParser& cfg) { - (void)gpuIdxs; + // Only ANE-only configurations may free the engine's in-memory weights: the + // GPU/MPSGraph path reads them via modelDescToSwift, so freeing is unsafe + // unless no GPU handle can ever be built from this model. + // INVARIANT: gpuIdxs must be the complete (deduplicated) set of device indices + // that will ever be passed as gpuIdxForThisThread to createComputeHandle for + // this context. aneOnly==true frees the in-memory weights, so if any thread + // later used a GPU (MPSGraph) index not represented here, it would read freed + // weights. KataGo derives both from the same gpuIdxByServerThread list, so the + // invariant holds today; preserve it if that wiring ever changes. + bool aneOnly = !gpuIdxs.empty(); + for(int idx : gpuIdxs) { + if(idx != METAL_MUX_ANE) { aneOnly = false; break; } + } (void)logger; (void)homeDataDirOverride; (void)loadedModel; (void)cfg; - return new ComputeContext(nnXLen, nnYLen, useFP16Mode); + ComputeContext* context = new ComputeContext(nnXLen, nnYLen, useFP16Mode); + context->aneOnly = aneOnly; + return context; } void NeuralNet::freeComputeContext(ComputeContext* computeContext) { @@ -541,6 +555,17 @@ static swift::Optional convertAndCreateCoreMLO bool useFP16 = (context->useFP16Mode != enabled_t::False); bool optimizeMask = requireExactNNLen; + // On a confirmed ANE-only run, free the engine's in-memory ModelDesc weight + // arrays. This function converts from loadedModel->modelPath (disk), + // so the in-memory weights are not read here; the GPU/MPSGraph path (which + // DOES read them via modelDescToSwift) is never built when aneOnly is true. + // The whole ComputeHandle ctor runs under computeHandleMutex, so this is not + // racy; releaseWeights() clears only weight vectors, leaving the scalar dims + // read by the ComputeHandle ctor / InputBuffers valid. + if(context->aneOnly) { + const_cast(loadedModel)->modelDesc.releaseWeights(); + } + // Convert model to CoreML format in temp directory string coremlModelPath = CoreMLConversion::convertModelToTemp( loadedModel->modelPath, diff --git a/cpp/neuralnet/metalbackend.h b/cpp/neuralnet/metalbackend.h index b7f751e63..e77dd18d9 100644 --- a/cpp/neuralnet/metalbackend.h +++ b/cpp/neuralnet/metalbackend.h @@ -113,6 +113,14 @@ struct ComputeContext { */ MetalComputeContext metalContext; + /** + * @brief True only when every configured device is METAL_MUX_ANE, so no + * MPSGraph (GPU) handle will ever read modelDesc weights. Gates the call to + * ModelDesc::releaseWeights() so a mixed GPU+ANE config can never free live + * weights. + */ + bool aneOnly = false; + /** * @brief Constructs a ComputeContext object. * @param nnX The width of the input tensor. @@ -179,6 +187,18 @@ struct ComputeHandle { */ bool maskIdentityChecked = false; + // Weight-release safety is guaranteed by ComputeContext::aneOnly, NOT by the + // declaration order below: within a single ComputeHandle exactly one handle is + // built (the two paths are mutually exclusive on gpuIdx, enforced by the + // ctor's exactly-one check), and releaseWeights() only ever fires on an + // aneOnly context, where no MPSGraph handle is built for any thread. + // That said, keep mpsGraphOnlyHandle declared before coremlOnlyHandle. C++ + // initializes members in DECLARATION order, so createMPSGraphHandleIfNeeded + // (which reads modelDesc weights via modelDescToSwift) is sequenced before + // createCoreMLOnlyHandleIfNeeded (which may call modelDesc.releaseWeights()). + // This ordering is belt-and-suspenders that preserves the natural read-then- + // release sequence should the aneOnly invariant ever be weakened; don't rely + // on it as the primary guarantee, but don't reorder it either. /** * @brief The MPSGraph-only handle instance from Swift (GPU-only mode). */