Skip to content

Make matrix layout invariants compile-time #29643

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
GopalakrishnanN wants to merge 3 commits into
base: main
Choose a base branch
from
+104 −39
Open

Make matrix layout invariants compile-time #29643

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
bd37fa5
Make matrix layout invariants compile-time
Jul 9, 2026
e766da1
Make Position narrowing conversion explicit
Jul 10, 2026
e173a96
Merge branch 'main' into GopalakrishnanN/compile-time-invariants
GopalakrishnanN Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 46 additions & 39 deletions onnxruntime/core/util/matrix_layout.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,10 @@

#pragma once

#include <algorithm>
#include <cstdint>
#include <type_traits>

#include <gsl/gsl>

#include "core/util/force_inline.h"
Expand All @@ -38,8 +41,10 @@ template <
>
struct Position {
public:
static_assert(Rank_ > 0, "Position rank must be positive.");

/// Number of elements in Position
static int const kRank = Rank_;
static constexpr int kRank = Rank_;

/// Index type used to store elements
using Index = Index_;
Expand All @@ -51,29 +56,29 @@ struct Position {
Index idx[kRank];

public:
ORT_FORCEINLINE explicit Position(Index value = Index(0)) {
ORT_FORCEINLINE constexpr explicit Position(Index value = Index(0)) {
for (int i = 0; i < kRank; ++i) {
idx[i] = value;
}
}

/// Constructs from an array of integers
ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position(Index const (&_idx)[kRank]) {
for (int i = 0; i < kRank; ++i) {
idx[i] = _idx[i];
}
}

template <int R, typename I, typename L>
ORT_FORCEINLINE
Position(Position<R, I, L> other) {
template <typename I, typename L>
ORT_FORCEINLINE constexpr
Position(Position<Rank_, I, L> const& other) {
for (int i = 0; i < kRank; ++i) {
idx[i] = other[i];
idx[i] = gsl::narrow_cast<Index>(other[i]);
}
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position operator+(Position const& b) const {
Position c;
for (int i = 0; i < kRank; ++i) {
Expand All @@ -82,7 +87,7 @@ struct Position {
return c;
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position operator-(Position const& b) const {
Position c;
for (int i = 0; i < kRank; ++i) {
Expand All @@ -91,7 +96,7 @@ struct Position {
return c;
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position operator*(Position const& b) const {
Position c;
for (int i = 0; i < kRank; ++i) {
Expand All @@ -100,7 +105,7 @@ struct Position {
return c;
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position operator/(Position const& b) const {
Position c;
for (int i = 0; i < kRank; ++i) {
Expand All @@ -109,61 +114,61 @@ struct Position {
return c;
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position& operator+=(Position const& b) {
for (int i = 0; i < kRank; ++i) {
idx[i] += b.idx[i];
}
return *this;
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position& operator-=(Position const& b) {
for (int i = 0; i < kRank; ++i) {
idx[i] -= b.idx[i];
}
return *this;
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position& operator*=(Position const& b) {
for (int i = 0; i < kRank; ++i) {
idx[i] *= b.idx[i];
}
return *this;
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position& operator/=(Position const& b) {
for (int i = 0; i < kRank; ++i) {
idx[i] /= b.idx[i];
}
return *this;
}

ORT_FORCEINLINE Index& operator[](int dim) { return idx[dim]; }
ORT_FORCEINLINE constexpr Index& operator[](int dim) { return idx[dim]; }

ORT_FORCEINLINE Index const& operator[](int dim) const { return idx[dim]; }
ORT_FORCEINLINE constexpr Index const& operator[](int dim) const { return idx[dim]; }

ORT_FORCEINLINE bool operator==(Position const& b) const {
ORT_FORCEINLINE constexpr bool operator==(Position const& b) const {
bool equal = true;
for (int i = 0; equal && i < kRank; ++i) {
equal = (idx[i] == b.idx[i]);
}
return equal;
}

ORT_FORCEINLINE bool operator!=(Position const& b) const { return !(*this == b); }
ORT_FORCEINLINE constexpr bool operator!=(Position const& b) const { return !(*this == b); }

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Position& clamp(Position const& max, Position const& min = Position()) {
for (int i = 0; i < kRank; ++i) {
idx[i] = std::max(std::min(idx[i], max.idx[i]), min.idx[i]);
}
return *this;
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Index sum() const {
Index sum_(idx[0]);
for (int i = 1; i < kRank; ++i) {
Expand All @@ -172,7 +177,7 @@ struct Position {
return sum_;
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
LongIndex product() const {
LongIndex product_(idx[0]);
for (int i = 1; i < kRank; ++i) {
Expand All @@ -183,15 +188,15 @@ struct Position {
};

template <typename T, typename L = int64_t>
Position<2, T, L> make_Position(T _0, T _1) {
constexpr Position<2, T, L> make_Position(T _0, T _1) {
T values[2] = {_0, _1};
return Position<2, T, L>(values);
}

template <typename T, typename L = int64_t>
Position<3, T, L> make_Position(T _0, T _1, T _2) {
constexpr Position<3, T, L> make_Position(T _0, T _1, T _2) {
T values[3] = {_0, _1, _2};
return Position<2, T, L>(values);
return Position<3, T, L>(values);
}

/// Describes the size of a matrix tile
Expand All @@ -200,11 +205,13 @@ template <
int Column_ ///< columns of a matrix
>
struct MatrixShape {
static int const kRow = Row_; ///< rows of a matrix
static int const kColumn = Column_; ///< columns of a matrix
static int const kCount = Row_ * Column_; ///< total number of elements in a matrix
static_assert(Row_ >= 0 && Column_ >= 0, "Matrix dimensions must not be negative.");

ORT_FORCEINLINE static Position<2> toCoord() {
static constexpr int kRow = Row_; ///< rows of a matrix
static constexpr int kColumn = Column_; ///< columns of a matrix
static constexpr int kCount = Row_ * Column_; ///< total number of elements in a matrix

ORT_FORCEINLINE static constexpr Position<2> toCoord() {
return make_Position(kRow, kColumn);
}
};
Expand All @@ -228,27 +235,27 @@ class RowMajorLayout {
Index stride_;

public:
ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
RowMajorLayout(Index ldm = 0) : stride_(ldm) {}

ORT_FORCEINLINE static RowMajorLayout packed(MatCoord const& extent) {
ORT_FORCEINLINE static constexpr RowMajorLayout packed(MatCoord const& extent) {
return RowMajorLayout(extent[1]);
}

/// Returns the offset of a coordinate in linear memory.
/// Assumes coordinate has convention (row, column)
ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
LongIndex operator()(MatCoord const& coord) const {
return LongIndex(coord[0]) * stride_ + coord[1];
}

/// Inverse of layout function, mapping linear offset to logical coordinate
ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
MatCoord inverse(LongIndex offset) const {
return make_Position(Index(offset / stride_), Index(offset % stride_));
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Index stride() const {
return stride_;
}
Expand All @@ -269,27 +276,27 @@ class ColumnMajorLayout {
Index stride_;

public:
ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
ColumnMajorLayout(Index ldm = 0) : stride_(ldm) {}

ORT_FORCEINLINE static ColumnMajorLayout packed(MatCoord const& extent) {
ORT_FORCEINLINE static constexpr ColumnMajorLayout packed(MatCoord const& extent) {
return ColumnMajorLayout(extent[0]);
}

/// Returns the offset of a coordinate in linear memory.
/// Assumes coordinate has convention (row, column)
ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
LongIndex operator()(MatCoord const& coord) const {
return LongIndex(coord[1]) * LongIndex(stride_) + coord[0];
}

/// Inverse of layout function, mapping linear offset to logical coordinate
ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
MatCoord inverse(LongIndex offset) const {
return make_Position(Index(offset % stride_), Index(offset / stride_));
}

ORT_FORCEINLINE
ORT_FORCEINLINE constexpr
Index stride() const {
return stride_;
}
Expand Down
58 changes: 58 additions & 0 deletions onnxruntime/test/common/matrix_layout_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include <type_traits>

#include "core/util/matrix_layout.h"
#include "gtest/gtest.h"

namespace onnxruntime {
namespace {

constexpr bool PositionOperationsAreConstexpr() {
auto position = make_Position(2, 3);
position += make_Position(4, 5);
position *= make_Position(2, 3);
position -= make_Position(1, 2);
position /= make_Position(11, 11);
position.clamp(make_Position(2, 3));

return position == make_Position(1, 2) && position.sum() == 3 && position.product() == 2;
}

static_assert(PositionOperationsAreConstexpr());

constexpr auto kPosition3D = make_Position(2, 3, 4);
static_assert(decltype(kPosition3D)::kRank == 3);
static_assert(kPosition3D[0] == 2 && kPosition3D[1] == 3 && kPosition3D[2] == 4);
static_assert(!std::is_constructible_v<Position<3>, Position<2>>);

constexpr Position<2, int64_t> kWidePosition = make_Position<int64_t>(5, 7);
constexpr Position<2> kConvertedPosition = kWidePosition;
static_assert(kConvertedPosition == make_Position(5, 7));

using TestMatrixShape = MatrixShape<3, 4>;
static_assert(TestMatrixShape::kCount == 12);
static_assert(TestMatrixShape::toCoord() == make_Position(3, 4));

constexpr auto kRowMajorLayout = RowMajorLayout::packed(TestMatrixShape::toCoord());
static_assert(kRowMajorLayout.stride() == 4);
static_assert(kRowMajorLayout(make_Position(2, 1)) == 9);
static_assert(kRowMajorLayout.inverse(9) == make_Position(2, 1));

constexpr auto kColumnMajorLayout = ColumnMajorLayout::packed(TestMatrixShape::toCoord());
static_assert(kColumnMajorLayout.stride() == 3);
static_assert(kColumnMajorLayout(make_Position(2, 1)) == 5);
static_assert(kColumnMajorLayout.inverse(5) == make_Position(2, 1));

TEST(MatrixLayoutTest, MapsRuntimeCoordinates) {
const int row = 2;
const int column = 1;
const auto coordinate = make_Position(row, column);

EXPECT_EQ(kRowMajorLayout(coordinate), 9);
EXPECT_EQ(kColumnMajorLayout(coordinate), 5);
}

} // namespace
} // namespace onnxruntime
Loading