Batch Matrix Operations System

Overview

You need to implement a system for processing multiple matrices simultaneously using GPU acceleration. The system should perform batch matrix operations efficiently, handling multiple independent matrix computations in parallel.

Requirements

Implement a module matrix_batch_processor.py that provides functionality for:

1. Batch Matrix Multiplication: Given two batches of matrices (batch A and batch B), compute the matrix product for each corresponding pair. The batches are represented as 3D arrays where:

   • The first dimension represents the batch size (number of matrices)
   • The second and third dimensions represent the matrix dimensions
   • Example: shape (N, m, k) @ (N, k, n) → (N, m, n)

2. Batch Matrix Inversion: Given a batch of square invertible matrices, compute their inverses. The batch is represented as a 3D array where:

   • The first dimension is the batch size
   • Each matrix along the batch dimension must be square
   • Example: shape (N, d, d) → (N, d, d)

3. Performance Comparison: Provide a function that compares batch operations against sequential (loop-based) operations to demonstrate the performance benefits of batched operations.

Implementation Details

Your implementation should include:

• batch_matrix_multiply(batch_a, batch_b): Performs batched matrix multiplication

  • Input: Two 3D arrays where the first dimension is the batch size
  • Output: 3D array of matrix products

• batch_matrix_inverse(batch_matrices): Computes inverses for a batch of matrices

  • Input: 3D array where first dimension is batch size, and each matrix is square
  • Output: 3D array of inverted matrices

• compare_batch_vs_sequential(batch_size, matrix_size): Benchmarks batch vs sequential operations

  • Input: Number of matrices and their dimensions
  • Output: Dictionary with timing results for batch and sequential approaches

Test Cases { .tests }

Test 1: Batch Matrix Multiplication @test

Input:

import cupy as cp

# Batch of 3 matrices (3x2) multiplied by batch of 3 matrices (2x4)
batch_a = cp.array([
    [[1, 2], [3, 4], [5, 6]],
    [[2, 1], [4, 3], [6, 5]],
    [[1, 1], [2, 2], [3, 3]]
])

batch_b = cp.array([
    [[1, 0, 1, 0], [0, 1, 0, 1]],
    [[1, 1, 0, 0], [0, 0, 1, 1]],
    [[2, 0, 0, 2], [0, 2, 2, 0]]
])

result = batch_matrix_multiply(batch_a, batch_b)

Expected Output:

# Shape: (3, 3, 4)
# First matrix result: [[1, 2, 1, 2], [3, 4, 3, 4], [5, 6, 5, 6]]
# Second matrix result: [[2, 2, 1, 1], [4, 4, 3, 3], [6, 6, 5, 5]]
# Third matrix result: [[2, 2, 2, 2], [4, 4, 4, 4], [6, 6, 6, 6]]

Test 2: Batch Matrix Inversion @test

Input:

import cupy as cp

# Batch of 2 identity-like matrices
batch_matrices = cp.array([
    [[2.0, 0.0], [0.0, 2.0]],
    [[4.0, 0.0], [0.0, 4.0]]
])

result = batch_matrix_inverse(batch_matrices)

Expected Output:

# Shape: (2, 2, 2)
# First inverse: [[0.5, 0.0], [0.0, 0.5]]
# Second inverse: [[0.25, 0.0], [0.0, 0.25]]

Test 3: Performance Comparison @test

Input:

result = compare_batch_vs_sequential(batch_size=100, matrix_size=50)

Expected Output:

# Dictionary with keys 'batch_time' and 'sequential_time'
# batch_time should be significantly lower than sequential_time
# Example: {'batch_time': 0.05, 'sequential_time': 0.45}
assert result['batch_time'] < result['sequential_time']

Dependencies { .dependencies }

cupy { .dependency }

Provides GPU-accelerated array operations and numerical computing capabilities.

Files to Create

• matrix_batch_processor.py: Main implementation module
• test_matrix_batch_processor.py: Test file containing the test cases above