CUDA_test

'''
This function generates two random square matrices of dimensions (size, size) using NumPy and performs matrix multiplication (@ is the Python 3.x operator for matrix multiplication). 
It returns the resultant matrix c.

Function: test_gpu(size=5000, chunk_size=1000, runs=500)
This function tests the GPU's performance in performing matrix multiplications.

Device Detection: It first checks whether a CUDA-compatible GPU is available. If so, it uses it; otherwise, it falls back to the CPU.
Initialization: times is an empty list that will hold the time taken for each GPU run.
The Loop: Runs the test runs times.
CPU Work: Calls do_cpu_work() to perform some CPU-intensive tasks, but the results aren't actually used in the subsequent steps.
Start Timer: Records the start time.
Chunk Generation: The script creates smaller "chunks" of tensors for matrix multiplication. The tensors are directly moved to the device (CPU or GPU).
Matrix Multiplication: For each pair of chunks, matrix multiplication is performed twice, each in its own CUDA stream. The results are appended to c1_list and c2_list.
Concatenation: The chunks are concatenated to form the final matrices c1 and c2.
End Timer: Records the end time and calculates the time taken for the GPU operations.
Average Time: Finally, it calculates the average time taken for all the runs.
''';


import torch
import time 
import numpy as np

def do_cpu_work(size=10000):
    a = np.random.rand(size, size)
    b = np.random.rand(size, size)
    c = a @ b
    return c

def test_gpu(size=5000, chunk_size=1000, runs=500):

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    times = []
    for i in range(runs):

        # CPU work
        a = do_cpu_work()
        b = do_cpu_work()

        start_time = time.time()

        # Generate matrices in smaller chunks
        b_chunks = [torch.randn(chunk_size, chunk_size, device=device) for _ in range(size // chunk_size)]
        a_chunks = [torch.randn(chunk_size, chunk_size, device=device) for _ in range(size // chunk_size)]

        c1_list = []
        c2_list = []

        for a_chunk, b_chunk in zip(a_chunks, b_chunks):
            with torch.cuda.stream(torch.cuda.Stream()):
                c_chunk = torch.matmul(a_chunk, b_chunk)
                c1_list.append(c_chunk)

            with torch.cuda.stream(torch.cuda.Stream()):
                c_chunk = torch.matmul(a_chunk, b_chunk)
                c2_list.append(c_chunk)

        c1 = torch.cat(c1_list)
        c2 = torch.cat(c2_list)

        end_time = time.time()

        times.append(end_time - start_time)
        print(f"Finished run {i+1}/{runs}", end='\r')

    avg_time = sum(times) / len(times)
    print(f"\nAverage time per run: {avg_time:.4f} seconds")

if __name__ == "__main__":
    test_gpu(size=5000, chunk_size=1000, runs=500)