Add support for slicing

Tensor Slicing API: Zero-Copy Views Strategy

Memory Layout: NCHW with Row-Major Storage

SKaiNET uses NCHW (batch, channels, height, width) memory layout with row-major storage for optimal performance and compatibility with deep learning frameworks:

NCHW Format: Tensors are organized as [batch, channels, height, width]
Row-Major Storage: Data is stored contiguously in memory with the rightmost dimension changing fastest
Stride Pattern: For a 4D tensor [N,C,H,W], strides are [C×H×W, H×W, W, 1]

This layout choice affects how zero-copy tensor views can be optimized, particularly for common ML operations like channel extraction and spatial region slicing.

Current Status

Existing Implementation

The SKaiNET tensor library currently provides:

Slicing DSL: A comprehensive numpy-like slicing syntax in TensorsRangeBuilder.kt
Slice Descriptors: Data classes that describe slice operations but don't create actual tensor views
Dense Tensors: CPU-based tensors (CpuTensorFP32, CpuTensorInt8, CpuTensorInt32) that store data in native arrays
Copy-Based Operations: All tensor operations currently create new tensors by copying data

Current Limitations

Memory Inefficiency: Every slice operation would require copying data into new tensors
No Zero-Copy Views: Slice descriptors exist but don't create actual tensor views
Missing Implementation: The slicing DSL returns List<Slice> but no mechanism to materialize sliced tensors

Zero-Copy Tensor Strategy

Core Concept

Create lightweight tensor wrappers that reference the original tensor's data with different indexing logic, avoiding data copying for common slicing operations.

Design Patterns

1. Tensor View Pattern

interface TensorView<T : DType, V> : Tensor<T, V> {
    val parentTensor: Tensor<T, V>
    val viewShape: Shape
    val indexMapping: IndexMapper
}

class SlicedTensorView<T : DType, V>(
    override val parentTensor: Tensor<T, V>,
    private val slices: List<Slice<T, V>>
) : TensorView<T, V> {
    override val shape: Shape by lazy { computeSlicedShape() }
    
    override fun get(vararg indices: Int): V {
        val mappedIndices = mapIndicesToParent(indices)
        return parentTensor.get(*mappedIndices)
    }
}

2. Index Mapping Strategy

interface IndexMapper {
    fun mapToParent(childIndices: IntArray): IntArray
    fun isContiguous(): Boolean
    fun getStride(): IntArray
}

class SliceIndexMapper(
    private val slices: List<Slice<*, *>>,
    private val parentShape: Shape
) : IndexMapper {
    override fun mapToParent(childIndices: IntArray): IntArray {
        // Map child tensor indices to parent tensor indices
        // considering slice start/end positions and dimension ordering
    }
}

3. Zero-Copy Operations

// Extension functions for zero-copy slicing
fun <T : DType, V> Tensor<T, V>.slice(slices: List<Slice<T, V>>): TensorView<T, V> {
    return SlicedTensorView(this, slices)
}

// Integrate with existing DSL
fun <T : DType, V> Tensor<T, V>.sliceView(builder: TensorSliceBuilder<T, V>.() -> Unit): TensorView<T, V> {
    val slices = slice(this, builder)  // Use existing DSL
    return SlicedTensorView(this, slices)
}

Common ML Use Cases

1. Batch Processing

Use Case: Processing mini-batches from larger datasets

// Original tensor: [1000, 784] (1000 samples, 784 features)
val dataset = backend.random(Shape(1000, 784))

// Zero-copy batch views
val batch1 = dataset.sliceView {
    segment { range(0, 32) }    // First 32 samples
    segment { all() }           // All features
}
// Memory: Original tensor + small view metadata (~64 bytes)
// vs Copy approach: Original + 32*784*4 = ~100KB additional memory

2. Feature Extraction

Use Case: Extracting specific feature channels or regions

// Image tensor: [batch, channels, height, width] - NCHW format
val images = backend.random(Shape(64, 3, 224, 224))

// Extract red channel separately (zero-copy, contiguous in NCHW)
val redChannel = images.sliceView {
    segment { all() }          // All batches
    segment { at(0) }          // Red channel only
    segment { all() }          // All height
    segment { all() }          // All width
}

// Extract spatial region (efficient due to row-major layout)
val centerRegion = images.sliceView {
    segment { all() }          // All batches
    segment { all() }          // All channels
    segment { range(112, 140) }  // Center height region
    segment { range(112, 140) }  // Center width region
}

3. Sequence Processing

Use Case: Time series and NLP sliding windows

// Sequence: [batch, sequence_length, features]
val sequences = backend.random(Shape(32, 1000, 512))

// Sliding window views (zero-copy)
val windows = (0 until 900 step 100).map { start ->
    sequences.sliceView {
        segment { all() }              // All batches
        segment { range(start, start + 100) }  // Window of 100 timesteps
        segment { all() }              // All features
    }
}

4. Model Layer Outputs

Use Case: Multi-head attention, feature map processing

// Feature maps: [batch, channels, height, width] - NCHW format
val featureMaps = backend.random(Shape(32, 256, 56, 56))

// Extract specific feature channels (zero-copy, efficient in NCHW)
val firstHalfChannels = featureMaps.sliceView {
    segment { all() }          // All batches
    segment { range(0, 128) }  // First 128 channels
    segment { all() }          // All height
    segment { all() }          // All width
}

// Multi-head attention reshape: [batch, heads, seq_len, head_dim]
val attention = backend.random(Shape(32, 12, 512, 64))
val singleHead = attention.sliceView {
    segment { all() }          // All batches
    segment { at(0) }          // First attention head
    segment { all() }          // All sequence positions
    segment { all() }          // All head dimensions
}

5. Data Augmentation Views

Use Case: Creating different views of the same data

// Original image: [batch, channels, height, width] - NCHW format
val image = backend.random(Shape(1, 3, 224, 224))

// Center crop (zero-copy, optimal for row-major NCHW)
val centerCrop = image.sliceView {
    segment { all() }            // All batches
    segment { all() }            // All channels
    segment { range(16, 208) }   // Crop height
    segment { range(16, 208) }   // Crop width
}

// Region of interest (zero-copy)
val roi = image.sliceView {
    segment { all() }            // All batches
    segment { all() }            // All channels
    segment { range(50, 150) }   // ROI height
    segment { range(75, 175) }   // ROI width
}

// Channel-wise processing (highly efficient in NCHW)
val redChannelOnly = image.sliceView {
    segment { all() }            // All batches
    segment { at(0) }            // Red channel only
    segment { all() }            // All height
    segment { all() }            // All width
}

NCHW Performance Considerations

Memory Access Patterns

The NCHW layout with row-major storage provides specific optimization opportunities:

// For tensor shape [N, C, H, W], element at [n, c, h, w] is at index:
// index = n * (C * H * W) + c * (H * W) + h * W + w

class NCHWIndexMapper(
    private val shape: Shape  // [N, C, H, W]
) : IndexMapper {
    private val strides = intArrayOf(
        shape[1] * shape[2] * shape[3],  // N stride: C*H*W
        shape[2] * shape[3],             // C stride: H*W  
        shape[3],                        // H stride: W
        1                                // W stride: 1
    )
    
    override fun mapToParent(childIndices: IntArray): IntArray {
        // Efficient stride-based mapping for NCHW layout
    }
}

Optimal Slicing Patterns

Most Efficient (Contiguous Memory Access):

Channel Slicing: tensor[:, 0:k, :, :] - Highly efficient
Batch Slicing: tensor[0:k, :, :, :] - Most efficient
Spatial Slicing: tensor[:, :, h1:h2, w1:w2] - Good efficiency

Less Efficient (Strided Access):

Width-only Slicing: tensor[:, :, :, w1:w2] - Requires stride calculations
Non-contiguous Channel: tensor[:, [0,2,4], :, :] - May require materialization

Zero-Copy Efficiency Matrix

Slice Type	Memory Pattern	Efficiency	Zero-Copy
Channel extraction	Contiguous blocks	⭐⭐⭐⭐⭐	Always
Batch slicing	Large contiguous	⭐⭐⭐⭐⭐	Always
Spatial regions	Row-aligned	⭐⭐⭐⭐	Always
Width-only	Strided access	⭐⭐⭐	Usually
Height-only	Large strides	⭐⭐	Usually
Scattered channels	Non-contiguous	⭐	Fallback to copy

Implementation Strategy

Phase 1: Core View Infrastructure

TensorView Interface: Define the contract for tensor views
SlicedTensorView: Implement basic slicing views
IndexMapper: Create efficient index translation
Integration: Connect with existing slicing DSL

Phase 2: Optimized Operations

Contiguous Detection: Identify when slices create contiguous memory views
Stride Calculations: Optimize memory access patterns
View Chaining: Support views of views efficiently
Memory Layout: Consider row-major vs column-major optimizations

Phase 3: Advanced Features

Lazy Materialization: Defer actual tensor creation until needed
View Caching: Cache frequently accessed view metadata
Operation Fusion: Combine multiple view operations
Backend Integration: Optimize for different compute backends

Fallback Mechanisms

When to Create New Tensors

class TensorViewStrategy<T : DType, V> {
    fun shouldCreateView(slices: List<Slice<T, V>>): Boolean {
        return when {
            // Simple contiguous slices -> zero-copy view
            slices.all { it.isContiguous() } -> true
            
            // Complex non-contiguous slices -> may need materialization
            slices.any { it.hasNonTrivialStride() } -> false
            
            // Very small slices -> copy might be faster
            estimateViewSize(slices) < SMALL_TENSOR_THRESHOLD -> false
            
            // Memory pressure -> prefer views
            systemMemoryPressure() > HIGH_PRESSURE_THRESHOLD -> true
            
            else -> true  // Default to zero-copy when possible
        }
    }
}

Materialization Strategies

interface MaterializationStrategy<T : DType, V> {
    fun materialize(view: TensorView<T, V>): Tensor<T, V>
}

class CopyMaterializationStrategy<T : DType, V> : MaterializationStrategy<T, V> {
    override fun materialize(view: TensorView<T, V>): Tensor<T, V> {
        // Create new tensor and copy data
        val newTensor = view.backend.zeros(view.shape)
        copyViewData(view, newTensor)
        return newTensor
    }
}

class LazyMaterializationStrategy<T : DType, V> : MaterializationStrategy<T, V> {
    override fun materialize(view: TensorView<T, V>): Tensor<T, V> {
        // Return a lazy tensor that materializes on first access
        return LazyTensor(view)
    }
}

Performance Considerations

Memory Benefits

Zero-Copy Views: ~64 bytes metadata vs potentially MBs of copied data
Reduced Allocations: Fewer GC pressure points
Cache Efficiency: Better cache locality for related data views

Computational Benefits

Lazy Evaluation: Defer expensive operations until needed
View Composition: Chain operations without intermediate copies
Memory Bandwidth: Reduce memory transfer bottlenecks

Trade-offs

Index Translation Overhead: Each access requires coordinate mapping
Memory Fragmentation: Views may prevent parent tensor garbage collection
Complexity: More complex implementation and debugging

API Evolution Path

Current API Enhancement

// Extend existing DSL to return views instead of slice descriptors
fun <T : DType, V> Tensor<T, V>.sliceView(
    builder: TensorSliceBuilder<T, V>.() -> Unit
): TensorView<T, V>

// Backward compatibility with materialization
fun <T : DType, V> Tensor<T, V>.sliceCopy(
    builder: TensorSliceBuilder<T, V>.() -> Unit
): Tensor<T, V>

Advanced API Features

// View management
fun <T : DType, V> TensorView<T, V>.materialize(): Tensor<T, V>
fun <T : DType, V> TensorView<T, V>.isContiguous(): Boolean
fun <T : DType, V> TensorView<T, V>.optimize(): TensorView<T, V>

// Batch operations on views
fun <T : DType, V> List<TensorView<T, V>>.batchMaterialize(): List<Tensor<T, V>>

Implementation Priority

High Priority (Core Functionality)

✅ Document current status and strategy
SlicedTensorView implementation
Basic IndexMapper for contiguous slices
Integration with existing slicing DSL
Simple fallback to copy-based tensors

Medium Priority (Optimizations)

Stride-based optimizations
Contiguous detection algorithms
View composition support
Performance benchmarking

Low Priority (Advanced Features)

Lazy materialization
View caching strategies
Custom materialization policies
Cross-backend view support

Testing Strategy

Unit Tests

Index mapping correctness
Edge cases (empty slices, out-of-bounds)
Memory usage verification
Performance regression tests

Integration Tests

Compatibility with existing tensor operations
View composition scenarios
ML workload simulations
Memory pressure testing

Benchmarks

Memory usage comparison (view vs copy)
Access pattern performance
Large tensor slicing scenarios
Real ML model integration tests

This document outlines the comprehensive strategy for implementing zero-copy tensor views in SKaiNET. The approach prioritizes memory efficiency while maintaining compatibility with existing APIs and providing clear fallback mechanisms for complex scenarios.

Add slicing and rework Tensors logic to use TensorData to have the sa…

Add initial version of slicing DSL and underlaying implementaion

Add unut tests

Comprehensive slicing infrastructure implemented. Closing.

Add support for slicing #120

Description