#16134: add test case for pre-allocated CreateBuffer

tenstorrent · Dec 18, 2024 · 99fbfe1 · 99fbfe1
1 parent 05fee29
commit 99fbfe1
Showing 1 changed file with 66 additions and 0 deletions.
diff --git a/tests/tt_metal/tt_metal/api/test_dram.cpp b/tests/tt_metal/tt_metal/api/test_dram.cpp
@@ -122,6 +122,54 @@ bool dram_single_core(
     fixture->ReadBuffer(device, output_dram_buffer, result_vec);
     return result_vec == src_vec;
 }
+
+bool dram_single_core_pre_allocated(
+    DispatchFixture* fixture, tt_metal::Device* device, const DRAMConfig& cfg, std::vector<uint32_t> src_vec) {
+    // Create a program
+    tt_metal::Program program = CreateProgram();
+
+    tt_metal::InterleavedBufferConfig dram_config{
+        .device = device,
+        .size = cfg.dram_buffer_size,
+        .page_size = cfg.dram_buffer_size,
+        .buffer_type = tt_metal::BufferType::DRAM};
+
+    auto input_dram_buffer = tt_metal::CreateBuffer(dram_config);
+    uint32_t input_dram_buffer_addr = input_dram_buffer->address();
+    auto input_dram_pre_allocated_buffer = tt_metal::CreateBuffer(dram_config, input_dram_buffer_addr);
+    uint32_t input_dram_pre_allocated_buffer_addr = input_dram_pre_allocated_buffer->address();
+
+    TT_FATAL(input_dram_buffer_addr == input_dram_pre_allocated_buffer_addr, "Error");
+
+    auto output_dram_buffer = tt_metal::CreateBuffer(dram_config);
+    uint32_t output_dram_buffer_addr = output_dram_buffer->address();
+    auto output_dram_pre_allocated_buffer = tt_metal::CreateBuffer(dram_config, output_dram_buffer_addr);
+    uint32_t output_dram_pre_allocated_buffer_addr = output_dram_pre_allocated_buffer->address();
+
+    TT_FATAL(output_dram_buffer_addr == output_dram_pre_allocated_buffer_addr, "Error");
+
+    // Create the kernel
+    auto dram_kernel = tt_metal::CreateKernel(program, cfg.kernel_file, cfg.core_range, cfg.data_movement_cfg);
+    fixture->WriteBuffer(device, input_dram_pre_allocated_buffer, src_vec);
+
+    tt_metal::SetRuntimeArgs(
+        program,
+        dram_kernel,
+        cfg.core_range,
+        {cfg.l1_buffer_addr,
+         input_dram_pre_allocated_buffer_addr,
+         (std::uint32_t)0,
+         output_dram_pre_allocated_buffer_addr,
+         (std::uint32_t)0,
+         cfg.dram_buffer_size});
+
+    fixture->RunProgram(device, program);
+
+    std::vector<uint32_t> result_vec;
+    fixture->ReadBuffer(device, output_dram_pre_allocated_buffer, result_vec);
+
+    return result_vec == src_vec;
+}
 }  // namespace unit_tests_common::dram::test_dram
 
 TEST_F(DispatchFixture, TensixDRAMLoopbackSingleCore) {
@@ -142,6 +190,24 @@ TEST_F(DispatchFixture, TensixDRAMLoopbackSingleCore) {
     }
 }
 
+TEST_F(DispatchFixture, TensixDRAMLoopbackSingleCorePreAllocated) {
+    uint32_t buffer_size = 2 * 1024 * 25;
+    std::vector<uint32_t> src_vec =
+        create_random_vector_of_bfloat16(buffer_size, 100, std::chrono::system_clock::now().time_since_epoch().count());
+    unit_tests_common::dram::test_dram::DRAMConfig dram_test_config = {
+        .core_range = {{0, 0}, {0, 0}},
+        .kernel_file = "tests/tt_metal/tt_metal/test_kernels/dataflow/dram_copy.cpp",
+        .dram_buffer_size = buffer_size,
+        .l1_buffer_addr = 400 * 1024,
+        .data_movement_cfg =
+            {.processor = tt_metal::DataMovementProcessor::RISCV_0, .noc = tt_metal::NOC::RISCV_0_default},
+    };
+    for (unsigned int id = 0; id < devices_.size(); id++) {
+        ASSERT_TRUE(unit_tests_common::dram::test_dram::dram_single_core_pre_allocated(
+            this, devices_.at(id), dram_test_config, src_vec));
+    }
+}
+
 TEST_F(DispatchFixture, TensixDRAMLoopbackSingleCoreDB) {
     if (!this->IsSlowDispatch()) {
         tt::log_info(tt::LogTest, "This test is only supported in slow dispatch mode");