Skip to content

Add TwoOffsets_CombinedIndexer, and UnpackedStridedIndexer #1162

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

Merged
merged 1 commit into from
Apr 8, 2023
Merged

Add TwoOffsets_CombinedIndexer, and UnpackedStridedIndexer #1162

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
133 changes: 114 additions & 19 deletions dpctl/tensor/libtensor/include/utils/offset_utils.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,7 @@ struct NoOpIndexer
}
};

/* @brief Indexer with shape and strides arrays of same size are packed */
struct StridedIndexer
{
StridedIndexer(int _nd,
Expand All @@ -143,24 +144,76 @@ struct StridedIndexer
{
}

size_t operator()(py::ssize_t gid) const
{
return compute_offset(gid);
}

size_t operator()(size_t gid) const
{
return compute_offset(static_cast<py::ssize_t>(gid));
}

private:
int nd;
py::ssize_t starting_offset;
py::ssize_t const *shape_strides;

size_t compute_offset(py::ssize_t gid) const
{
using dpctl::tensor::strides::CIndexer_vector;

CIndexer_vector _ind(nd);
py::ssize_t relative_offset(0);
_ind.get_displacement<const py::ssize_t *, const py::ssize_t *>(
static_cast<py::ssize_t>(gid),
gid,
shape_strides, // shape ptr
shape_strides + nd, // strides ptr
relative_offset);
return starting_offset + relative_offset;
}
};

/* @brief Indexer with shape, strides provided separately */
struct UnpackedStridedIndexer
{
UnpackedStridedIndexer(int _nd,
py::ssize_t _offset,
py::ssize_t const *_shape,
py::ssize_t const *_strides)
: nd(_nd), starting_offset(_offset), shape(_shape), strides(_strides)
{
}

size_t operator()(py::ssize_t gid) const
{
return compute_offset(gid);
}

size_t operator()(size_t gid) const
{
return compute_offset(static_cast<py::ssize_t>(gid));
}

private:
int nd;
py::ssize_t starting_offset;
py::ssize_t const *shape_strides;
py::ssize_t const *shape;
py::ssize_t const *strides;

size_t compute_offset(py::ssize_t gid) const
{
using dpctl::tensor::strides::CIndexer_vector;

CIndexer_vector _ind(nd);
py::ssize_t relative_offset(0);
_ind.get_displacement<const py::ssize_t *, const py::ssize_t *>(
gid,
shape, // shape ptr
strides, // strides ptr
relative_offset);
return starting_offset + relative_offset;
}
};

struct Strided1DIndexer
Expand Down Expand Up @@ -206,7 +259,8 @@ struct Strided1DCyclicIndexer
template <typename displacementT> struct TwoOffsets
{
TwoOffsets() : first_offset(0), second_offset(0) {}
TwoOffsets(displacementT first_offset_, displacementT second_offset_)
TwoOffsets(const displacementT &first_offset_,
const displacementT &second_offset_)
: first_offset(first_offset_), second_offset(second_offset_)
{
}
Expand Down Expand Up @@ -238,6 +292,22 @@ struct TwoOffsets_StridedIndexer
}

TwoOffsets<py::ssize_t> operator()(py::ssize_t gid) const
{
return compute_offsets(gid);
}

TwoOffsets<py::ssize_t> operator()(size_t gid) const
{
return compute_offsets(static_cast<py::ssize_t>(gid));
}

private:
int nd;
py::ssize_t starting_first_offset;
py::ssize_t starting_second_offset;
py::ssize_t const *shape_strides;

TwoOffsets<py::ssize_t> compute_offsets(py::ssize_t gid) const
{
using dpctl::tensor::strides::CIndexer_vector;

Expand All @@ -254,12 +324,6 @@ struct TwoOffsets_StridedIndexer
starting_first_offset + relative_first_offset,
starting_second_offset + relative_second_offset);
}

private:
int nd;
py::ssize_t starting_first_offset;
py::ssize_t starting_second_offset;
py::ssize_t const *shape_strides;
};

struct TwoZeroOffsets_Indexer
Expand All @@ -272,12 +336,33 @@ struct TwoZeroOffsets_Indexer
}
};

template <typename FirstIndexerT, typename SecondIndexerT>
struct TwoOffsets_CombinedIndexer
{
private:
FirstIndexerT first_indexer_;
SecondIndexerT second_indexer_;

public:
TwoOffsets_CombinedIndexer(const FirstIndexerT &first_indexer,
const SecondIndexerT &second_indexer)
: first_indexer_(first_indexer), second_indexer_(second_indexer)
{
}

TwoOffsets<py::ssize_t> operator()(py::ssize_t gid) const
{
return TwoOffsets<py::ssize_t>(first_indexer_(gid),
second_indexer_(gid));
}
};

template <typename displacementT> struct ThreeOffsets
{
ThreeOffsets() : first_offset(0), second_offset(0), third_offset(0) {}
ThreeOffsets(displacementT first_offset_,
displacementT second_offset_,
displacementT third_offset_)
ThreeOffsets(const displacementT &first_offset_,
const displacementT &second_offset_,
const displacementT &third_offset_)
: first_offset(first_offset_), second_offset(second_offset_),
third_offset(third_offset_)
{
Expand Down Expand Up @@ -317,6 +402,23 @@ struct ThreeOffsets_StridedIndexer
}

ThreeOffsets<py::ssize_t> operator()(py::ssize_t gid) const
{
return compute_offsets(gid);
}

ThreeOffsets<py::ssize_t> operator()(size_t gid) const
{
return compute_offsets(static_cast<py::ssize_t>(gid));
}

private:
int nd;
py::ssize_t starting_first_offset;
py::ssize_t starting_second_offset;
py::ssize_t starting_third_offset;
py::ssize_t const *shape_strides;

ThreeOffsets<py::ssize_t> compute_offsets(py::ssize_t gid) const
{
using dpctl::tensor::strides::CIndexer_vector;

Expand All @@ -337,13 +439,6 @@ struct ThreeOffsets_StridedIndexer
starting_second_offset + relative_second_offset,
starting_third_offset + relative_third_offset);
}

private:
int nd;
py::ssize_t starting_first_offset;
py::ssize_t starting_second_offset;
py::ssize_t starting_third_offset;
py::ssize_t const *shape_strides;
};

struct ThreeZeroOffsets_Indexer
Expand Down