From b05f5594ee1db8412efe389de15d416f7d1e442e Mon Sep 17 00:00:00 2001 From: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Date: Sat, 30 May 2026 17:11:04 +0800 Subject: [PATCH 1/8] Reduce CoreML conversion peak and ANE steady-state memory Cuts memory during the on-device KataGo -> CoreML conversion and while running the ANE/CoreML path, with byte-identical converter output: - The converter's weight tensors become non-owning views into the parsed model instead of owning extra FP32 copies; derived/transposed tensors keep an owned buffer. This drops redundant resident weight copies during conversion. CoreML model serialization is made deterministic (SetSerializationDeterministic) so the output is byte-stable. - The KataGo model parser streams the gzip through a bounded ~1 MB refill buffer instead of decompressing the whole file into memory, while preserving the existing NaN/Inf weight validation. - ModelDesc gains releaseWeights(), which frees the in-memory weight arrays (keeping scalar shape metadata). The Metal backend calls it on the ANE (CoreML) path after converting from the model file on disk, gated by a new ComputeContext::aneOnly flag so it only fires when every configured device is ANE -- the GPU/MPSGraph path keeps its weights. The call is serialized under computeHandleMutex and only scalar dims are read afterward. Measured on b18c384nbt (19x19) over the ANE path: idle steady-state RSS 0.59 GB -> 0.19 GB; peak (load+convert) 0.87 GB -> 0.48 GB. Cross-backend parity vs an Eigen reference is unchanged on both the GPU and ANE paths. Co-Authored-By: Claude Opus 4.8 (1M context) --- cpp/external/katagocoreml/src/Converter.cpp | 16 +- .../katagocoreml/src/builder/MILBuilder.cpp | 18 +- .../katagocoreml/src/builder/MILBuilder.hpp | 13 +- .../katagocoreml/src/builder/Operations.cpp | 20 ++- .../katagocoreml/src/builder/Operations.hpp | 24 ++- .../katagocoreml/src/parser/KataGoParser.cpp | 168 ++++++++---------- .../katagocoreml/src/parser/KataGoParser.hpp | 16 +- .../src/serializer/CoreMLSerializer.cpp | 11 +- .../src/serializer/WeightSerializer.cpp | 10 +- cpp/neuralnet/desc.cpp | 69 +++++++ cpp/neuralnet/desc.h | 5 + cpp/neuralnet/metalbackend.cpp | 29 ++- cpp/neuralnet/metalbackend.h | 14 ++ 13 files changed, 287 insertions(+), 126 deletions(-) diff --git a/cpp/external/katagocoreml/src/Converter.cpp b/cpp/external/katagocoreml/src/Converter.cpp index cb6ca80d9..72b78e736 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; @@ -82,7 +84,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); } ModelInfo KataGoConverter::getModelInfo(const std::string& input_path) { diff --git a/cpp/external/katagocoreml/src/builder/MILBuilder.cpp b/cpp/external/katagocoreml/src/builder/MILBuilder.cpp index db0c6c4b1..a30d2ce43 100644 --- a/cpp/external/katagocoreml/src/builder/MILBuilder.cpp +++ b/cpp/external/katagocoreml/src/builder/MILBuilder.cpp @@ -212,9 +212,23 @@ 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 + // Register weight for blob storage (non-owning view into the model) m_ops.registerWeight(name, data, shape); + 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 weight; KataGoOps takes ownership of the buffer + m_ops.registerOwnedWeight(name, std::move(data), shape); + 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"); @@ -958,7 +972,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)}; diff --git a/cpp/external/katagocoreml/src/builder/MILBuilder.hpp b/cpp/external/katagocoreml/src/builder/MILBuilder.hpp index 042f9fc16..640864579 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; } @@ -80,6 +80,15 @@ class MILBuilder { const std::vector& data, const std::vector& shape); + 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 c0c036292..148c44089 100644 --- a/cpp/external/katagocoreml/src/builder/Operations.cpp +++ b/cpp/external/katagocoreml/src/builder/Operations.cpp @@ -17,9 +17,25 @@ std::string KataGoOps::registerWeight(const std::string& name, const std::vector& shape) { WeightEntry entry; entry.name = name; - entry.data = data; + entry.data = data.data(); + entry.count = data.size(); entry.shape = shape; - entry.blob_offset = 0; // Will be set during serialization + entry.blob_offset = 0; + m_weights.push_back(std::move(entry)); + return name; +} + +std::string KataGoOps::registerOwnedWeight(const std::string& name, + std::vector&& data, + const std::vector& shape) { + m_owned.push_back(std::move(data)); + const std::vector& stored = m_owned.back(); + WeightEntry entry; + entry.name = name; + entry.data = stored.data(); + entry.count = stored.size(); + entry.shape = shape; + entry.blob_offset = 0; m_weights.push_back(std::move(entry)); return name; } diff --git a/cpp/external/katagocoreml/src/builder/Operations.hpp b/cpp/external/katagocoreml/src/builder/Operations.hpp index 3fc72ad88..9649cb8e6 100644 --- a/cpp/external/katagocoreml/src/builder/Operations.hpp +++ b/cpp/external/katagocoreml/src/builder/Operations.hpp @@ -5,15 +5,18 @@ #include "../types/KataGoTypes.hpp" #include +#include #include #include namespace katagocoreml { -/// Weight entry for blob file storage +/// Weight entry for blob file storage. `data`/`count` are a NON-OWNING view into +/// the live KataGoModelDesc (or into KataGoOps::m_owned for derived tensors). struct WeightEntry { std::string name; - std::vector data; + const float* data = nullptr; + size_t count = 0; std::vector shape; uint64_t blob_offset = 0; // Set during serialization }; @@ -51,16 +54,22 @@ class KataGoOps { /// Get precomputed mask constants const MaskConstants& getMaskConstants() const { return m_mask_constants; } - /// Register a weight tensor and return its reference name + /// Register a weight that lives in the model (stored as a non-owning view). std::string registerWeight(const std::string& name, const std::vector& data, const std::vector& shape); - /// Get all registered weights - const std::vector& getWeights() const { return m_weights; } + /// Register a derived/temporary weight; KataGoOps takes ownership so the + /// view stays valid through serialization. + std::string registerOwnedWeight(const std::string& name, + std::vector&& data, + const std::vector& shape); - /// Clear all registered weights - void clearWeights() { m_weights.clear(); } + /// Get all registered weights (mutable; serialization sets blob_offset) + std::vector& getWeightsMutable() { return m_weights; } + + /// 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); @@ -71,6 +80,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 2d06c27e5..19b26e90d 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 @@ -30,54 +29,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 == nullptr) return false; + int n = gzread(m_gz, m_refill.data(), (unsigned)m_refill.size()); + if(n < 0) { + int errnum; + const char* errmsg = gzerror(m_gz, &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; } } @@ -86,16 +72,27 @@ 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()); - - return parseModel(); + // Allocate the refill buffer before opening the file so a bad_alloc here + // cannot leak an open gzFile handle. + m_refill.resize(1024 * 1024); + m_gz = gzopen(m_model_path.c_str(), "rb"); + if(m_gz == nullptr) + 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; + KataGoModelDesc model; + try { + model = parseModel(); + } catch(...) { + gzclose(m_gz); + m_gz = nullptr; + throw; + } + gzclose(m_gz); + m_gz = nullptr; + return model; } // ============================================================================ @@ -103,24 +100,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++; } } @@ -147,37 +140,28 @@ bool KataGoParser::readBool() { std::vector KataGoParser::readFloats(size_t count, const std::string& name) { 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 - while (m_pos < m_buffer.size()) { - if (m_buffer[m_pos] == '@') { - break; - } - 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) - size_t num_bytes = count * 4; - if (m_pos + num_bytes > m_buffer.size()) { - throw std::runtime_error(name + ": not enough bytes for " + std::to_string(count) + " floats"); - } - // 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 cbcfdefa8..2d3f1c47a 100644 --- a/cpp/external/katagocoreml/src/parser/KataGoParser.hpp +++ b/cpp/external/katagocoreml/src/parser/KataGoParser.hpp @@ -7,6 +7,7 @@ #include #include #include +#include namespace katagocoreml { @@ -31,9 +32,17 @@ class KataGoParser { private: std::string m_model_path; - std::vector m_buffer; - size_t m_pos = 0; + gzFile m_gz = nullptr; + 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); @@ -65,9 +74,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 2ac23a3da..86e41aaec 100644 --- a/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp +++ b/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp @@ -17,18 +17,18 @@ size_t WeightSerializer::serialize(std::vector& weights, for (auto& entry : weights) { if (use_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(entry.count); + for (size_t i = 0; i < entry.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 += entry.count * sizeof(MILBlob::Fp16); } else { // Write FP32 weights - MILBlob::Util::Span span(entry.data.data(), entry.data.size()); + MILBlob::Util::Span span(entry.data, entry.count); entry.blob_offset = writer.WriteData(span); - total_bytes += entry.data.size() * sizeof(float); + total_bytes += entry.count * sizeof(float); } } diff --git a/cpp/neuralnet/desc.cpp b/cpp/neuralnet/desc.cpp index eda55111a..e59d2e585 100644 --- a/cpp/neuralnet/desc.cpp +++ b/cpp/neuralnet/desc.cpp @@ -1783,6 +1783,75 @@ void ModelDesc::applyScale8ToReduceActivations() { postProcessParams.outputScaleMultiplier *= 8.0f; } +static void releaseVec(std::vector& v) { std::vector().swap(v); } + +static void releaseConv(ConvLayerDesc& c) { releaseVec(c.weights); } + +static void releaseBN(BatchNormLayerDesc& b) { + releaseVec(b.mean); releaseVec(b.variance); releaseVec(b.scale); + releaseVec(b.bias); releaseVec(b.mergedScale); releaseVec(b.mergedBias); +} + +static void releaseMatMul(MatMulLayerDesc& m) { releaseVec(m.weights); } +static void releaseMatBias(MatBiasLayerDesc& m) { releaseVec(m.weights); } + +static void releaseResidual(ResidualBlockDesc& b) { + releaseBN(b.preBN); releaseConv(b.regularConv); + releaseBN(b.midBN); releaseConv(b.finalConv); +} + +static void releaseGPool(GlobalPoolingResidualBlockDesc& b) { + releaseBN(b.preBN); releaseConv(b.regularConv); releaseConv(b.gpoolConv); + releaseBN(b.gpoolBN); releaseMatMul(b.gpoolToBiasMul); + releaseBN(b.midBN); releaseConv(b.finalConv); +} + +static void releaseBlocks(std::vector>& blocks); + +static void releaseNested(NestedBottleneckResidualBlockDesc& b) { + releaseBN(b.preBN); releaseConv(b.preConv); + releaseBlocks(b.blocks); + releaseBN(b.postBN); releaseConv(b.postConv); +} + +static void releaseBlocks(std::vector>& blocks) { + for(size_t i = 0; i < blocks.size(); i++) { + if(blocks[i].first == ORDINARY_BLOCK_KIND) + releaseResidual(*(ResidualBlockDesc*)blocks[i].second.get()); + else if(blocks[i].first == GLOBAL_POOLING_BLOCK_KIND) + releaseGPool(*(GlobalPoolingResidualBlockDesc*)blocks[i].second.get()); + else if(blocks[i].first == NESTED_BOTTLENECK_BLOCK_KIND) + releaseNested(*(NestedBottleneckResidualBlockDesc*)blocks[i].second.get()); + else + ASSERT_UNREACHABLE; + } +} + +static void releaseSGFEncoder(SGFMetadataEncoderDesc& e) { + releaseMatMul(e.mul1); releaseMatBias(e.bias1); + releaseMatMul(e.mul2); releaseMatBias(e.bias2); + releaseMatMul(e.mul3); +} + +void ModelDesc::releaseWeights() { + releaseConv(trunk.initialConv); + releaseMatMul(trunk.initialMatMul); + if(trunk.metaEncoderVersion > 0) + releaseSGFEncoder(trunk.sgfMetadataEncoder); + releaseBlocks(trunk.blocks); + releaseBN(trunk.trunkTipBN); + releaseConv(policyHead.p1Conv); releaseConv(policyHead.g1Conv); + releaseBN(policyHead.g1BN); releaseMatMul(policyHead.gpoolToBiasMul); + releaseBN(policyHead.p1BN); releaseConv(policyHead.p2Conv); + releaseMatMul(policyHead.gpoolToPassMul); releaseMatBias(policyHead.gpoolToPassBias); + releaseMatMul(policyHead.gpoolToPassMul2); + releaseConv(valueHead.v1Conv); releaseBN(valueHead.v1BN); + releaseMatMul(valueHead.v2Mul); releaseMatBias(valueHead.v2Bias); + releaseMatMul(valueHead.v3Mul); releaseMatBias(valueHead.v3Bias); + releaseMatMul(valueHead.sv3Mul); releaseMatBias(valueHead.sv3Bias); + releaseConv(valueHead.vOwnershipConv); +} + struct NonCopyingStreamBuf : public std::streambuf { NonCopyingStreamBuf(string& str) { diff --git a/cpp/neuralnet/desc.h b/cpp/neuralnet/desc.h index 86676c011..9536283f2 100644 --- a/cpp/neuralnet/desc.h +++ b/cpp/neuralnet/desc.h @@ -389,6 +389,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 95adf5da4..1e7db2fba 100644 --- a/cpp/neuralnet/metalbackend.cpp +++ b/cpp/neuralnet/metalbackend.cpp @@ -425,14 +425,28 @@ ComputeContext* NeuralNet::createComputeContext( enabled_t useNHWCMode, const LoadedModel* loadedModel) { - (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)openCLTunerFile; (void)homeDataDirOverride; (void)openCLReTunePerBoardSize; (void)loadedModel; - return new ComputeContext(nnXLen, nnYLen, useFP16Mode, useNHWCMode); + ComputeContext* context = new ComputeContext(nnXLen, nnYLen, useFP16Mode, useNHWCMode); + context->aneOnly = aneOnly; + return context; } void NeuralNet::freeComputeContext(ComputeContext* computeContext) { @@ -459,6 +473,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 a00f21864..db161d28c 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. @@ -180,6 +188,12 @@ struct ComputeHandle { */ bool maskIdentityChecked = false; + // IMPORTANT (weight-release safety): mpsGraphOnlyHandle MUST be declared + // before coremlOnlyHandle. C++ initializes members in DECLARATION order, so + // createMPSGraphHandleIfNeeded (which reads modelDesc weights via + // modelDescToSwift) runs before createCoreMLOnlyHandleIfNeeded (which may call + // modelDesc.releaseWeights() on an ANE-only run). Reordering these would let a + // GPU handle read freed weights. Do not reorder. /** * @brief The MPSGraph-only handle instance from Swift (GPU-only mode). */ From 971fa9d8c0bd9fafd7987f25edfcb5cc96c38c1d Mon Sep 17 00:00:00 2001 From: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Date: Sat, 30 May 2026 18:37:44 +0800 Subject: [PATCH 2/8] Enforce non-owning weight-view contract at compile time WeightEntry stores a non-owning view (const float*, count) into the live KataGoModelDesc, so the backing std::vector must outlive serialization. addConstOp/registerWeight took the data by const& and silently stored a pointer to it; a caller passing a temporary would bind to that const& and leave the view dangling, read much later during serialization. Delete the rvalue overloads of both so any such call fails to compile, forcing temporaries through addOwnedConstOp/registerOwnedWeight (which take ownership). Named lvalues (the model-member call sites) still bind to the const& overload, so no existing caller changes. Co-Authored-By: Claude Opus 4.8 (1M context) --- cpp/external/katagocoreml/src/builder/MILBuilder.hpp | 9 +++++++++ cpp/external/katagocoreml/src/builder/Operations.hpp | 7 +++++++ 2 files changed, 16 insertions(+) diff --git a/cpp/external/katagocoreml/src/builder/MILBuilder.hpp b/cpp/external/katagocoreml/src/builder/MILBuilder.hpp index 640864579..a25c3f537 100644 --- a/cpp/external/katagocoreml/src/builder/MILBuilder.hpp +++ b/cpp/external/katagocoreml/src/builder/MILBuilder.hpp @@ -80,6 +80,15 @@ 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, diff --git a/cpp/external/katagocoreml/src/builder/Operations.hpp b/cpp/external/katagocoreml/src/builder/Operations.hpp index 9649cb8e6..f5431f79f 100644 --- a/cpp/external/katagocoreml/src/builder/Operations.hpp +++ b/cpp/external/katagocoreml/src/builder/Operations.hpp @@ -59,6 +59,13 @@ class KataGoOps { const std::vector& data, const std::vector& shape); + /// 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. std::string registerOwnedWeight(const std::string& name, From eeefc976222fcaf67fc4092300b6e39db38c634d Mon Sep 17 00:00:00 2001 From: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Date: Sun, 31 May 2026 10:31:12 +0800 Subject: [PATCH 3/8] RAII the gzFile handle in KataGoParser Own the gzFile with a custom-deleter unique_ptr so it closes on every exit path (normal return, exception, bad_alloc); removes the manual try/catch+gzclose in parse() and the ordering caveat on buffer allocation. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../katagocoreml/src/parser/KataGoParser.cpp | 26 ++++++------------- .../katagocoreml/src/parser/KataGoParser.hpp | 10 ++++++- 2 files changed, 17 insertions(+), 19 deletions(-) diff --git a/cpp/external/katagocoreml/src/parser/KataGoParser.cpp b/cpp/external/katagocoreml/src/parser/KataGoParser.cpp index 19b26e90d..37497f6cc 100644 --- a/cpp/external/katagocoreml/src/parser/KataGoParser.cpp +++ b/cpp/external/katagocoreml/src/parser/KataGoParser.cpp @@ -33,11 +33,11 @@ bool KataGoParser::isVersionSupported(int version) { // ============================================================================ bool KataGoParser::refill() { - if(m_gz == nullptr) return false; - int n = gzread(m_gz, m_refill.data(), (unsigned)m_refill.size()); + 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, &errnum); + const char* errmsg = gzerror(m_gz.get(), &errnum); throw std::runtime_error("Error reading gzip stream: " + std::string(errmsg)); } m_refillPos = 0; @@ -72,27 +72,17 @@ void KataGoParser::readExact(uint8_t* dst, size_t n, const std::string& name) { // ============================================================================ KataGoModelDesc KataGoParser::parse() { - // Allocate the refill buffer before opening the file so a bad_alloc here - // cannot leak an open gzFile handle. + // Allocate the refill buffer first; if this throws, no handle has been opened. m_refill.resize(1024 * 1024); - m_gz = gzopen(m_model_path.c_str(), "rb"); - if(m_gz == nullptr) + 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; - KataGoModelDesc model; - try { - model = parseModel(); - } catch(...) { - gzclose(m_gz); - m_gz = nullptr; - throw; - } - gzclose(m_gz); - m_gz = nullptr; - return model; + // ~GzHandle closes the file on normal return OR exception — no try/catch needed. + return parseModel(); } // ============================================================================ diff --git a/cpp/external/katagocoreml/src/parser/KataGoParser.hpp b/cpp/external/katagocoreml/src/parser/KataGoParser.hpp index 2d3f1c47a..8ee1a90ab 100644 --- a/cpp/external/katagocoreml/src/parser/KataGoParser.hpp +++ b/cpp/external/katagocoreml/src/parser/KataGoParser.hpp @@ -5,7 +5,9 @@ #include "../types/KataGoTypes.hpp" #include +#include #include +#include #include #include @@ -32,7 +34,13 @@ class KataGoParser { private: std::string m_model_path; - gzFile m_gz = nullptr; + // 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 From 6bfa617b9fafe0be9b40b04196be7f94ed22a8f6 Mon Sep 17 00:00:00 2001 From: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Date: Sun, 31 May 2026 10:31:14 +0800 Subject: [PATCH 4/8] Replace WeightEntry raw ptr+count with a local FloatView Introduce a KataGo-local non-owning FloatView for WeightEntry::data instead of a raw const float*/size_t pair; convert to MILBlob::Util::Span only inside WeightSerializer, keeping the MILBlob dependency out of Operations.hpp. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../katagocoreml/src/builder/Operations.cpp | 6 ++---- .../katagocoreml/src/builder/Operations.hpp | 19 +++++++++++++++---- .../src/serializer/WeightSerializer.cpp | 13 +++++++------ 3 files changed, 24 insertions(+), 14 deletions(-) diff --git a/cpp/external/katagocoreml/src/builder/Operations.cpp b/cpp/external/katagocoreml/src/builder/Operations.cpp index 148c44089..4cbd1038a 100644 --- a/cpp/external/katagocoreml/src/builder/Operations.cpp +++ b/cpp/external/katagocoreml/src/builder/Operations.cpp @@ -17,8 +17,7 @@ std::string KataGoOps::registerWeight(const std::string& name, const std::vector& shape) { WeightEntry entry; entry.name = name; - entry.data = data.data(); - entry.count = data.size(); + entry.data = FloatView{data.data(), data.size()}; entry.shape = shape; entry.blob_offset = 0; m_weights.push_back(std::move(entry)); @@ -32,8 +31,7 @@ std::string KataGoOps::registerOwnedWeight(const std::string& name, const std::vector& stored = m_owned.back(); WeightEntry entry; entry.name = name; - entry.data = stored.data(); - entry.count = stored.size(); + entry.data = FloatView{stored.data(), stored.size()}; entry.shape = shape; entry.blob_offset = 0; m_weights.push_back(std::move(entry)); diff --git a/cpp/external/katagocoreml/src/builder/Operations.hpp b/cpp/external/katagocoreml/src/builder/Operations.hpp index f5431f79f..5bc8378e2 100644 --- a/cpp/external/katagocoreml/src/builder/Operations.hpp +++ b/cpp/external/katagocoreml/src/builder/Operations.hpp @@ -11,12 +11,23 @@ namespace katagocoreml { -/// Weight entry for blob file storage. `data`/`count` are a NON-OWNING view into -/// the live KataGoModelDesc (or into KataGoOps::m_owned for derived tensors). +/// 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; - const float* data = nullptr; - size_t count = 0; + FloatView data; // non-owning view (replaces raw ptr + count) std::vector shape; uint64_t blob_offset = 0; // Set during serialization }; diff --git a/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp b/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp index 86e41aaec..e27a342f7 100644 --- a/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp +++ b/cpp/external/katagocoreml/src/serializer/WeightSerializer.cpp @@ -15,20 +15,21 @@ size_t WeightSerializer::serialize(std::vector& weights, size_t total_bytes = 0; for (auto& entry : weights) { + const size_t count = entry.data.size(); if (use_fp16) { // Convert FP32 weights to FP16 - std::vector fp16_data(entry.count); - for (size_t i = 0; i < entry.count; ++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.count * sizeof(MILBlob::Fp16); + total_bytes += count * sizeof(MILBlob::Fp16); } else { - // Write FP32 weights - MILBlob::Util::Span span(entry.data, entry.count); + // 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.count * sizeof(float); + total_bytes += count * sizeof(float); } } From 415993015e80674a03d3c34b06af6a2773047483 Mon Sep 17 00:00:00 2001 From: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Date: Sun, 31 May 2026 21:57:00 +0800 Subject: [PATCH 5/8] Clarify weight-release safety comment: aneOnly is the guarantee The ComputeHandle member-order comment claimed that declaring mpsGraphOnlyHandle before coremlOnlyHandle is what prevents a GPU handle from reading freed weights. That overstates the ordering's role: within a single ComputeHandle exactly one handle is built (mutually exclusive on gpuIdx, enforced by the ctor's exactly-one check), and releaseWeights() only fires on an aneOnly context where no MPSGraph handle is ever built. Reframe the declaration order as belt-and-suspenders and point at ComputeContext::aneOnly as the actual invariant. Comment-only change. Co-Authored-By: Claude Opus 4.8 (1M context) --- cpp/neuralnet/metalbackend.h | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/cpp/neuralnet/metalbackend.h b/cpp/neuralnet/metalbackend.h index db161d28c..3922d6828 100644 --- a/cpp/neuralnet/metalbackend.h +++ b/cpp/neuralnet/metalbackend.h @@ -188,12 +188,18 @@ struct ComputeHandle { */ bool maskIdentityChecked = false; - // IMPORTANT (weight-release safety): mpsGraphOnlyHandle MUST be declared - // before coremlOnlyHandle. C++ initializes members in DECLARATION order, so - // createMPSGraphHandleIfNeeded (which reads modelDesc weights via - // modelDescToSwift) runs before createCoreMLOnlyHandleIfNeeded (which may call - // modelDesc.releaseWeights() on an ANE-only run). Reordering these would let a - // GPU handle read freed weights. Do not reorder. + // 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). */ From 44342a388c3629ded3ed6aa6a5ca184614e6f2ab Mon Sep 17 00:00:00 2001 From: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Date: Sun, 31 May 2026 22:58:36 +0800 Subject: [PATCH 6/8] Refactor weight release into per-struct releaseWeights() methods Replace the file-local releaseXXX free functions in desc.cpp (which reached into each desc struct's internals from outside) with releaseWeights() member methods on each weight-bearing struct, matching the existing OO convention used by applyScale8ToReduceActivations() and iterConvLayers(). Each container delegates to its members; type-erased block dispatch is inlined with the same cast pattern those methods use. Behavior-preserving: same set of freed vectors, same block recursion, same metaEncoderVersion guard. ModelDesc::releaseWeights() keeps its signature, so the metalbackend.cpp call site is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- cpp/neuralnet/desc.cpp | 152 ++++++++++++++++++++++++++++------------- cpp/neuralnet/desc.h | 22 ++++++ 2 files changed, 126 insertions(+), 48 deletions(-) diff --git a/cpp/neuralnet/desc.cpp b/cpp/neuralnet/desc.cpp index e59d2e585..0c4943bb5 100644 --- a/cpp/neuralnet/desc.cpp +++ b/cpp/neuralnet/desc.cpp @@ -1783,73 +1783,129 @@ void ModelDesc::applyScale8ToReduceActivations() { postProcessParams.outputScaleMultiplier *= 8.0f; } -static void releaseVec(std::vector& v) { std::vector().swap(v); } +void ConvLayerDesc::releaseWeights() { + std::vector().swap(weights); +} -static void releaseConv(ConvLayerDesc& c) { releaseVec(c.weights); } +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); +} -static void releaseBN(BatchNormLayerDesc& b) { - releaseVec(b.mean); releaseVec(b.variance); releaseVec(b.scale); - releaseVec(b.bias); releaseVec(b.mergedScale); releaseVec(b.mergedBias); +void MatMulLayerDesc::releaseWeights() { + std::vector().swap(weights); } -static void releaseMatMul(MatMulLayerDesc& m) { releaseVec(m.weights); } -static void releaseMatBias(MatBiasLayerDesc& m) { releaseVec(m.weights); } +void MatBiasLayerDesc::releaseWeights() { + std::vector().swap(weights); +} -static void releaseResidual(ResidualBlockDesc& b) { - releaseBN(b.preBN); releaseConv(b.regularConv); - releaseBN(b.midBN); releaseConv(b.finalConv); +void ResidualBlockDesc::releaseWeights() { + preBN.releaseWeights(); + regularConv.releaseWeights(); + midBN.releaseWeights(); + finalConv.releaseWeights(); } -static void releaseGPool(GlobalPoolingResidualBlockDesc& b) { - releaseBN(b.preBN); releaseConv(b.regularConv); releaseConv(b.gpoolConv); - releaseBN(b.gpoolBN); releaseMatMul(b.gpoolToBiasMul); - releaseBN(b.midBN); releaseConv(b.finalConv); +void GlobalPoolingResidualBlockDesc::releaseWeights() { + preBN.releaseWeights(); + regularConv.releaseWeights(); + gpoolConv.releaseWeights(); + gpoolBN.releaseWeights(); + gpoolToBiasMul.releaseWeights(); + midBN.releaseWeights(); + finalConv.releaseWeights(); } -static void releaseBlocks(std::vector>& blocks); +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 { + ASSERT_UNREACHABLE; + } + } + postBN.releaseWeights(); + postConv.releaseWeights(); +} -static void releaseNested(NestedBottleneckResidualBlockDesc& b) { - releaseBN(b.preBN); releaseConv(b.preConv); - releaseBlocks(b.blocks); - releaseBN(b.postBN); releaseConv(b.postConv); +void SGFMetadataEncoderDesc::releaseWeights() { + mul1.releaseWeights(); + bias1.releaseWeights(); + mul2.releaseWeights(); + bias2.releaseWeights(); + mul3.releaseWeights(); } -static void releaseBlocks(std::vector>& blocks) { - for(size_t i = 0; i < blocks.size(); i++) { - if(blocks[i].first == ORDINARY_BLOCK_KIND) - releaseResidual(*(ResidualBlockDesc*)blocks[i].second.get()); - else if(blocks[i].first == GLOBAL_POOLING_BLOCK_KIND) - releaseGPool(*(GlobalPoolingResidualBlockDesc*)blocks[i].second.get()); - else if(blocks[i].first == NESTED_BOTTLENECK_BLOCK_KIND) - releaseNested(*(NestedBottleneckResidualBlockDesc*)blocks[i].second.get()); - else +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 { ASSERT_UNREACHABLE; + } } + trunkTipBN.releaseWeights(); +} + +void PolicyHeadDesc::releaseWeights() { + p1Conv.releaseWeights(); + g1Conv.releaseWeights(); + g1BN.releaseWeights(); + gpoolToBiasMul.releaseWeights(); + p1BN.releaseWeights(); + p2Conv.releaseWeights(); + gpoolToPassMul.releaseWeights(); + gpoolToPassBias.releaseWeights(); + gpoolToPassMul2.releaseWeights(); } -static void releaseSGFEncoder(SGFMetadataEncoderDesc& e) { - releaseMatMul(e.mul1); releaseMatBias(e.bias1); - releaseMatMul(e.mul2); releaseMatBias(e.bias2); - releaseMatMul(e.mul3); +void ValueHeadDesc::releaseWeights() { + v1Conv.releaseWeights(); + v1BN.releaseWeights(); + v2Mul.releaseWeights(); + v2Bias.releaseWeights(); + v3Mul.releaseWeights(); + v3Bias.releaseWeights(); + sv3Mul.releaseWeights(); + sv3Bias.releaseWeights(); + vOwnershipConv.releaseWeights(); } void ModelDesc::releaseWeights() { - releaseConv(trunk.initialConv); - releaseMatMul(trunk.initialMatMul); - if(trunk.metaEncoderVersion > 0) - releaseSGFEncoder(trunk.sgfMetadataEncoder); - releaseBlocks(trunk.blocks); - releaseBN(trunk.trunkTipBN); - releaseConv(policyHead.p1Conv); releaseConv(policyHead.g1Conv); - releaseBN(policyHead.g1BN); releaseMatMul(policyHead.gpoolToBiasMul); - releaseBN(policyHead.p1BN); releaseConv(policyHead.p2Conv); - releaseMatMul(policyHead.gpoolToPassMul); releaseMatBias(policyHead.gpoolToPassBias); - releaseMatMul(policyHead.gpoolToPassMul2); - releaseConv(valueHead.v1Conv); releaseBN(valueHead.v1BN); - releaseMatMul(valueHead.v2Mul); releaseMatBias(valueHead.v2Bias); - releaseMatMul(valueHead.v3Mul); releaseMatBias(valueHead.v3Bias); - releaseMatMul(valueHead.sv3Mul); releaseMatBias(valueHead.sv3Bias); - releaseConv(valueHead.vOwnershipConv); + trunk.releaseWeights(); + policyHead.releaseWeights(); + valueHead.releaseWeights(); } struct NonCopyingStreamBuf : public std::streambuf diff --git a/cpp/neuralnet/desc.h b/cpp/neuralnet/desc.h index 9536283f2..6b612207d 100644 --- a/cpp/neuralnet/desc.h +++ b/cpp/neuralnet/desc.h @@ -34,6 +34,8 @@ struct ConvLayerDesc { double getSpatialConvDepth() const; void scaleOutputChannels(const std::vector& scaling); + + void releaseWeights(); }; struct BatchNormLayerDesc { @@ -64,6 +66,8 @@ struct BatchNormLayerDesc { void extractChannelFactorsAbsLtOne(std::vector& channelFactors); void extractChannelFactorsAbsLtOneWithInverses(std::vector& channelFactors, std::vector& invChannelFactors); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct ActivationLayerDesc { @@ -99,6 +103,8 @@ struct MatMulLayerDesc { MatMulLayerDesc& operator=(MatMulLayerDesc&& other); void scaleOutputChannels(const std::vector& scaling); + + void releaseWeights(); }; struct MatBiasLayerDesc { @@ -115,6 +121,8 @@ struct MatBiasLayerDesc { MatBiasLayerDesc& operator=(MatBiasLayerDesc&& other); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct ResidualBlockDesc { @@ -140,6 +148,8 @@ struct ResidualBlockDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct GlobalPoolingResidualBlockDesc { @@ -170,6 +180,8 @@ struct GlobalPoolingResidualBlockDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct NestedBottleneckResidualBlockDesc { @@ -200,6 +212,8 @@ struct NestedBottleneckResidualBlockDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct SGFMetadataEncoderDesc { @@ -223,6 +237,8 @@ struct SGFMetadataEncoderDesc { SGFMetadataEncoderDesc& operator=(const SGFMetadataEncoderDesc&) = delete; SGFMetadataEncoderDesc& operator=(SGFMetadataEncoderDesc&& other); + + void releaseWeights(); }; @@ -263,6 +279,8 @@ struct TrunkDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct PolicyHeadDesc { @@ -296,6 +314,8 @@ struct PolicyHeadDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct ValueHeadDesc { @@ -327,6 +347,8 @@ struct ValueHeadDesc { void transformToReduceActivations(); void applyScale8ToReduceActivations(); + + void releaseWeights(); }; struct ModelPostProcessParams { From 98b17ebbf2459e00dcca8120e83e766aff335ab0 Mon Sep 17 00:00:00 2001 From: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Date: Sun, 31 May 2026 23:22:13 +0800 Subject: [PATCH 7/8] Co-locate releaseWeights() defs with each struct's other methods Move the 11 leaf/container releaseWeights() definitions in desc.cpp out of the bottom cluster (inherited from the old free-function layout) and place each immediately after its struct's last existing method, matching the file's per-struct grouping convention used by every other method. ModelDesc::releaseWeights() stays put, already adjacent to its siblings. Pure relocation: function bodies and desc.h are unchanged; only two stray double-blank lines were normalized to single. Verified clean Metal build, testgpuerror vs Eigen reference (g170-b6c96) at <0.0004% winrate error, and runtests all pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- cpp/neuralnet/desc.cpp | 236 ++++++++++++++++++++--------------------- 1 file changed, 117 insertions(+), 119 deletions(-) diff --git a/cpp/neuralnet/desc.cpp b/cpp/neuralnet/desc.cpp index 0c4943bb5..2fc47f340 100644 --- a/cpp/neuralnet/desc.cpp +++ b/cpp/neuralnet/desc.cpp @@ -193,6 +193,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) {} @@ -377,6 +381,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); } @@ -487,6 +500,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); } @@ -504,6 +521,10 @@ void MatBiasLayerDesc::applyScale8ToReduceActivations() { } } +void MatBiasLayerDesc::releaseWeights() { + std::vector().swap(weights); +} + //----------------------------------------------------------------------------- ResidualBlockDesc::ResidualBlockDesc() {} @@ -575,6 +596,13 @@ void ResidualBlockDesc::applyScale8ToReduceActivations() { midActivation.applyScale8ToReduceActivations(); } +void ResidualBlockDesc::releaseWeights() { + preBN.releaseWeights(); + regularConv.releaseWeights(); + midBN.releaseWeights(); + finalConv.releaseWeights(); +} + //----------------------------------------------------------------------------- GlobalPoolingResidualBlockDesc::GlobalPoolingResidualBlockDesc() {} @@ -685,6 +713,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() {} @@ -847,6 +885,30 @@ 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 { + ASSERT_UNREACHABLE; + } + } + postBN.releaseWeights(); + postConv.releaseWeights(); +} + //----------------------------------------------------------------------------- static void parseResidualBlockStack( @@ -1009,6 +1071,14 @@ SGFMetadataEncoderDesc& SGFMetadataEncoderDesc::operator=(SGFMetadataEncoderDesc return *this; } +void SGFMetadataEncoderDesc::releaseWeights() { + mul1.releaseWeights(); + bias1.releaseWeights(); + mul2.releaseWeights(); + bias2.releaseWeights(); + mul3.releaseWeights(); +} + //----------------------------------------------------------------------------- TrunkDesc::TrunkDesc() @@ -1259,6 +1329,30 @@ 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 { + ASSERT_UNREACHABLE; + } + } + trunkTipBN.releaseWeights(); +} //----------------------------------------------------------------------------- @@ -1406,6 +1500,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) {} @@ -1541,6 +1647,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(); +} //----------------------------------------------------------------------------- @@ -1783,125 +1900,6 @@ void ModelDesc::applyScale8ToReduceActivations() { postProcessParams.outputScaleMultiplier *= 8.0f; } -void ConvLayerDesc::releaseWeights() { - std::vector().swap(weights); -} - -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); -} - -void MatMulLayerDesc::releaseWeights() { - std::vector().swap(weights); -} - -void MatBiasLayerDesc::releaseWeights() { - std::vector().swap(weights); -} - -void ResidualBlockDesc::releaseWeights() { - preBN.releaseWeights(); - regularConv.releaseWeights(); - midBN.releaseWeights(); - finalConv.releaseWeights(); -} - -void GlobalPoolingResidualBlockDesc::releaseWeights() { - preBN.releaseWeights(); - regularConv.releaseWeights(); - gpoolConv.releaseWeights(); - gpoolBN.releaseWeights(); - gpoolToBiasMul.releaseWeights(); - midBN.releaseWeights(); - finalConv.releaseWeights(); -} - -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 { - ASSERT_UNREACHABLE; - } - } - postBN.releaseWeights(); - postConv.releaseWeights(); -} - -void SGFMetadataEncoderDesc::releaseWeights() { - mul1.releaseWeights(); - bias1.releaseWeights(); - mul2.releaseWeights(); - bias2.releaseWeights(); - mul3.releaseWeights(); -} - -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 { - ASSERT_UNREACHABLE; - } - } - trunkTipBN.releaseWeights(); -} - -void PolicyHeadDesc::releaseWeights() { - p1Conv.releaseWeights(); - g1Conv.releaseWeights(); - g1BN.releaseWeights(); - gpoolToBiasMul.releaseWeights(); - p1BN.releaseWeights(); - p2Conv.releaseWeights(); - gpoolToPassMul.releaseWeights(); - gpoolToPassBias.releaseWeights(); - gpoolToPassMul2.releaseWeights(); -} - -void ValueHeadDesc::releaseWeights() { - v1Conv.releaseWeights(); - v1BN.releaseWeights(); - v2Mul.releaseWeights(); - v2Bias.releaseWeights(); - v3Mul.releaseWeights(); - v3Bias.releaseWeights(); - sv3Mul.releaseWeights(); - sv3Bias.releaseWeights(); - vOwnershipConv.releaseWeights(); -} - void ModelDesc::releaseWeights() { trunk.releaseWeights(); policyHead.releaseWeights(); From 8481a9411854618befefdc0f25ff15402e2d6c70 Mon Sep 17 00:00:00 2001 From: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Date: Thu, 4 Jun 2026 12:04:44 +0800 Subject: [PATCH 8/8] Conform CoreML transformer derived consts to the owned-weight + FP32 contract The transformer attention builder emits four function-local std::vector tensors: RoPE cos/sin tables, the rotation matrix R, and per-head out-projection weight slices. After merging the transformer support onto the FloatView branch, these needed two fixes: 1. Dangling view. #1202 made WeightEntry::data a non-owning FloatView, so addConstOp registers a view whose backing buffer must outlive serialization. These locals were passed to addConstOp and would dangle once the build function returns (serialization runs afterwards). Route them through addOwnedConstOp so KataGoOps owns the buffer until serialization. (Under #1205's owning WeightEntry they were copied, so this only surfaces post-merge.) 2. dtype mismatch. emitConstOp declares each const's dtype as m_weight_dtype, but addOwnedConstOp / registerOwnedWeight stored at the global mode (is_fp32 hardcoded false). In an FP16 model these derived consts land in the attention / value-head FP32 sub-region (m_weight_dtype == FLOAT32), so they were declared FP32 but stored FP16. CoreML/ANE then rejects the model at load ("Metadata data type does not match requested type", BNNS error -14), which SIGABRT'd every FP16 ANE transformer. Thread is_fp32 through registerOwnedWeight and have addOwnedConstOp pass is_fp32 = (m_weight_dtype == FLOAT32), mirroring addConstOp so the stored dtype always matches the declared dtype. This also fixes the same latent mismatch for addLinearOp's transposed value-head weights. Verified with testgpuerror against fresh Eigen FP32 references: b7c96h3tfrs and b7c96h6gqa, which previously SIGABRT'd on the FP16 ANE path, now load and match to <0.0005% winrate; convnet ANE output is byte-identical and the Metal GPU path is unchanged. katago runtests and runnnlayertests also pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../katagocoreml/src/builder/MILBuilder.cpp | 21 +++++++++++++------ .../katagocoreml/src/builder/Operations.cpp | 4 +++- .../katagocoreml/src/builder/Operations.hpp | 6 ++++-- 3 files changed, 22 insertions(+), 9 deletions(-) diff --git a/cpp/external/katagocoreml/src/builder/MILBuilder.cpp b/cpp/external/katagocoreml/src/builder/MILBuilder.cpp index ba3db1a19..f86181d87 100644 --- a/cpp/external/katagocoreml/src/builder/MILBuilder.cpp +++ b/cpp/external/katagocoreml/src/builder/MILBuilder.cpp @@ -281,8 +281,12 @@ void MILBuilder::addOwnedConstOp(CoreML::Specification::MILSpec::Block* block, const std::string& name, std::vector&& data, const std::vector& shape) { - // Register derived weight; KataGoOps takes ownership of the buffer - m_ops.registerOwnedWeight(name, std::move(data), 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); } @@ -2230,10 +2234,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"); @@ -2363,7 +2370,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/Operations.cpp b/cpp/external/katagocoreml/src/builder/Operations.cpp index 5de42d09c..e86364943 100644 --- a/cpp/external/katagocoreml/src/builder/Operations.cpp +++ b/cpp/external/katagocoreml/src/builder/Operations.cpp @@ -28,7 +28,8 @@ std::string KataGoOps::registerWeight(const std::string& name, std::string KataGoOps::registerOwnedWeight(const std::string& name, std::vector&& data, - const std::vector& shape) { + const std::vector& shape, + bool is_fp32) { m_owned.push_back(std::move(data)); const std::vector& stored = m_owned.back(); WeightEntry entry; @@ -36,6 +37,7 @@ std::string KataGoOps::registerOwnedWeight(const std::string& 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; } diff --git a/cpp/external/katagocoreml/src/builder/Operations.hpp b/cpp/external/katagocoreml/src/builder/Operations.hpp index 1fb0d92a8..385648d19 100644 --- a/cpp/external/katagocoreml/src/builder/Operations.hpp +++ b/cpp/external/katagocoreml/src/builder/Operations.hpp @@ -82,10 +82,12 @@ class KataGoOps { const std::vector& shape) = delete; /// Register a derived/temporary weight; KataGoOps takes ownership so the - /// view stays valid through serialization. + /// 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); + const std::vector& shape, + bool is_fp32 = false); /// Get all registered weights (mutable; serialization sets blob_offset) std::vector& getWeightsMutable() { return m_weights; }