Reduce CoreML conversion peak and ANE steady-state memory

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<WeightEntry> weights_copy(weights.begin(), weights.end());
+    // Serialize directly from the builder's weight views (no copy).
+    std::vector<WeightEntry>& 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) {

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<float>& data,
                             const std::vector<int64_t>& 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<float>&& data,
+                                 const std::vector<int64_t>& 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<int64_t>& 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<int64_t> transposed_shape = {static_cast<int64_t>(out_ch), static_cast<int64_t>(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<int64_t> bias_shape = {static_cast<int64_t>(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<size_t>(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) {

cpp/external/katagocoreml/src/builder/MILBuilder.hpp

@@ -29,8 +29,8 @@ class MILBuilder {
     /// @return Unique pointer to MIL Program protobuf
     std::unique_ptr<CoreML::Specification::MILSpec::Program> build();
 
-    /// Get weight entries for blob serialization
-    const std::vector<WeightEntry>& getWeights() const { return m_ops.getWeights(); }
+    /// Get weight entries for blob serialization (mutable; serialization sets blob_offset)
+    std::vector<WeightEntry>& 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<float>& data,
                     const std::vector<int64_t>& 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<float>&& data,
+                    const std::vector<int64_t>& shape) = delete;
+
+    void addOwnedConstOp(CoreML::Specification::MILSpec::Block* block,
+                         const std::string& name,
+                         std::vector<float>&& data,
+                         const std::vector<int64_t>& shape);
+
+    void emitConstOp(CoreML::Specification::MILSpec::Block* block,
+                     const std::string& name,
+                     const std::vector<int64_t>& shape);
+
     void addIntArrayConstOp(CoreML::Specification::MILSpec::Block* block,
                             const std::string& name,
                             const std::vector<int32_t>& values);

cpp/external/katagocoreml/src/builder/Operations.cpp

@@ -18,14 +18,30 @@ 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;
     m_weights.push_back(std::move(entry));
     return name;
 }
 
+std::string KataGoOps::registerOwnedWeight(const std::string& name,
+                                            std::vector<float>&& data,
+                                            const std::vector<int64_t>& shape,
+                                            bool is_fp32) {
+    m_owned.push_back(std::move(data));
+    const std::vector<float>& 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++);
 }

cpp/external/katagocoreml/src/builder/Operations.hpp

@@ -5,15 +5,29 @@
 
 #include "../types/KataGoTypes.hpp"
 #include <cmath>
+#include <deque>
 #include <string>
 #include <vector>
 
 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<float> data;
+    FloatView data;            // non-owning view (replaces raw ptr + count)
     std::vector<int64_t> 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<float>& data,
                                const std::vector<int64_t>& shape,
                                bool is_fp32 = false);
 
-    /// Get all registered weights
-    const std::vector<WeightEntry>& 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<float>&& data,
+                               const std::vector<int64_t>& 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<float>&& data,
+                                    const std::vector<int64_t>& shape,
+                                    bool is_fp32 = false);
+
+    /// Get all registered weights (mutable; serialization sets blob_offset)
+    std::vector<WeightEntry>& 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<WeightEntry> m_weights;
+    std::deque<std::vector<float>> m_owned;
     int m_op_counter = 0;
 };