fix(graph): preserve recurrent graph inputs and allocation lifetimes

src/infinicore-test/memory_test.cc

@@ -1,4 +1,5 @@
 #include "memory_test.h"
+#include "../infinicore/context/allocators/pinnable_block_allocator.hpp"
 #include <algorithm>
 #include <cstring>
 #include <random>
@@ -78,6 +79,27 @@ TestResult BasicMemoryTest::run() {
             }
             spdlog::debug("BasicMemoryTest: Pinned memory test completed");
 
+            PinnableBlockAllocator allocator(current_device);
+            auto reinstantiated = allocator.allocate(1024);
+            allocator.deallocate(reinstantiated);
+            allocator.mark_in_use_(reinstantiated, true);
+            allocator.trim();
+            allocator.deallocate(reinstantiated);
+            allocator.trim();
+
+            auto captured = allocator.allocate(1024);
+            allocator.deallocate(captured);
+            allocator.set_pin_mode(true);
+            if (allocator.allocate(1024) != captured) {
+                return false;
+            }
+            allocator.deallocate(captured);
+            allocator.set_pin_mode(false);
+            allocator.trim();
+            // A graph may reinstantiate storage after its temporary owner expires.
+            allocator.mark_in_use_(captured, true);
+            allocator.deallocate(captured);
+
             return true;
         } catch (const std::exception &e) {
             std::cerr << "BasicMemoryTest failed with exception: " << e.what() << std::endl;

src/infinicore/context/allocators/pinnable_block_allocator.cc

@@ -64,6 +64,7 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) {
                     cls.free_blocks.pop_back();
                     block->in_use = true;
                     block->use_count = 1;
+                    block->frozen = block->frozen || pinned_mode_;
                     return reinterpret_cast<std::byte *>(block->ptr);
                 }
             }
@@ -151,6 +152,16 @@ size_t PinnableBlockAllocator::mark_in_use_(void *ptr, bool in_use) {
 
     auto block = it->second;
     if (in_use) {
+        if (!block->in_use) {
+            // Reinstantiated graph storage is no longer available for trimming.
+            for (auto &cls : size_classes_) {
+                if (block->size == cls.block_size) {
+                    auto &free = cls.free_blocks;
+                    free.erase(std::remove(free.begin(), free.end(), block), free.end());
+                    break;
+                }
+            }
+        }
         block->in_use = true;
         ++block->use_count;
     } else if (block->use_count > 0) {
@@ -166,7 +177,7 @@ void PinnableBlockAllocator::trim() {
     // Free non-frozen size-class blocks
     for (auto &cls : size_classes_) {
         for (auto it = cls.free_blocks.begin(); it != cls.free_blocks.end();) {
-            if (!(*it)->frozen) {
+            if (!(*it)->frozen && !(*it)->in_use) {
                 INFINICORE_CHECK_ERROR(infinirtFree((*it)->ptr));
                 all_blocks_.erase((*it)->ptr);
                 it = cls.free_blocks.erase(it);

src/infinicore/ops/cat/cat.cc

@@ -12,6 +12,7 @@ bool use_slice_copy_cat(Device::Type device_type, int dim, int ndim) {
     // correct through copy_from, but their performance impact is unverified.
     return dim == ndim - 1
         && (device_type == Device::Type::NVIDIA
+            || device_type == Device::Type::METAX
             || device_type == Device::Type::HYGON
             || device_type == Device::Type::ILUVATAR
             || device_type == Device::Type::ALI);
@@ -172,7 +173,8 @@ void cat_(Tensor out, std::vector<Tensor> tensors, int dim) {
         // index. Concatenating MLA tensors on the last dimension can therefore
         // enqueue hundreds of tiny D2D copy calls per layer. A strided output
         // slice is semantically identical and copy_from lowers it to a single
-        // rearrange kernel per input tensor on CUDA-like backends.
+        // rearrange kernel per input tensor on CUDA-like backends. Unlike
+        // direct memcpy calls, these copies are also recorded for graph replay.
         size_t offset = 0;
         for (auto &tensor : tensors) {
             if (tensor->ndim() == 1) {

src/infinicore/ops/float_power/float_power.cc

@@ -1,7 +1,25 @@
 #include "infinicore/ops/float_power.hpp"
+#include "infinicore/graph/graph.hpp"
 #include "infinicore/tensor.hpp"
 
 namespace infinicore::op {
+namespace {
+
+class RecordedScalarPower final : public graph::GraphOperator {
+public:
+    RecordedScalarPower(Tensor output, Tensor input, double exponent)
+        : output_(output), input_(input), exponent_(exponent) {}
+
+    void run() const override {
+        FloatPower::dispatcher_scalar().lookup(input_->device().getType())(output_, input_, exponent_);
+    }
+
+private:
+    graph::GraphTensor output_, input_;
+    double exponent_;
+};
+
+} // namespace
 
 // =======================================================================
 // 1. Dispatcher 单例
@@ -22,6 +40,10 @@ common::OpDispatcher<FloatPower::schema_tensor> &FloatPower::dispatcher_tensor()
 // =======================================================================
 
 void FloatPower::execute(Tensor output, Tensor input, double exponent) {
+    if (context::isGraphRecording()) {
+        context::addGraphOperator(std::make_shared<RecordedScalarPower>(output, input, exponent));
+        return;
+    }
     dispatcher_scalar()
         .lookup(context::getDevice().getType())(output, input, exponent);
 }

src/infinicore/ops/sum/sum_infiniop.cc

@@ -1,47 +1,44 @@
-#include "../../utils.hpp"
-#include "infinicore/common/hash.hpp"
-#include "infinicore/ops/common/cache.hpp"
+#include "../infiniop_impl.hpp"
+#include "infinicore/graph/graph.hpp"
 #include "infinicore/ops/sum.hpp"
-#include <infiniop.h>
 
 namespace infinicore::op::sum_impl::infiniop {
 
-thread_local common::OpCache<size_t, infiniopSumDescriptor_t> caches(
-    100, // capacity
-    [](infiniopSumDescriptor_t &desc) {
-        if (desc != nullptr) {
-            INFINICORE_CHECK_ERROR(infiniopDestroySumDescriptor(desc));
-            desc = nullptr;
-        }
-    });
-
-void calculate(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim) {
-    size_t seed = hash_combine(output, input, dim.size(), keepdim);
-
-    auto device_type = context::getDevice().getType();
-    auto device_index = context::getDevice().getIndex();
+INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Sum, 100);
 
-    auto &cache = caches.getCache(device_type, device_index);
-
-    auto desc_opt = cache.get(seed);
-    infiniopSumDescriptor_t desc = nullptr;
+class RecordedSum final : public graph::GraphOperator {
+public:
+    RecordedSum(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim)
+        : output_(output), input_(input), dim_(std::move(dim)), keepdim_(keepdim) {
+        size_t seed = hash_combine(output, input, dim_.size(), keepdim_);
+        for (auto axis : dim_) {
+            hash_combine(seed, axis);
+        }
+        INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(
+            Descriptor, descriptor, Sum, seed,
+            output->desc(), input->desc(), dim_.data(), dim_.size(), keepdim_);
+        INFINIOP_WORKSPACE_TENSOR(workspace, Sum, descriptor);
+        descriptor_ = std::move(descriptor);
+        workspace_ = std::make_unique<graph::GraphTensor>(workspace);
+    }
 
-    if (!desc_opt) {
-        INFINICORE_CHECK_ERROR(infiniopCreateSumDescriptor(
-            context::getInfiniopHandle(output->device()), &desc,
-            output->desc(), input->desc(), dim.data(), dim.size(), keepdim));
-        cache.put(seed, desc);
-    } else {
-        desc = *desc_opt;
+    void run() const override {
+        auto output = output_;
+        INFINICORE_CHECK_ERROR(infiniopSum(
+            descriptor_->desc, (*workspace_)->data(), (*workspace_)->numel(),
+            output->data(), input_->data(), dim_.data(), dim_.size(), keepdim_, context::getStream()));
     }
 
-    size_t workspace_size = 0;
-    INFINICORE_CHECK_ERROR(infiniopGetSumWorkspaceSize(desc, &workspace_size));
-    std::shared_ptr<Memory> workspace = context::allocateMemory(workspace_size);
+private:
+    graph::GraphTensor output_, input_;
+    mutable std::vector<size_t> dim_;
+    bool keepdim_;
+    std::shared_ptr<Descriptor> descriptor_;
+    std::unique_ptr<graph::GraphTensor> workspace_;
+};
 
-    INFINICORE_CHECK_ERROR(infiniopSum(
-        desc, workspace->data(), workspace_size,
-        output->data(), input->data(), dim.data(), dim.size(), keepdim, context::getStream()));
+void calculate(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(RecordedSum, output, input, std::move(dim), keepdim);
 }
 
 static bool registered = []() {

src/infiniop/ops/sum/nvidia/sum_nvidia.cu

@@ -59,8 +59,7 @@ infiniStatus_t launchKernel(
     CHECK_CUDA(cudaMemcpyAsync(output_strides_cuda, info.output_strides.data(), output_ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, stream));
 
     if (info.reduce_num == input_size) {
-        T zero = static_cast<T>(0.0f);
-        CHECK_CUDA(cudaMemcpyAsync(output, &zero, sizeof(T), cudaMemcpyHostToDevice, stream));
+        CHECK_CUDA(cudaMemsetAsync(output, 0, sizeof(T), stream));
         size_t grid_size = (input_size + BLOCK_SIZE - 1) / BLOCK_SIZE;
         sumAllKernel<BLOCK_SIZE, T, T><<<grid_size, BLOCK_SIZE, BLOCK_SIZE * sizeof(T), stream>>>(
             output, input, input_size, input_ndim, permuted_input_shape_cuda, permuted_input_strides_cuda);

test/infinicore/graph/test_cat.py

@@ -0,0 +1,48 @@
+"""Last-axis concatenation must replay its producers and strided copies."""
+
+import pytest
+import torch
+
+import infinicore
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="A GPU is required.")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("batch", [1, 4])
+def test_last_axis_cat_replays_changed_inputs(dtype, batch):
+    device = infinicore.device("cuda", 0)
+    infinicore.set_device(device)
+    source = (
+        torch.linspace(-1, 1, batch * 129, device="cuda")
+        .reshape(1, batch, 129)
+        .to(dtype)
+    )
+    # Match projection views with a larger row stride than the sliced width.
+    inputs = [source[..., :64], source[..., 64:96], source[..., 96:]]
+    tensors = [
+        infinicore.strided_from_blob(
+            tensor.data_ptr(),
+            list(tensor.shape),
+            list(tensor.stride()),
+            dtype=infinicore.utils.to_infinicore_dtype(dtype),
+            device=device,
+        )
+        for tensor in inputs
+    ]
+    output = torch.empty_like(source)
+    target = infinicore.from_torch(output)
+    torch.cuda.synchronize()
+
+    infinicore.start_graph_recording(device)
+    projected = infinicore.mul(tensors[0], tensors[0])
+    joined = infinicore.cat([projected, *tensors[1:]], dim=-1)
+    infinicore.mul(joined, joined, out=target)
+    graph = infinicore.stop_graph_recording()
+
+    for scale in (1.0, -0.5, 2.0):
+        source.mul_(scale)
+        torch.cuda.synchronize()
+        graph.run()
+        infinicore.sync_stream()
+        expected = torch.cat([inputs[0].square(), *inputs[1:]], dim=-1).square()
+        torch.testing.assert_close(output, expected)

test/infinicore/graph/test_sum_scalar_power.py

@@ -0,0 +1,45 @@
+"""Exercise the reduction and reciprocal norm used by tensor-parallel models."""
+
+import pytest
+import torch
+
+import infinicore
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="A CUDA device is required.")
+def test_sum_descriptor_cache_distinguishes_axes():
+    source = torch.arange(16, device="cuda", dtype=torch.float32).reshape(4, 4)
+    tensor = infinicore.from_torch(source)
+    for axis in (0, 1, 0):
+        output = torch.empty(4, device="cuda")
+        torch.cuda.synchronize()
+        infinicore.sum(tensor, dim=axis, out=infinicore.from_torch(output))
+        infinicore.sync_device()
+        torch.testing.assert_close(output, source.sum(axis))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="A CUDA device is required.")
+@pytest.mark.parametrize("keepdim", [False, True])
+@pytest.mark.parametrize("batch", [1, 4])
+def test_sum_scalar_power_replay(keepdim, batch):
+    device = infinicore.device("cuda", 0)
+    infinicore.set_device(device)
+    source = torch.linspace(0.25, 2.0, batch * 768, device="cuda").reshape(batch, 768)
+    output = torch.empty((batch, 1) if keepdim else (batch,), device="cuda")
+    input_tensor = infinicore.from_torch(source)
+    output_tensor = infinicore.from_torch(output)
+    torch.cuda.synchronize()
+
+    infinicore.start_graph_recording(device)
+    squared = infinicore.mul(input_tensor, input_tensor)
+    reduced = infinicore.sum(squared, dim=1, keepdim=keepdim)
+    infinicore.float_power(reduced, -0.5, out=output_tensor)
+    graph = infinicore.stop_graph_recording()
+
+    for scale in (1.0, 0.5, 3.0):
+        source.mul_(scale)
+        torch.cuda.synchronize()
+        graph.run()
+        infinicore.sync_stream()
+        expected = source.square().sum(dim=1, keepdim=keepdim).rsqrt()
+        torch.testing.assert_close(output, expected, atol=1e-6, rtol=1e-5)