BRL-CAD/oidn
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SYCL: use 2D block load/store instructions

Browse Source
This commit is contained in:
Attila Afra
2026-04-16 01:00:28 +03:00
parent 5763847cca
commit 588eef3b75
3 changed files with 174 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
devices/sycl/sycl_concat_conv_xe.h
+8 -19
View File
@@ -15,7 +15,6 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
using Base::KH, Base::KW, Base::PH, Base::PW;
using Base::blockC, Base::blockOH, Base::blockOW;
using typename Base::AccumRows, typename Base::InRows, typename Base::OutRows;
using Base::loadRow, Base::loadUpsampleRow, Base::mmaBlockKW, Base::finalizeBlock, Base::storeRow;
TensorAccessor3D<SrcDstT, srcDstLayout> src0, src1;
TensorAccessor4D<WeightT, weightLayout> weight;
@@ -43,7 +42,7 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
for (int boh = 0; boh < blockOH - 1; ++boh)
{
if (boh == 0 || (boh + PH) % 2 == 0) // only need to load every other row due to 2x upsampling
loadUpsampleRow(inRows[boh], src0, ic, ih + boh, iw);
Base::loadUpsampleRow(inRows[boh], src0, ic, ih + boh, iw);
else
inRows[boh] = inRows[boh - 1]; // duplicate previous row
}
@@ -54,12 +53,12 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
{
// Load next input row into ring buffer, upsampling by 2x
if ((kh + blockOH - 1 + PH) % 2 == 0) // only need to load every other row due to 2x upsampling
loadUpsampleRow(inRows[(kh + blockOH - 1) % blockOH], src0, ic, ih + kh + blockOH - 1, iw);
Base::loadUpsampleRow(inRows[(kh + blockOH - 1) % blockOH], src0, ic, ih + kh + blockOH - 1, iw);
else
inRows[(kh + blockOH - 1) % blockOH] = inRows[(kh + blockOH - 2) % blockOH]; // duplicate previous row
// Multiply + accumulate kernel row
mmaBlockKW(accumRows, inRows, &weight(oc, ic, kh, 0), kh);
Base::mmaBlockKW(accumRows, inRows, &weight(oc, ic, kh, 0), kh);
}
}
@@ -67,36 +66,26 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
for (int ic = 0; ic < src1.C; ic += blockC)
{
// Load input rows into ring buffer
#pragma unroll
for (int boh = 0; boh < blockOH - 1; ++boh)
loadRow(inRows[boh], src1, ic, ih + boh, iw);
Base::template loadRows<blockOH - 1>(inRows, src1, ic, ih, iw);
// Iterate over kernel height
#pragma unroll
for (int kh = 0; kh < KH; ++kh)
{
// Load next input row into ring buffer
loadRow(inRows[(kh + blockOH - 1) % blockOH], src1, ic, ih + kh + blockOH - 1, iw);
Base::loadRow(inRows[(kh + blockOH - 1) % blockOH], src1, ic, ih + kh + blockOH - 1, iw);
// Multiply + accumulate kernel row
mmaBlockKW(accumRows, inRows, &weight(oc, src0.C + ic, kh, 0), kh);
Base::mmaBlockKW(accumRows, inRows, &weight(oc, src0.C + ic, kh, 0), kh);
}
}
// Convert accumulator to output data type, apply bias and activation
OutRows outRows;
finalizeBlock(outRows, accumRows, &bias(oc));
Base::finalizeBlock(outRows, accumRows, &bias(oc));
// Store output rows
#pragma unroll
for (int boh = 0; boh < blockOH; ++boh)
{
if (oh + boh >= dst.H)
break;
// Store output row
storeRow(outRows[boh], dst, oc, oh + boh, ow);
}
Base::template storeRows<blockOH>(outRows, dst, oc, oh, ow);
}
};
devices/sycl/sycl_conv_base_xe.h
+159 -37
View File
@@ -55,15 +55,76 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
using InRows = simd<MatmulT, blockIW * blockC>[blockOH];
using OutRows = simd<SrcDstT, blockOW * blockC>[blockOH];
// Loads a row from a src tensor, assuming the row is within the tensor bounds
#if defined(OIDN_ARCH_XEHPC) || defined(OIDN_ARCH_XE2)
// Loads a row from a tensor using LSC 2D block loads with implicit zero padding
template<int N>
static oidn_inline void loadInnerRow(simd<MatmulT, N>& row,
const TensorAccessor3D<SrcDstT, srcDstLayout>& src,
int ic, int ih, int iw)
static oidn_inline void loadRow(simd<MatmulT, N>& row,
const TensorAccessor3D<SrcDstT, srcDstLayout>& src,
int ic, int ih, int iw)
{
constexpr int lscBlockN = 64 / sizeof(SrcDstT); // LSC 2D block width in elements
static_assert(N % lscBlockN == 0, "row size must be multiple of LSC 2D block width");
const SrcDstT* surfPtr = &src(ic, 0, 0);
const uint surfWidth = src.getByteOffset.hByteStride - 1; // bytes - 1
const uint surfHeight = src.H - 1; // rows - 1
const uint surfPitch = surfWidth;
#pragma unroll
for (int n = 0; n < N; n += lscBlockN)
{
row.template select<lscBlockN, 1>(n) =
load_2d<SrcDstT, lscBlockN, 1>(surfPtr, surfWidth, surfHeight, surfPitch,
iw * blockC + n, ih);
}
}
// Loads multiple rows from a tensor using LSC 2D block loads with implicit zero padding
template<int numRows>
static oidn_inline void loadRows(InRows& rows,
const TensorAccessor3D<SrcDstT, srcDstLayout>& src,
int ic, int ih, int iw)
{
constexpr int N = blockIW * blockC;
constexpr int lscBlockN = 64 / sizeof(SrcDstT);
static_assert(N % lscBlockN == 0, "row size must be multiple of LSC 2D block width");
const SrcDstT* surfPtr = &src(ic, 0, 0);
const uint surfWidth = src.getByteOffset.hByteStride - 1; // bytes - 1
const uint surfHeight = src.H - 1; // rows - 1
const uint surfPitch = surfWidth;
#pragma unroll
for (int n = 0; n < N; n += lscBlockN)
{
auto lscBlock = load_2d<SrcDstT, lscBlockN, numRows>(
surfPtr, surfWidth, surfHeight, surfPitch,
iw * blockC + n, ih);
#pragma unroll
for (int r = 0; r < numRows; ++r)
{
rows[r].template select<lscBlockN, 1>(n) =
lscBlock.template select<lscBlockN, 1>(r * lscBlockN);
}
}
}
#else
// Loads a row from a tensor using regular LSC block loads with explicit bounds checking and padding
template<int N>
static oidn_inline void loadRow(simd<MatmulT, N>& row,
const TensorAccessor3D<SrcDstT, srcDstLayout>& src,
int ic, int ih, int iw)
{
static_assert(N % blockC == 0, "non-integer width");
constexpr int W = N / blockC;
if (ih < 0 || ih >= src.H)
{
row = 0;
return;
}
if (iw >= 0 && iw + W <= src.W)
{
// Fast path: load the entire row
@@ -89,21 +150,18 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
}
}
template<int N>
static oidn_inline void loadRow(simd<MatmulT, N>& row,
const TensorAccessor3D<SrcDstT, srcDstLayout>& src,
int ic, int ih, int iw)
template<int numRows>
static oidn_inline void loadRows(InRows& rows,
const TensorAccessor3D<SrcDstT, srcDstLayout>& src,
int ic, int ih, int iw)
{
if (ih < 0 || ih >= src.H)
{
row = 0;
return;
}
loadInnerRow(row, src, ic, ih, iw);
#pragma unroll
for (int r = 0; r < numRows; ++r)
loadRow(rows[r], src, ic, ih + r, iw);
}
#endif
// Loads a row from a virtually 2x upsampled src tensor
// Loads a row from a virtually 2x upsampled tensor
template<int N>
static oidn_inline void loadUpsampleRow(simd<MatmulT, N>& row,
const TensorAccessor3D<SrcDstT, srcDstLayout>& src,
@@ -116,19 +174,14 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
ih >>= 1;
iw >>= 1;
if (ih < 0 || ih >= src.H)
{
row = 0;
return;
}
constexpr int alignW = PW & 1; // assumption: iw%2 == PW%2
constexpr int srcW = (W >> 1) + alignW; // number of source pixels to load
simd<MatmulT, srcW * blockC> srcRow;
loadInnerRow(srcRow, src, ic, ih, iw);
loadRow(srcRow, src, ic, ih, iw);
// Upsample by copying the pixels from the source row to the destination row
#pragma unroll
for (int i = 0; i < W; ++i)
{
row.template select<blockC, 1>(i * blockC) =
@@ -136,16 +189,73 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
}
}
// Stores a row to the dst tensor
// Pixels can be stored in chunks of K to improve performance
template<int K = 1, int N>
#if defined(OIDN_ARCH_XEHPC) || defined(OIDN_ARCH_XE2)
// Stores a row to a tensor using LSC 2D block stores which silently drop out-of-bounds writes
template<int N>
static oidn_inline void storeRow(simd<SrcDstT, N>& row,
const TensorAccessor3D<SrcDstT, srcDstLayout>& dst,
int oc, int oh, int ow)
{
constexpr int lscBlockN = 64 / sizeof(SrcDstT); // LSC 2D block width in elements
static_assert(N % lscBlockN == 0, "row size must be multiple of LSC 2D block width");
SrcDstT* surfPtr = &dst(oc, 0, 0);
const uint surfWidth = dst.W * blockC * sizeof(SrcDstT) - 1; // bytes - 1
const uint surfHeight = dst.H - 1; // rows - 1
const uint surfPitch = surfWidth;
#pragma unroll
for (int n = 0; n < N; n += lscBlockN)
{
store_2d<SrcDstT, lscBlockN, 1>(
surfPtr, surfWidth, surfHeight, surfPitch,
ow * blockC + n, oh,
row.template select<lscBlockN, 1>(n).read());
}
}
// Stores multiple rows to a tensor using LSC 2D block stores which silently drop out-of-bounds writes
template<int numRows>
static oidn_inline void storeRows(OutRows& outRows,
const TensorAccessor3D<SrcDstT, srcDstLayout>& dst,
int oc, int oh, int ow)
{
constexpr int N = blockOW * blockC;
constexpr int lscBlockN = 64 / sizeof(SrcDstT); // LSC 2D block width in elements
static_assert(N % lscBlockN == 0, "row size must be multiple of LSC 2D block width");
SrcDstT* surfPtr = &dst(oc, 0, 0);
const uint surfWidth = dst.getByteOffset.hByteStride - 1; // bytes - 1
const uint surfHeight = dst.H - 1; // rows - 1
const uint surfPitch = surfWidth;
#pragma unroll
for (int n = 0; n < N; n += lscBlockN)
{
simd<SrcDstT, lscBlockN * numRows> lscBlock;
#pragma unroll
for (int r = 0; r < numRows; ++r)
{
lscBlock.template select<lscBlockN, 1>(r * lscBlockN) =
outRows[r].template select<lscBlockN, 1>(n);
}
store_2d<SrcDstT, lscBlockN, numRows>(
surfPtr, surfWidth, surfHeight, surfPitch,
ow * blockC + n, oh,
lscBlock);
}
}
#else
// Stores a row to a tensor using regular LSC block stores with explicit bounds checking
template<int N>
static oidn_inline void storeRow(simd<SrcDstT, N>& row,
const TensorAccessor3D<SrcDstT, srcDstLayout>& dst,
int oc, int oh, int ow)
{
static_assert(N % blockC == 0, "non-integer width");
constexpr int W = N / blockC;
static_assert(W % K == 0, "non-integer chunks");
//if (oh >= dst.H)
// return;
@@ -159,25 +269,37 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
else
{
// Slow path: store the in-bounds pixels of the row
constexpr int numChunks = W / K;
constexpr int chunkSize = blockC * K;
const simd<int, numChunks> owVec(ow, K); // ow, ow+K, ow+2*K, ...
simd_mask<numChunks> predVec = owVec < dst.W;
simd<uint32_t, numChunks> dstOffsetVec = dst.getByteOffset(oc, oh, 0) +
owVec * dst.getByteOffset.wByteStride;
simd<uintptr_t, numChunks> dstAddrVec = reinterpret_cast<uintptr_t>(dst.ptr) + dstOffsetVec;
const simd<int, W> owVec(ow, 1); // ow, ow+1, ow+2, ...
simd_mask<W> predVec = owVec < dst.W;
simd<uint32_t, W> dstOffsetVec = dst.getByteOffset(oc, oh, 0) +
owVec * dst.getByteOffset.wByteStride;
simd<uintptr_t, W> dstAddrVec = reinterpret_cast<uintptr_t>(dst.ptr) + dstOffsetVec;
#pragma unroll
for (int i = 0; i < numChunks; ++i)
for (int i = 0; i < W; ++i)
{
SrcDstT* dstPtr = reinterpret_cast<SrcDstT*>(uintptr_t(dstAddrVec[i]));
storeBlock(dstPtr, row.template select<chunkSize, 1>(i * chunkSize).read(),
storeBlock(dstPtr, row.template select<blockC, 1>(i * blockC).read(),
predVec.template select<1, 1>(i));
}
}
}
template<int numRows>
static oidn_inline void storeRows(OutRows& outRows,
const TensorAccessor3D<SrcDstT, srcDstLayout>& dst,
int oc, int oh, int ow)
{
#pragma unroll
for (int r = 0; r < numRows; ++r)
{
if (oh + r >= dst.H)
break;
storeRow(outRows[r], dst, oc, oh + r, ow);
}
}
#endif
// Multiply + accumulate block for one kernel row (kh)
static oidn_inline void mmaBlockKW(AccumRows& accumRows, InRows& inRows, const WeightT* weightPtr, int kh)
{
devices/sycl/sycl_conv_xe.h
+7 -17
View File
@@ -16,7 +16,6 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
using Base::KH, Base::KW, Base::PH, Base::PW;
using Base::blockC, Base::blockOH, Base::blockOW;
using typename Base::AccumRows, typename Base::InRows, typename Base::OutRows;
using Base::loadRow, Base::loadUpsampleRow, Base::mmaBlockKW, Base::finalizeBlock, Base::storeRow;
TensorAccessor3D<SrcDstT, srcDstLayout> src;
TensorAccessor4D<WeightT, weightLayout> weight;
@@ -40,25 +39,23 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
const int iw = ow - PW;
// Load input rows into a ring buffer
#pragma unroll
for (int boh = 0; boh < blockOH - 1; ++boh)
loadRow(inRows[boh], src, ic, ih + boh, iw);
Base::template loadRows<blockOH - 1>(inRows, src, ic, ih, iw);
// Iterate over kernel height
#pragma unroll
for (int kh = 0; kh < KH; ++kh)
{
// Load next input row into ring buffer
loadRow(inRows[(kh + blockOH - 1) % blockOH], src, ic, ih + (kh + blockOH - 1), iw);
Base::loadRow(inRows[(kh + blockOH - 1) % blockOH], src, ic, ih + (kh + blockOH - 1), iw);
// Multiply + accumulate kernel row
mmaBlockKW(accumRows, inRows, &weight(oc, ic, kh, 0), kh);
Base::mmaBlockKW(accumRows, inRows, &weight(oc, ic, kh, 0), kh);
}
}
// Convert accumulator to output data type, apply bias and activation
OutRows outRows;
finalizeBlock(outRows, accumRows, &bias(oc));
Base::finalizeBlock(outRows, accumRows, &bias(oc));
// Store output rows
if constexpr (fusion == Fusion::PoolDst)
@@ -75,20 +72,13 @@ OIDN_ARCH_XE_NAMESPACE_BEGIN
poolRow2x1.template replicate_vs_w<blockOW / 2, blockC * 2, blockC>(blockC));
// Store pooled row
storeRow(poolRow2x2, dst, oc, (oh + boh) / 2, ow / 2);
Base::storeRow(poolRow2x2, dst, oc, (oh + boh) / 2, ow / 2);
}
}
else
{
#pragma unroll
for (int boh = 0; boh < blockOH; ++boh)
{
if (oh + boh >= dst.H)
break;
// Store output row
storeRow(outRows[boh], dst, oc, oh + boh, ow);
}
// Store output rows
Base::template storeRows<blockOH>(outRows, dst, oc, oh, ow);
}
}
};