Skip to content

[Matrix][SYCL] Add support for bf16's wi_element #5397

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 3 commits into from
Jan 29, 2022
Merged

[Matrix][SYCL] Add support for bf16's wi_element #5397

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
92b20f0
[Matrix][SYCL] add support for bf16's wi_element
yubingex007-a11y Jan 26, 2022
22542aa
Add some comments
yubingex007-a11y Jan 27, 2022
a5f902a
Modify some comments
yubingex007-a11y Jan 28, 2022
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
280 changes: 280 additions & 0 deletions sycl/include/sycl/ext/oneapi/matrix/matrix-jit.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -451,6 +451,286 @@ class wi_element {
}
};

// Note that similarly to the other matrix functions, uint16_t is used here to
// represent bf16 type. Since the AMX and DPAS implementations don't support
// uint16_t, this interpretation is possible. This design choice was made before
// the introduction of SYCL experimental bfloat16 type. Our plan is to move
// towards using the SYCL bfloat16. But since it is still experimental, we will
// probably keep both uint16 interpretation and SYCL bfloat16.
template <size_t NumRows, size_t NumCols, matrix_layout Layout, typename Group>
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add a comment that uint16_t here is representing bf16 type as we have been doing for the previous matrix functions.
Since the AMX and DPAS implementations don't support uint16_t, this should raise no problem.
Add a comment that the plan is to move towards SYCL bfloat16 as it makes it to the specification.

class wi_element<uint16_t, NumRows, NumCols, Layout, Group> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should specialize this class on bfloat16 type.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alexey, we are adding the support of bfloat16 type as the next step in a separate PR.
This PR is to extend the support of the matrix bf16 type represented by unsigned short.
so we are not changing the way we treat bf16 right now.
Instead, we plan on keeping both: unsigned short as bf16, then SYCL bfloat16 type.
We will keep both until the SYCL bfloat16 makes it to the specification (it is experimental right now)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@AlexeySotkin @dkhaldi For now we still consider uint16_t as bf16 in matrix and in the future we will support sycl's bf16.

joint_matrix<uint16_t, NumRows, NumCols, Layout, Group> &M;
std::size_t idx;

public:
wi_element(joint_matrix<uint16_t, NumRows, NumCols, Layout, Group> &Mat,
std::size_t i)
: M(Mat), idx(i) {}
operator uint16_t() {
#ifdef __SYCL_DEVICE_ONLY__
return __spirv_VectorExtractDynamic(M.spvm, idx);
#else
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

explicit operator bool() {
#ifdef __SYCL_DEVICE_ONLY__
return __spirv_VectorExtractDynamic(M.spvm, idx) !=
static_cast<uint16_t>(0);
#else
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

wi_element &operator=(const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
M.spvm = __spirv_VectorInsertDynamic(M.spvm, rhs, idx);
return *this;
#else
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

wi_element &
operator=(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
M.spvm = __spirv_VectorInsertDynamic(
M.spvm, __spirv_VectorExtractDynamic(rhs.M.spvm, rhs.idx), idx);
return *this;
#else
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

// We use here the following functions for conversion (bf16=>fp32 and
// fp32=>bf16). This is a workaround until we are able to use
// __spirv_ConvertFToBF16INTEL and __spirv_ConvertBF16ToFINTEL once these are
// supported in the CPU backend
static float make_fp32(uint16_t x) {
unsigned int y = x;
y = y << 16;
float *res = reinterpret_cast<float *>(&y);
return *res;
}

static uint16_t make_bf16(float x) {
int *res = reinterpret_cast<int *>(&x);
*res = *res >> 16;
return (uint16_t)*res;
}
Comment on lines +516 to +527
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why we can't use __spirv_ConvertBF16ToFINTEL and __spirv_ConvertFToBF16INTEL here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We use this workaround until we support __spirv_ConvertBF16ToFINTEL and __spirv_ConvertFToBF16INTEL in cpu side.

Copy link
Contributor

@Nuullll Nuullll Mar 8, 2022 •
edited by bader
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI. I'm working on supporting SYCL_INTEL_bf16_conversion capability on CPU.


friend uint16_t
operator+(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_bf16(
make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) +
make_fp32(rhs));
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

wi_element &operator+=(const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
M.spvm = __spirv_VectorInsertDynamic(
M.spvm,
make_bf16(make_fp32(__spirv_VectorExtractDynamic(M.spvm, idx)) +
make_fp32(rhs)),
idx);
return *this;
#else
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

friend uint16_t
operator-(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_bf16(
make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) -
make_fp32(rhs));
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

wi_element &operator-=(const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
M.spvm = __spirv_VectorInsertDynamic(
M.spvm,
make_bf16(make_fp32(__spirv_VectorExtractDynamic(M.spvm, idx)) -
make_fp32(rhs)),
idx);
return *this;
#else
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

friend uint16_t
operator*(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_bf16(
make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) *
make_fp32(rhs));
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

wi_element &operator*=(const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
M.spvm = __spirv_VectorInsertDynamic(
M.spvm,
make_bf16(make_fp32(__spirv_VectorExtractDynamic(M.spvm, idx)) *
make_fp32(rhs)),
idx);
return *this;
#else
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

friend uint16_t
operator/(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_bf16(
make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) /
make_fp32(rhs));
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

wi_element &operator/=(const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
M.spvm = __spirv_VectorInsertDynamic(
M.spvm,
make_bf16(make_fp32(__spirv_VectorExtractDynamic(M.spvm, idx)) /
make_fp32(rhs)),
idx);
return *this;
#else
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

friend bool
operator<(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) <
make_fp32(rhs);
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

friend bool
operator<=(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) <=
make_fp32(rhs);
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

friend bool
operator>(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) >
make_fp32(rhs);
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

friend bool
operator>=(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) >=
make_fp32(rhs);
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

friend bool
operator==(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) ==
make_fp32(rhs);
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}

friend bool
operator!=(const wi_element<uint16_t, NumRows, NumCols, Layout, Group> &lhs,
const uint16_t &rhs) {
#ifdef __SYCL_DEVICE_ONLY__
return make_fp32(__spirv_VectorExtractDynamic(lhs.M.spvm, lhs.idx)) !=
make_fp32(rhs);
#else
(void)lhs;
(void)rhs;
throw runtime_error("joint matrix is not supported on host device.",
PI_INVALID_DEVICE);
#endif // __SYCL_DEVICE_ONLY__
}
};

template <typename T, size_t NumRows, size_t NumCols, matrix_layout Layout,
typename Group>
class wi_slice {
Expand Down