diff --git a/src/infinicore-test/memory_test.cc b/src/infinicore-test/memory_test.cc index 2029e0a81..59f753ba2 100644 --- a/src/infinicore-test/memory_test.cc +++ b/src/infinicore-test/memory_test.cc @@ -1,4 +1,5 @@ #include "memory_test.h" +#include "../infinicore/context/allocators/pinnable_block_allocator.hpp" #include #include #include @@ -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; diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.cc b/src/infinicore/context/allocators/pinnable_block_allocator.cc index 32e5c5e9b..a4f2f7273 100644 --- a/src/infinicore/context/allocators/pinnable_block_allocator.cc +++ b/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(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); diff --git a/src/infinicore/ops/cat/cat.cc b/src/infinicore/ops/cat/cat.cc index 217ac7ec9..c48e48718 100644 --- a/src/infinicore/ops/cat/cat.cc +++ b/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 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) { diff --git a/src/infinicore/ops/float_power/float_power.cc b/src/infinicore/ops/float_power/float_power.cc index c3bf5003b..a8e17b514 100644 --- a/src/infinicore/ops/float_power/float_power.cc +++ b/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::dispatcher_tensor() // ======================================================================= void FloatPower::execute(Tensor output, Tensor input, double exponent) { + if (context::isGraphRecording()) { + context::addGraphOperator(std::make_shared(output, input, exponent)); + return; + } dispatcher_scalar() .lookup(context::getDevice().getType())(output, input, exponent); } diff --git a/src/infinicore/ops/sum/sum_infiniop.cc b/src/infinicore/ops/sum/sum_infiniop.cc index 700ab04fa..c1f960959 100644 --- a/src/infinicore/ops/sum/sum_infiniop.cc +++ b/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 namespace infinicore::op::sum_impl::infiniop { -thread_local common::OpCache caches( - 100, // capacity - [](infiniopSumDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySumDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, std::vector 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 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(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 workspace = context::allocateMemory(workspace_size); +private: + graph::GraphTensor output_, input_; + mutable std::vector dim_; + bool keepdim_; + std::shared_ptr descriptor_; + std::unique_ptr 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 dim, bool keepdim) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(RecordedSum, output, input, std::move(dim), keepdim); } static bool registered = []() { diff --git a/src/infiniop/ops/sum/nvidia/sum_nvidia.cu b/src/infiniop/ops/sum/nvidia/sum_nvidia.cu index 9f165d271..224a57f6c 100644 --- a/src/infiniop/ops/sum/nvidia/sum_nvidia.cu +++ b/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(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<<>>( output, input, input_size, input_ndim, permuted_input_shape_cuda, permuted_input_strides_cuda); diff --git a/test/infinicore/graph/test_cat.py b/test/infinicore/graph/test_cat.py new file mode 100644 index 000000000..1448755ac --- /dev/null +++ b/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) diff --git a/test/infinicore/graph/test_sum_scalar_power.py b/test/infinicore/graph/test_sum_scalar_power.py new file mode 100644 index 000000000..0cd9ed442 --- /dev/null +++ b/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)