CPU: CPUConv cleanup

This commit is contained in:
Attila Afra
2025-12-20 19:39:52 +02:00
parent 2a52037abb
commit e78c64093e
6 changed files with 185 additions and 244 deletions
-2
View File
@@ -106,8 +106,6 @@ else()
list(APPEND OIDN_ISPC_SOURCES_COMMON
cpu_conv.ispc
cpu_conv_compute.isph
cpu_conv_compute_block.isph
)
endif()
+24 -24
View File
@@ -23,41 +23,41 @@ OIDN_NAMESPACE_BEGIN
throw std::invalid_argument("unsupported convolution bias layout/data type");
const int blockC = getTensorLayoutInfo(dstDesc.layout).blockC;
const int IC = srcDesc.getPaddedC();
const int OC = dstDesc.getPaddedC();
const int OH = dstDesc.getH();
const int OW = dstDesc.getW();
const int IC = srcDesc.getPaddedC(); // input channels
const int OC = dstDesc.getPaddedC(); // output channels
const int OH = dstDesc.getH(); // output height
const int OW = dstDesc.getW(); // output width
const int OCB = OC / blockC;
const int OCB = OC / blockC; // number of output channel blocks
blockOCB = min(OCB, ispc::CPUConvKernel_getMaxBlockOCB());
while (OCB % blockOCB != 0)
blockOCB--;
OCBB = OCB / blockOCB;
OCBB = OCB / blockOCB; // number of output channel block blocks
blockOW = ispc::CPUConvKernel_getBlockOW(blockOCB);
// Split the output width into tiles to fit into the L2 cache
// Split the output width into chunks to fit into the L2 cache
const size_t cacheSize = 512 * 1024; // FIXME: query actual size but this also works well
const int tileOW = max(
const int chunkOW = max(
static_cast<int>((cacheSize - weightDesc.getByteSize() - biasDesc.getByteSize()) /
(IC * getDataTypeSize(srcDesc.dataType) * weightDesc.getH() +
OC * getDataTypeSize(dstDesc.dataType)) - weightDesc.getW() + 1),
1);
const int maxOWT = max(OW / (2*blockOW), 1); // max number of OW tiles
OWT = min(ceil_div(OW, tileOW), maxOWT); // number of OW tiles
const int maxOWC = max(OW / (2*blockOW), 1); // max number of output width chunks
OWC = min(ceil_div(OW, chunkOW), maxOWC); // number of output width chunks
// Tweak the number of OW tiles to maximize threading efficiency
// Tweak the number of output width chunks to maximize threading efficiency
const int numThreads = engine->getNumThreads();
double bestThreadEff = 0;
for (int curOWT = OWT; curOWT < maxOWT; ++curOWT)
for (int curOWC = OWC; curOWC < maxOWC; ++curOWC)
{
const size_t N = size_t(OCBB) * OH * curOWT; // work amount
const size_t N = size_t(OCBB) * OH * curOWC; // work amount
const double threadEff = 1. - double(N % numThreads) / N;
if (threadEff > bestThreadEff)
{
OWT = curOWT;
OWC = curOWC;
bestThreadEff = threadEff;
if (threadEff >= 0.99)
break;
@@ -79,23 +79,23 @@ OIDN_NAMESPACE_BEGIN
engine->submitFunc([=]
{
const int OH = kernel.dst.H;
const int OW = kernel.dst.W;
const size_t N = size_t(OCBB) * OH * OWT;
const int OH = kernel.dst.H; // output height
const int OW = kernel.dst.W; // output width
const size_t N = size_t(OCBB) * OH * OWC; // total number of work items
parallel_for(N, [&](size_t i)
{
const size_t j = i / OCBB;
const int ocbb = int(i % OCBB);
const int oh = int(j % OH);
const int owt = int(j / OH);
const int ocbb = int(i % OCBB); // output channel block block index
const int oh = int(j % OH); // output height index
const int owc = int(j / OH); // output width chunk index
constexpr int PW = 1; // KW = 3
const int owr = OWT * (blockOW - PW - 1);
const int owBegin = owt > 0 ? (owt * OW + owr) / (OWT*blockOW) * blockOW + PW : 0;
const int owEnd = owt+1 < OWT ? ((owt+1) * OW + owr) / (OWT*blockOW) * blockOW + PW : OW;
const int owr = OWC * (blockOW - PW - 1);
const int owBegin = owc > 0 ? (owc * OW + owr) / (OWC*blockOW) * blockOW + PW : 0;
const int owEnd = owc+1 < OWC ? ((owc+1) * OW + owr) / (OWC*blockOW) * blockOW + PW : OW;
ispc::CPUConvKernel_run(&kernel, blockOCB, ocbb * blockOCB, oh, owBegin, owEnd);
ispc::CPUConvKernel_run_f32(&kernel, blockOCB, ocbb * blockOCB, oh, owBegin, owEnd);
});
}, ct);
}
+2 -2
View File
@@ -19,9 +19,9 @@ OIDN_NAMESPACE_BEGIN
private:
CPUEngine* engine;
int blockOCB; // block of output channel blocks
int blockOW; // block of output width
int blockOW; // output block width
int OCBB; // number of output channel block blocks
int OWT; // number of output width tiles
int OWC; // number of output width chunks
};
OIDN_NAMESPACE_END
+159 -64
View File
@@ -12,20 +12,163 @@ struct CPUConvKernel
uniform bool relu;
};
#define _CPUConvKernel_compute(T, blockOCB) CPUConvKernel_compute_##T##_##blockOCB
#define CPUConvKernel_compute(T, blockOCB) _CPUConvKernel_compute(T, blockOCB)
#define _CPUConvKernel_computeBlock(T, blockOCB, blockOW) CPUConvKernel_computeBlock_##T##_##blockOCB##_##blockOW
#define CPUConvKernel_computeBlock(T, blockOCB, blockOW) _CPUConvKernel_computeBlock(T, blockOCB, blockOW)
#define T float
#define blockC programCount
#define KW 3 // kerned width
#define KH 3 // kerned height
#define PW 1 // padding width on each side
#define PH 1 // padding height on each side
#define blockC programCount
template<typename T, int blockOCB, int blockOW>
inline unmasked void CPUConvKernel_computeBlock(
const uniform int8* uniform srcPtr,
uniform size_t srcHByteStride,
const uniform int8* uniform weightPtr,
const uniform int8* uniform biasPtr,
uniform int8* uniform dstPtr,
uniform size_t dstCByteStride,
uniform size_t khEnd,
uniform size_t kwBegin, uniform size_t kwEnd,
uniform bool relu)
{
varying T accum[blockOCB][blockOW];
if (biasPtr)
{
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
accum[bocb][bow] = *((const varying T* uniform)biasPtr + bocb);
}
}
else
{
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
accum[bocb][bow] = *((const varying T* uniform)(dstPtr + bocb * dstCByteStride) + bow);
}
}
#pragma nounroll
for (uniform size_t kh = 0; kh < khEnd; ++kh)
{
#pragma nounroll
for (uniform size_t kw = kwBegin; kw < kwEnd; ++kw)
{
#pragma unroll
for (uniform size_t i = 0; i < blockC; ++i)
{
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
const varying T weightVec =
*((const varying T* uniform)weightPtr + (bocb * KW * KH + kw) * blockC + i);
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
{
const varying T srcVec = *((const uniform T* uniform)srcPtr + (bow + kw - PW) * blockC + i);
accum[bocb][bow] += srcVec * weightVec;
}
}
}
}
srcPtr += srcHByteStride;
weightPtr += KW * blockC * blockC * sizeof(uniform T);
}
if (relu)
{
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
accum[bocb][bow] = max(accum[bocb][bow], 0);
}
}
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
*((varying T* uniform)(dstPtr + bocb * dstCByteStride) + bow) = accum[bocb][bow];
}
}
template<typename T, int blockOCB, int blockOW>
unmasked void CPUConvKernel_compute(const uniform CPUConvKernel* uniform self,
uniform int ocb, uniform int oh,
uniform int owBegin, uniform int owEnd)
{
const uniform int oc = ocb * blockC;
#if KH == 3 && PH == 1
const uniform int khBegin = oh > 0 ? 0 : 1;
const uniform int khEnd = oh < self->dst.H-1 ? 3 : 2;
#else
const uniform int khBegin = max(PH - oh, 0);
const uniform int khEnd = KH - max(PH + oh - (self->dst.H-1), 0);
#endif
for (uniform int ic = 0; ic < self->src.C; ic += blockC)
{
const uniform int8* uniform srcPtr = Tensor_getPtr<T>(self->src, ic, oh + khBegin - PH, owBegin);
const uniform int8* uniform weightPtr = Tensor_getPtr<T>(self->weight, oc, ic, khBegin, 0);
const uniform int8* uniform biasPtr = (ic == 0) ? Tensor_getPtr<T>(self->bias, oc) : NULL;
uniform int8* uniform dstPtr = Tensor_getPtr<T>(self->dst, oc, oh, owBegin);
const uniform bool relu = self->relu && ic == (self->src.C - blockC);
uniform int ow = owBegin; // owBegin/owEnd *must* be aligned to block boundaries
while (ow < owEnd)
{
if (ow > PW - 1 && ow + blockOW + PW - 1 < self->dst.W)
{
// Fast path (no padding, width blocking)
CPUConvKernel_computeBlock<T, blockOCB, blockOW>(
srcPtr, self->src.hByteStride,
weightPtr, biasPtr,
dstPtr, self->dst.CByteStride,
khEnd - khBegin,
0, KW,
relu);
srcPtr += blockOW * blockC * sizeof(uniform T);
dstPtr += blockOW * blockC * sizeof(uniform T);
ow += blockOW;
}
else
{
// Slow path (padding, no width blocking)
CPUConvKernel_computeBlock<T, blockOCB, 1>(
srcPtr, self->src.hByteStride,
weightPtr, biasPtr,
dstPtr, self->dst.CByteStride,
khEnd - khBegin,
#if KW == 3 && PW == 1
ow > 0 ? 0 : 1,
ow < self->dst.W-1 ? 3 : 2,
#else
max(PW - ow, 0),
KW - max(PW + ow - (self->dst.W-1), 0),
#endif
relu);
srcPtr += blockC * sizeof(uniform T);
dstPtr += blockC * sizeof(uniform T);
ow++;
}
}
}
}
// Kernel variants optimized for different ISAs
#if defined(OIDN_ARCH_AVX512)
#define maxBlockOCB 4
@@ -51,54 +194,6 @@ struct CPUConvKernel
#error "Unsupported architecture"
#endif
#if maxBlockOCB >= 1
#define blockOCB 1
#define blockOW 1
#include "cpu_conv_compute_block.isph"
#undef blockOW
#define blockOW blockOW1
#include "cpu_conv_compute_block.isph"
#include "cpu_conv_compute.isph"
#undef blockOW
#undef blockOCB
#endif
#if maxBlockOCB >= 2
#define blockOCB 2
#define blockOW 1
#include "cpu_conv_compute_block.isph"
#undef blockOW
#define blockOW blockOW2
#include "cpu_conv_compute_block.isph"
#include "cpu_conv_compute.isph"
#undef blockOW
#undef blockOCB
#endif
#if maxBlockOCB >= 3
#define blockOCB 3
#define blockOW 1
#include "cpu_conv_compute_block.isph"
#undef blockOW
#define blockOW blockOW3
#include "cpu_conv_compute_block.isph"
#include "cpu_conv_compute.isph"
#undef blockOW
#undef blockOCB
#endif
#if maxBlockOCB >= 4
#define blockOCB 4
#define blockOW 1
#include "cpu_conv_compute_block.isph"
#undef blockOW
#define blockOW blockOW4
#include "cpu_conv_compute_block.isph"
#include "cpu_conv_compute.isph"
#undef blockOW
#undef blockOCB
#endif
export uniform int CPUConvKernel_getMaxBlockOCB()
{
return maxBlockOCB;
@@ -122,21 +217,21 @@ export uniform int CPUConvKernel_getBlockOW(uniform int blockOCB)
}
}
export void CPUConvKernel_run(const uniform CPUConvKernel* uniform self,
uniform int blockOCB, uniform int ocb, uniform int oh,
uniform int owBegin, uniform int owEnd)
export void CPUConvKernel_run_f32(const uniform CPUConvKernel* uniform self,
uniform int blockOCB, uniform int ocb, uniform int oh,
uniform int owBegin, uniform int owEnd)
{
switch (blockOCB)
{
case 1: CPUConvKernel_compute(T, 1)(self, ocb, oh, owBegin, owEnd); break;
case 1: CPUConvKernel_compute<float, 1, blockOW1>(self, ocb, oh, owBegin, owEnd); break;
#if maxBlockOCB >= 2
case 2: CPUConvKernel_compute(T, 2)(self, ocb, oh, owBegin, owEnd); break;
case 2: CPUConvKernel_compute<float, 2, blockOW2>(self, ocb, oh, owBegin, owEnd); break;
#endif
#if maxBlockOCB >= 3
case 3: CPUConvKernel_compute(T, 3)(self, ocb, oh, owBegin, owEnd); break;
case 3: CPUConvKernel_compute<float, 3, blockOW3>(self, ocb, oh, owBegin, owEnd); break;
#endif
#if maxBlockOCB >= 4
case 4: CPUConvKernel_compute(T, 4)(self, ocb, oh, owBegin, owEnd); break;
case 4: CPUConvKernel_compute<float, 4, blockOW4>(self, ocb, oh, owBegin, owEnd); break;
#endif
}
}
-67
View File
@@ -1,67 +0,0 @@
// Copyright 2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
unmasked void CPUConvKernel_compute(T, blockOCB)(const uniform CPUConvKernel* uniform self,
uniform int ocb, uniform int oh,
uniform int owBegin, uniform int owEnd)
{
const uniform int oc = ocb * blockC;
#if KH == 3 && PH == 1
const uniform int khBegin = oh > 0 ? 0 : 1;
const uniform int khEnd = oh < self->dst.H-1 ? 3 : 2;
#else
const uniform int khBegin = max(PH - oh, 0);
const uniform int khEnd = KH - max(PH + oh - (self->dst.H-1), 0);
#endif
for (uniform int ic = 0; ic < self->src.C; ic += blockC)
{
const uniform int8* uniform srcPtr = Tensor_getPtr<T>(self->src, ic, oh + khBegin - PH, owBegin);
const uniform int8* uniform weightPtr = Tensor_getPtr<T>(self->weight, oc, ic, khBegin, 0);
const uniform int8* uniform biasPtr = (ic == 0) ? Tensor_getPtr<T>(self->bias, oc) : NULL;
uniform int8* uniform dstPtr = Tensor_getPtr<T>(self->dst, oc, oh, owBegin);
const uniform bool relu = self->relu && ic == (self->src.C - blockC);
uniform int ow = owBegin; // owBegin/owEnd *must* be aligned to block boundaries
while (ow < owEnd)
{
if (ow > PW - 1 && ow + blockOW + PW - 1 < self->dst.W)
{
// Fast path (no padding, width blocking)
CPUConvKernel_computeBlock(T, blockOCB, blockOW)(
srcPtr, self->src.hByteStride,
weightPtr, biasPtr,
dstPtr, self->dst.CByteStride,
khEnd - khBegin,
0, KW,
relu);
srcPtr += blockOW * blockC * sizeof(uniform T);
dstPtr += blockOW * blockC * sizeof(uniform T);
ow += blockOW;
}
else
{
// Slow path (padding, no width blocking)
CPUConvKernel_computeBlock(T, blockOCB, 1)(
srcPtr, self->src.hByteStride,
weightPtr, biasPtr,
dstPtr, self->dst.CByteStride,
khEnd - khBegin,
#if KW == 3 && PW == 1
ow > 0 ? 0 : 1,
ow < self->dst.W-1 ? 3 : 2,
#else
max(PW - ow, 0),
KW - max(PW + ow - (self->dst.W-1), 0),
#endif
relu);
srcPtr += blockC * sizeof(uniform T);
dstPtr += blockC * sizeof(uniform T);
ow++;
}
}
}
}
-85
View File
@@ -1,85 +0,0 @@
// Copyright 2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
inline unmasked void CPUConvKernel_computeBlock(T, blockOCB, blockOW)(
const uniform int8* uniform srcPtr,
uniform size_t srcHByteStride,
const uniform int8* uniform weightPtr,
const uniform int8* uniform biasPtr,
uniform int8* uniform dstPtr,
uniform size_t dstCByteStride,
uniform size_t khEnd,
uniform size_t kwBegin, uniform size_t kwEnd,
uniform bool relu)
{
varying T accum[blockOCB][blockOW];
if (biasPtr)
{
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
accum[bocb][bow] = *((const varying T* uniform)biasPtr + bocb);
}
}
else
{
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
accum[bocb][bow] = *((const varying T* uniform)(dstPtr + bocb * dstCByteStride) + bow);
}
}
#pragma nounroll
for (uniform size_t kh = 0; kh < khEnd; ++kh)
{
#pragma nounroll
for (uniform size_t kw = kwBegin; kw < kwEnd; ++kw)
{
#pragma unroll
for (uniform size_t i = 0; i < blockC; ++i)
{
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
const varying T weightVec =
*((const varying T* uniform)weightPtr + (bocb * KW * KH + kw) * blockC + i);
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
{
const varying T srcVec = *((const uniform T* uniform)srcPtr + (bow + kw - PW) * blockC + i);
accum[bocb][bow] += srcVec * weightVec;
}
}
}
}
srcPtr += srcHByteStride;
weightPtr += KW * blockC * blockC * sizeof(uniform T);
}
if (relu)
{
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
accum[bocb][bow] = max(accum[bocb][bow], 0);
}
}
#pragma unroll
for (uniform size_t bocb = 0; bocb < blockOCB; ++bocb)
{
#pragma unroll
for (uniform size_t bow = 0; bow < blockOW; ++bow)
*((varying T* uniform)(dstPtr + bocb * dstCByteStride) + bow) = accum[bocb][bow];
}
}