slices.hpp

Go to the documentation of this file.

// Copyright 2025 Eric Hermosis
//
// This file is part of the Tannic Tensor Library.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// 
 
#ifndef SLICES_HPP
#define SLICES_HPP
 
#include <tuple>
#include <utility>
#include <cstddef>
#include <vector>
 
#include "types.hpp" 
#include "traits.hpp"
#include "shape.hpp"
#include "strides.hpp"
#include "indexing.hpp"  
 
namespace tannic {
class Tensor;
}
 
namespace tannic::expression {   
 
template <Expression Source, class... Indexes>
class Slice {  
public:    
 
    constexpr Slice(typename Trait<Source>::Reference source, std::tuple<Indexes...> indexes)
    :   dtype_(source.dtype())  
    ,   source_(source)
    ,   indexes_(indexes)
    {   
        if (source.rank() == 0) { 
            shape_ = Shape{};
            strides_ = Strides{};
            offset_ = 0;  
        }
 
        else {
            std::array<Shape::size_type, Shape::limit> shape{};
            std::array<Strides::size_type, Strides::limit> strides{};
            Shape::rank_type dimension = 0;
            Shape::rank_type rank = 0;   
            offset_ = 0;
            auto process = [&](const auto& argument) {
                using Argument = std::decay_t<decltype(argument)>;
                if constexpr (std::is_same_v<Argument, indexing::Range>) { 
                    auto range = normalize(argument, source.shape()[dimension]);
                    auto size = range.stop - range.start; 
                    shape[rank] = size;
                    strides[rank] = source.strides()[dimension]; 
                    offset_ += range.start * source.strides()[dimension] * dsizeof(dtype_); 
                    rank++; dimension++;
                } 
                
                else if constexpr (std::is_integral_v<Argument>) { 
                    auto index = indexing::normalize(argument, source.shape()[dimension]);
                    offset_ += index * source.strides()[dimension] * dsizeof(dtype_);
                    dimension++;
                } 
                
                else {         
                    throw Exception("Unknown index type"); 
                } 
            };
 
            std::apply([&](const auto&... arguments) {
                (process(arguments), ...);
            }, indexes);
 
            while (dimension < source.rank()) {
                shape[rank] = source.shape()[dimension];
                strides[rank] = source.strides()[dimension];
                rank++; dimension++;
            } 
 
            shape_ = Shape(shape.begin(), shape.begin() + rank);
            strides_ = Strides(strides.begin(), strides.begin() + rank); 
        } 
    }    
 
    template<Integral Index>
    constexpr auto operator[](Index index) const { 
        return Slice<Source, Indexes..., Index>(source_, std::tuple_cat(indexes_, std::make_tuple(index)));
    } 
 
    constexpr auto operator[](indexing::Range range) const {    
        return Slice<Source, Indexes..., indexing::Range>(source_, std::tuple_cat(indexes_, std::make_tuple(range)));
    }  
 
    template<typename T>
    void operator=(T value);
 
 
 
    template<typename T>
    bool operator==(T value) const;  
 
    constexpr auto dtype() const {
        return dtype_;
    } 
 
    constexpr auto rank() const {
        return shape_.rank();
    } 
 
    constexpr Shape const& shape() const {
        return shape_;
    }
 
 
    constexpr Strides const& strides() const {
        return strides_;
    }  
 
    std::ptrdiff_t offset() const {
        return offset_ + source_.offset();
    }
 
    std::byte* bytes() {
        return source_.bytes() + offset_;
    }
 
    std::byte const* bytes() const {
        return source_.bytes() + offset_;
    } 
 
    Tensor forward() const;
 
    void assign(std::byte const* value, std::ptrdiff_t offset); 
    void assign(bool const*, std::ptrdiff_t); 
    bool compare(std::byte const* value, std::ptrdiff_t offset) const;
 
private: 
    type dtype_;
    Shape shape_;
    Strides strides_;
    std::ptrdiff_t offset_;
    typename Trait<Source>::Reference source_;              
    std::tuple<Indexes...> indexes_;    
};   
 
template<typename T>
inline std::byte const* tobytes(T const& reference) { 
    return reinterpret_cast<std::byte const*>(&reference);
}  
 
template <Expression Source, class... Indexes>
template <typename T>
void Slice<Source, Indexes...>::operator=(T value) {    
    auto copy = [this](std::byte const* value, std::ptrdiff_t offset) {
        if(rank() == 0) { 
            assign(value, offset);
            return;
        } 
        std::vector<std::size_t> indexes(rank(), 0);
        bool done = false;  
        
        while (!done) {
            std::size_t position = offset;
            for (auto dimension = 0; dimension < rank(); ++dimension) {
                position += indexes[dimension] * strides_[dimension] * dsizeof(dtype_);
            } 
  
            assign(value, position); 
            done = true;
            for (int dimension = rank() - 1; dimension >= 0; --dimension) {
                if (++indexes[dimension] < shape_[dimension]) {
                    done = false;
                    break;
                }
                indexes[dimension] = 0;
            }
        }
    };
 
    switch (dtype_) {
        case int8:  {   int8_t casted = value; copy(tobytes(casted), offset()); break; }
        case int16: {  int16_t casted = value; copy(tobytes(casted), offset()); break; }
        case int32: {  int32_t casted = value; copy(tobytes(casted), offset()); break; }
        case int64: {  int64_t casted = value; copy(tobytes(casted), offset()); break; }
        case float32: {  float casted = value; copy(tobytes(casted), offset()); break; }
        case float64: { double casted = value; copy(tobytes(casted), offset()); break; } 
        default: throw Exception("Unsupported dtype for assignment");
    } 
}
 
template <Expression Source, class... Indexes>
template <typename T>
bool Slice<Source, Indexes...>::operator==(T value) const {    
    if (rank() != 0)
        throw Exception("Cannot compare an scalar to a non scalar slice");
 
    switch (dtype_) { 
        case int8:  {    int8_t casted = value; return compare(tobytes(casted), offset()); }
        case int16: {   int16_t casted = value; return compare(tobytes(casted), offset()); }
        case int32: {   int32_t casted = value; return compare(tobytes(casted), offset()); }
        case int64: {   int64_t casted = value; return compare(tobytes(casted), offset()); }
        case float32:  {  float casted = value; return compare(tobytes(casted), offset()); }
        case float64:  { double casted = value; return compare(tobytes(casted), offset()); }  
        default: throw Exception("Unsupported dtype for comparison");
    }
}   
 
} // namespace tannic::expression
 
#endif // SLICES_HPP