File size: 6,945 Bytes
7a0c684 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 |
"""
Virtual GPU interface with modality support.
Provides a high-level interface for multi-modal tensor operations.
"""
from typing import Optional, Dict, Any, List, Union, Tuple
import numpy as np
from .modality import ModalityType, ModalityConfig
from .modality_aware_tensor_core import ModalityAwareTensorCore
class ModalityAwareVirtualGPU:
"""High-level interface for modality-aware virtual GPU operations"""
def __init__(
self,
num_tensor_cores: int = 1,
bits: int = 2,
memory_size: Optional[int] = None,
bandwidth_tbps: float = 10000
):
# Initialize tensor cores
self.tensor_cores = [
ModalityAwareTensorCore(
bits=bits,
memory_size=memory_size,
bandwidth_tbps=bandwidth_tbps
)
for _ in range(num_tensor_cores)
]
# Track modality assignments to cores
self.modality_core_map: Dict[ModalityType, int] = {}
self.core_modality_map: Dict[int, ModalityType] = {}
# Initialize modality-specific memory pools
self.modality_memory: Dict[ModalityType, Dict[str, np.ndarray]] = {
modality: {} for modality in ModalityType
}
def assign_modality_to_core(
self,
modality: ModalityType,
core_idx: int
) -> None:
"""Assign a modality to a specific tensor core"""
if core_idx >= len(self.tensor_cores):
raise ValueError(f"Invalid core index: {core_idx}")
# Update mappings
self.modality_core_map[modality] = core_idx
self.core_modality_map[core_idx] = modality
# Configure the core
self.tensor_cores[core_idx].set_modality(modality)
def get_core_for_modality(
self,
modality: ModalityType
) -> Optional[ModalityAwareTensorCore]:
"""Get the tensor core assigned to a modality"""
core_idx = self.modality_core_map.get(modality)
if core_idx is not None:
return self.tensor_cores[core_idx]
return None
def matmul(
self,
x: np.ndarray,
y: np.ndarray,
x_modality: ModalityType,
y_modality: Optional[ModalityType] = None
) -> np.ndarray:
"""Matrix multiplication with modality handling"""
y_mod = y_modality or x_modality
core = self.get_core_for_modality(x_modality)
if core is None:
raise ValueError(f"No core assigned for modality: {x_modality}")
return core.matmul(x, y, x_modality, y_mod)
def conv(
self,
x: np.ndarray,
weight: np.ndarray,
stride: Union[int, Tuple[int, ...]] = 1,
padding: Union[int, Tuple[int, ...]] = 0,
modality: ModalityType = None
) -> np.ndarray:
"""Convolution operation"""
core = self.get_core_for_modality(modality)
if core is None:
raise ValueError(f"No core assigned for modality: {modality}")
return core.conv(x, weight, stride, padding, modality)
def attention(
self,
q: np.ndarray,
k: np.ndarray,
v: np.ndarray,
mask: Optional[np.ndarray] = None,
modality: ModalityType = None
) -> np.ndarray:
"""Attention operation"""
core = self.get_core_for_modality(modality)
if core is None:
raise ValueError(f"No core assigned for modality: {modality}")
return core.attention(q, k, v, mask, modality)
def pool(
self,
x: np.ndarray,
kernel_size: Union[int, Tuple[int, ...]],
stride: Optional[Union[int, Tuple[int, ...]]] = None,
mode: str = 'max',
modality: ModalityType = None
) -> np.ndarray:
"""Pooling operation"""
core = self.get_core_for_modality(modality)
if core is None:
raise ValueError(f"No core assigned for modality: {modality}")
return core.pool(x, kernel_size, stride, mode, modality)
def normalize(
self,
x: np.ndarray,
modality: ModalityType = None,
eps: float = 1e-5
) -> np.ndarray:
"""Normalization operation"""
core = self.get_core_for_modality(modality)
if core is None:
raise ValueError(f"No core assigned for modality: {modality}")
return core.normalize(x, modality, eps)
def fuse_modalities(
self,
x: np.ndarray,
y: np.ndarray,
x_modality: ModalityType,
y_modality: ModalityType
) -> np.ndarray:
"""Fuse tensors from different modalities"""
x_core = self.get_core_for_modality(x_modality)
if x_core is None:
raise ValueError(f"No core assigned for modality: {x_modality}")
return x_core.fuse_modalities(x, y, x_modality, y_modality)
def unfuse_modalities(
self,
z: np.ndarray,
x_modality: ModalityType,
y_modality: ModalityType
) -> Tuple[np.ndarray, np.ndarray]:
"""Separate fused tensor back into modalities"""
x_core = self.get_core_for_modality(x_modality)
if x_core is None:
raise ValueError(f"No core assigned for modality: {x_modality}")
return x_core.unfuse_modalities(z, x_modality, y_modality)
def store_tensor(
self,
data: np.ndarray,
modality: ModalityType,
key: str
) -> None:
"""Store tensor in modality-specific memory"""
core = self.get_core_for_modality(modality)
if core is None:
raise ValueError(f"No core assigned for modality: {modality}")
core.store_modality_tensor(data, modality, key)
def load_tensor(
self,
modality: ModalityType,
key: str
) -> Optional[np.ndarray]:
"""Load tensor from modality-specific memory"""
core = self.get_core_for_modality(modality)
if core is None:
raise ValueError(f"No core assigned for modality: {modality}")
return core.load_modality_tensor(modality, key)
def clear_memory(
self,
modality: Optional[ModalityType] = None
) -> None:
"""Clear modality-specific memory"""
if modality:
core = self.get_core_for_modality(modality)
if core is not None:
core.clear_modality_memory(modality)
else:
for core in self.tensor_cores:
core.clear_modality_memory()
|