matrix.hh

Go to the documentation of this file.

/*
 * Copyright (c) 2022 Arm Limited
 * All rights reserved
 *
 * The license below extends only to copyright in the software and shall
 * not be construed as granting a license to any other intellectual
 * property including but not limited to intellectual property relating
 * to a hardware implementation of the functionality of the software
 * licensed hereunder.  You may use the software subject to the license
 * terms below provided that you ensure that this notice is replicated
 * unmodified and in its entirety in all distributions of the software,
 * modified or unmodified, in source code or in binary form.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met: redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer;
 * redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution;
 * neither the name of the copyright holders nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
 
#ifndef __ARCH_ARM_MATRIX_HH__
#define __ARCH_ARM_MATRIX_HH__
 
#include <array>
#include <cassert>
#include <cstring>
#include <iostream>
#include <type_traits>
 
#include "base/cprintf.hh"
#include "base/logging.hh"
#include "base/types.hh"
#include "sim/serialize_handlers.hh"
 
namespace gem5
{
 
constexpr unsigned MaxMatRegRowLenInBytes = 256;
constexpr unsigned MaxMatRegRows = 256;
 
// Forward declarations
template <size_t X, size_t Y>
class MatStore;
template <typename ElemType, typename Container>
class Tile;
 
template <size_t X, size_t Y>
struct ParseParam<MatStore<X, Y>>;
 
template <typename ElemType, typename Container, bool FromTile>
class HorizontalSlice
{
    template <size_t, size_t> friend class MatStore;
    template <typename, typename> friend class Tile;
 
  private:
    Container * container;
    size_t index;
    size_t xElems;
    size_t yElems;
    size_t startElts;
    size_t strideElts;
 
  private:
    HorizontalSlice(Container& cnt, size_t _startBytes, size_t _strideBytes,
                    size_t idx)
      : container(&cnt), index(idx),
        xElems(container->xSize() / sizeof(ElemType)),
        yElems(container->ySize() / (FromTile ? sizeof(ElemType): 1)),
        startElts(_startBytes / sizeof(ElemType)),
        strideElts(_strideBytes / sizeof(ElemType))
    {
        gem5_assert(xElems > 0, "The number of xElems cannot be 0");
        gem5_assert(yElems > 0, "The number of yElems cannot be 0");
 
        // Make sure that we have a whole multiple of an element size
        assert (_startBytes % sizeof(ElemType) == 0);
        assert (_strideBytes % sizeof(ElemType) == 0);
 
        if constexpr (!FromTile) {
            // If we are not operating on a tile, the stride must be the
            // same as the row length, X.
            assert(_strideBytes == container->xSize());
        } else {
            // If we are operating on a tile, then the stride must be
            // sizeof(ElemSize) greater than X.
            assert(_strideBytes / container->xSize() == sizeof(ElemType));
        }
    };
 
  public:
    ElemType&
    operator[](size_t elem_idx)
    {
        assert(elem_idx < xElems);
        size_t linear_index = startElts + index * strideElts + elem_idx;
        return container->template rawPtr<ElemType>()[linear_index];
    };
 
    void
    zero()
    {
        for (int i = 0; i < xElems; ++i) {
            (*this)[i] = (ElemType)0;
        }
    };
};
 
template <typename ElemType, typename Container, bool FromTile>
class VerticalSlice
{
    template <size_t, size_t> friend class MatStore;
    template <typename, typename> friend class Tile;
 
  private:
    Container * container;
    size_t index;
    size_t xElems;
    size_t yElems;
    size_t startElts;
    size_t strideElts;
 
  private:
    VerticalSlice(Container& cnt, size_t _startBytes, size_t _strideBytes, size_t idx)
      : container(&cnt), index(idx),
        xElems(container->xSize() / sizeof(ElemType)),
        yElems(container->ySize() / (FromTile ? sizeof(ElemType): 1)),
        startElts(_startBytes / sizeof(ElemType)),
        strideElts(_strideBytes / sizeof(ElemType))
    {
        gem5_assert(xElems > 0, "The number of xElems cannot be 0");
        gem5_assert(yElems > 0, "The number of yElems cannot be 0");
 
        // Make sure that we have a whole multiple of an element size
        assert (_startBytes % sizeof(ElemType) == 0);
        assert (_strideBytes % sizeof(ElemType) == 0);
 
        if constexpr (!FromTile) {
            // If we are not operating on a tile, the stride must be the
            // same as the row length, X.
            assert(_strideBytes == container->xSize());
        } else {
            // If we are operating on a tile, then the stride must be
            // sizeof(ElemSize) greater than X.
            assert(_strideBytes / container->xSize() == sizeof(ElemType));
        }
    };
 
  public:
    ElemType&
    operator[](size_t elem_idx)
    {
        assert(elem_idx < yElems);
        size_t linear_index = startElts + elem_idx * strideElts + index;
        return container->template rawPtr<ElemType>()[linear_index];
    };
 
    void
    zero()
    {
        for (int i = 0; i < yElems; ++i) {
            (*this)[i] = (ElemType)0;
        }
    };
};
 
template <typename ElemType, typename Container>
class Tile
{
    template <size_t, size_t> friend class MatStore;
 
    // We "calculate" the number of possible tiles based on the element size
    static constexpr size_t NUM_TILES = sizeof(ElemType);
 
  private:
    Container * container;
    size_t index;
    size_t startBytes;
    size_t strideBytes;
 
  private:
    Tile(Container& cnt, size_t idx)
      : container(&cnt), index(idx)
    {
        assert(index < NUM_TILES);
        startBytes = container->xSize() * index;
        strideBytes = NUM_TILES * container->xSize();
    };
 
  public:
    auto
    operator[](size_t idx)
    {
        assert(idx < (container->ySize() / NUM_TILES));
        return asHSlice(idx);
    };
 
    Container*
    getContainer()
    {
        return container;
    };
 
    auto
    asHSlice(size_t row_idx)
    {
        assert(row_idx < container->ySize() / NUM_TILES);
        return HorizontalSlice<ElemType, Container, true>(*container,
                                                          startBytes,
                                                          strideBytes,
                                                          row_idx);
    };
 
    auto
    asVSlice(size_t col_idx)
    {
        assert(col_idx < container->xSize());
        return VerticalSlice<ElemType, Container, true>(*container, startBytes,
                                                        strideBytes, col_idx);
    };
 
    void
    zero()
    {
        for (int i = 0; i < container->ySize() / NUM_TILES; ++i) {
            // We zero the tile by rows. We need to do it this way due
            // to the interleaving.
            auto row = this->asHSlice(i);
            row.zero();
        }
    };
};
 
// Base container class for a matrix. Allows for non-square matricies.
template <size_t X, size_t Y>
class MatStore
{
    static_assert(X > 0, "X size cannot be 0");
    static_assert(Y > 0, "Y size cannot be 0");
 
    static constexpr size_t LINEAR_SIZE = X * Y;
 
    template <typename, typename, bool> friend class HorizontalSlice;
    template <typename, typename, bool> friend class VerticalSlice;
 
  public:
    static constexpr inline size_t xSize() { return X; };
    static constexpr inline size_t ySize() { return Y; };
    static constexpr inline size_t linearSize() { return LINEAR_SIZE; };
 
    using Container = std::array<uint8_t, LINEAR_SIZE>;
    using MyClass = MatStore<X, Y>;
  private:
    // We need to be able to handle 128-bit types; align accordingly
    alignas(16) Container container;
 
  public:
    MatStore() {};
 
    MatStore(const MatStore&) = default;
 
    void
    zero()
    {
        memset(container.data(), 0 , LINEAR_SIZE);
    }
 
    MyClass&
    operator=(const MyClass& that)
    {
        if (&that == this)
            return *this;
        memcpy(container.data(), that.container.data(), LINEAR_SIZE);
        return *this;
    }
    template<size_t X2, size_t Y2>
    inline bool
    operator==(const MatStore<X2, Y2>& that) const
    {
        return X == X2 && Y == Y2 &&
               !memcmp(container.data(), that.container.data(), LINEAR_SIZE);
    }
 
    template<size_t X2, size_t Y2>
    bool
    operator!=(const MatStore<X2, Y2>& that) const
    {
        return !operator==(that);
    }
 
  private:
    template <typename ElemType>
    const ElemType* rawPtr() const
    {
        return reinterpret_cast<const ElemType*>(container.data());
    }
 
    template <typename ElemType>
    ElemType* rawPtr() { return reinterpret_cast<ElemType*>(container.data()); }
 
  public:
    template <typename ElemType>
    auto
    asTile(size_t index)
    {
        return Tile<ElemType, MyClass>(*this, index);
    }
 
    template <typename ElemType>
    auto
    asHSlice(size_t row_idx)
    {
        return HorizontalSlice<ElemType, MyClass, false>(*this, 0, X, row_idx);
    }
 
    template <typename ElemType>
    auto
    asVSlice(size_t col_idx)
    {
        return VerticalSlice<ElemType, MyClass, false>(*this, 0, X, col_idx);
    }
 
    friend std::ostream&
    operator<<(std::ostream& os, const MatStore<X, Y>& v)
    {
        // When printing for human consumption, break into 4 byte chunks.
        ccprintf(os, "[");
        size_t count = 0;
        for (auto& b: v.container) {
            if (count && (count % 4) == 0)
                os << "_";
            ccprintf(os, "%02x", b);
            count++;
        }
        ccprintf(os, "]");
        return os;
    }
 
    friend ParseParam<MatStore<X, Y>>;
    friend ShowParam<MatStore<X, Y>>;
 
};
 
template <size_t X, size_t Y>
struct ParseParam<MatStore<X, Y>>
{
    static bool
    parse(const std::string &str, MatStore<X, Y> &value)
    {
        fatal_if(str.size() > 2 * X * Y,
                 "Matrix register value overflow at unserialize");
        fatal_if(str.size() < 2 * X * Y,
                 "Matrix register value underflow at unserialize");
 
        for (int i = 0; i < X * Y; i++) {
            uint8_t b = 0;
            if (2 * i < str.size())
                b = stoul(str.substr(i * 2, 2), nullptr, 16);
            value.template rawPtr<uint8_t>()[i] = b;
        }
        return true;
    }
};
 
template <size_t X, size_t Y>
struct ShowParam<MatStore<X, Y>>
{
    static void
    show(std::ostream &os, const MatStore<X, Y> &value)
    {
        for (auto& b: value.container)
            ccprintf(os, "%02x", b);
    }
};
struct DummyMatRegContainer
{
    RegVal filler = 0;
    bool operator == (const DummyMatRegContainer &d) const { return true; }
    bool operator != (const DummyMatRegContainer &d) const { return true; }
};
template <>
struct ParseParam<DummyMatRegContainer>
{
    static bool
    parse(const std::string &s, DummyMatRegContainer &value)
    {
        return false;
    }
};
static_assert(sizeof(DummyMatRegContainer) == sizeof(RegVal));
static inline std::ostream &
operator<<(std::ostream &os, const DummyMatRegContainer &d)
{
    return os;
}
} // namespace gem5
 
#endif // __ARCH_ARM_MATRIX_HH__