diff --git a/dnn-providers/integration-tests/tests/gpu-ref/TestGpuFpReferenceConvolution.cpp b/dnn-providers/integration-tests/tests/gpu-ref/TestGpuFpReferenceConvolution.cpp index 0bcfbb42ef45..a76d04d9db0f 100644 --- a/dnn-providers/integration-tests/tests/gpu-ref/TestGpuFpReferenceConvolution.cpp +++ b/dnn-providers/integration-tests/tests/gpu-ref/TestGpuFpReferenceConvolution.cpp @@ -221,6 +221,30 @@ TEST(TestConvolutionValidation, ThrowsOnOutputDimValueMismatch) std::invalid_argument); } +TEST(TestConvolutionValidation, ThrowsOnBatchMismatch) +{ + // wgrad walks the batch of y and reads x at the same batch index. + const Tensor x({1, 1, 4, 4}); + const Tensor w({1, 1, 3, 3}); + const Tensor y({2, 1, 2, 2}); + + EXPECT_THROW(hipdnn_test_sdk::utilities::validateConvolutionParams( + x, w, y, Vec{1, 1}, Vec{1, 1}, Vec{0, 0}, Vec{0, 0}), + std::invalid_argument); +} + +TEST(TestConvolutionValidation, ThrowsOnOutputChannelMismatch) +{ + // fprop writes one y channel per weight output channel. + const Tensor x({1, 1, 4, 4}); + const Tensor w({2, 1, 3, 3}); + const Tensor y({1, 1, 2, 2}); + + EXPECT_THROW(hipdnn_test_sdk::utilities::validateConvolutionParams( + x, w, y, Vec{1, 1}, Vec{1, 1}, Vec{0, 0}, Vec{0, 0}), + std::invalid_argument); +} + // ============================================================================ // TestGpuConvFwdRefValidation — same validation error paths as above, but // exercised through the GPU code path (GpuFpReferenceConvolution::fprop). diff --git a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/ConvolutionValidation.hpp b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/ConvolutionValidation.hpp index 91f788c2be4f..0a99abb2cab4 100644 --- a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/ConvolutionValidation.hpp +++ b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/ConvolutionValidation.hpp @@ -15,7 +15,9 @@ namespace hipdnn_test_sdk::utilities /// Validates convolution parameters against tensor dimensions. /// Checks that strides/dilations/padding vectors have correct sizes and valid values, -/// and that the output tensor spatial dimensions match the expected convolution output. +/// that the output tensor's batch size matches the input's and its channel count matches +/// the weight's output-channel count, and that the output tensor spatial dimensions match +/// the expected convolution output. /// /// Callers are responsible for validating tensor dimension counts before calling this /// (e.g. restricting to 3D/4D/5D for GPU kernels, or >= 3 for a generic CPU path). @@ -43,6 +45,20 @@ inline void validateConvolutionParams(const std::vector& xDims, "Output tensor must have the same number of dimensions as input"); } + if(yDims[0] != xDims[0]) + { + throw std::invalid_argument("Output tensor batch size " + std::to_string(yDims[0]) + + " does not match input batch size " + + std::to_string(xDims[0])); + } + + if(yDims[1] != wDims[0]) + { + throw std::invalid_argument("Output tensor channel count " + std::to_string(yDims[1]) + + " does not match weight output-channel count " + + std::to_string(wDims[0])); + } + if(strides.size() != nSpatialDims) { throw std::invalid_argument("Strides must have exactly " + std::to_string(nSpatialDims) diff --git a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceConvolution.hpp b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceConvolution.hpp index 267497f5acf6..509a612dc075 100644 --- a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceConvolution.hpp +++ b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceConvolution.hpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -16,6 +17,8 @@ namespace hipdnn_test_sdk::utilities class CpuFpReferenceConvolution { + static constexpr auto PREFIX = "CpuFpReferenceConvolution: "; + public: // Check if this CPU implementation supports the given node configuration static bool isApplicable(const hipdnn_flatbuffers_sdk::data_objects::Node& node) @@ -61,7 +64,7 @@ class CpuFpReferenceConvolution const std::vector& prePadding, const std::vector& postPadding) { - validateInput(x, w, y, strides, dilations, prePadding, postPadding); + validateInput(x, w, y, {"x", "w", "y"}, strides, dilations, prePadding, postPadding); // Extract dimensions - NC[spatial...] format for x/y, [G*K][C][spatial...] for w const auto& xDims = x.dims(); @@ -73,7 +76,7 @@ class CpuFpReferenceConvolution const int64_t totalOutputChannels = wDims[0]; // G * K (flattened) int64_t channelsPerGroup = wDims[1]; // C - int64_t nSpatialDims = static_cast(xDims.size()) - 2; + const std::size_t nSpatialDims = xDims.size() - 2; std::vector xSpatialDims(xDims.begin() + 2, xDims.end()); std::vector kernelSpatialDims(wDims.begin() + 2, wDims.end()); std::vector ySpatialDims(yDims.begin() + 2, yDims.end()); @@ -82,88 +85,85 @@ class CpuFpReferenceConvolution const int64_t nGroups = nInputChannels / channelsPerGroup; const int64_t yChannelsPerGroup = totalOutputChannels / nGroups; - // This lambda computes a single element of the y tensor - auto convolutionFunc = [&](const std::vector& indices) { - auto gIdx = indices[0]; // group index - auto nIdx = indices[1]; // batch index - auto kIdx = indices[2]; // y channel within group + // Tensors are addressed through hoisted base pointers and strides, which + // validateInput has restricted to the dense layout. + const XDataType* xBase = x.memory().hostData(); + const WDataType* wBase = w.memory().hostData(); + YDataType* yBase = y.memory().hostData(); + + const auto& xStrides = x.strides(); + const auto& wStrides = w.strides(); + const auto& yStrides = y.strides(); + + // This lambda computes a single element of the y tensor. `window` is per-thread + // scratch owned by the parallel functor, so rebuilding it here costs no allocation. + auto convolutionFunc = [&](hipdnn_test_sdk::detail::ConvolutionWindow& window, + const std::vector& indices) { + const int64_t gIdx = indices[0]; // group index + const int64_t nIdx = indices[1]; // batch index + const int64_t kIdx = indices[2]; // y channel within group + + // Add 3 because [gIdx, nIdx, kIdx] are the first 3 elements + const int64_t* ySpatialIndices = indices.data() + 3; + + // Which kernel taps hit the logical x tensor depends only on the y spatial + // position, so resolve the whole window once instead of per channel. + window.build( + nSpatialDims, + kernelSpatialDims.data(), + wStrides.data() + 2, + xStrides.data() + 2, + [&](size_t dim, int64_t kernelIndex) { + const int64_t xIndex = (ySpatialIndices[dim] * strides[dim]) + + (kernelIndex * dilations[dim]) - prePadding[dim]; + + // In either case, this position does not exist in the logical x tensor. + // 1. (y_idx * stride) + (kernel_idx * dilation) - prePadding < 0 + // => (y_idx * stride) + (kernel_idx * dilation) < prePadding + // 2. (y_idx * stride) + (kernel_idx * dilation) - prePadding >= x_dim + // => (y_idx * stride) + (kernel_idx * dilation) >= x_dim + prePadding + // It is implicit in Case 2 that the position could be in the postPadding region. + return (xIndex < 0 || xIndex >= xSpatialDims[dim]) ? -1 : xIndex; + }); - // Add 3 because [gIdx, nIdx, cIdx] are the first 3 elements - std::vector ySpatialIndices(indices.begin() + 3, indices.end()); + // Weight dims: [yChannels, xChannels/groupCount, ...] + // Thus, we index via flattened y channel index and group-offset x channel index (c). + const int64_t yChannel = (gIdx * yChannelsPerGroup) + kIdx; + const WDataType* wFilter = wBase + (yChannel * wStrides[0]); + + // Input dims: [n, xChannel, ...], indexed via the global x channel index. + const XDataType* xBatch + = xBase + (nIdx * xStrides[0]) + (gIdx * channelsPerGroup * xStrides[1]); auto accumulator = static_cast(0.0f); - const int64_t baseInputChannel = gIdx * channelsPerGroup; + const auto& taps = window.expand(); for(int64_t c = 0; c < channelsPerGroup; ++c) { - const int64_t xChannel = baseInputChannel + c; - - // Iterate kernel spatial positions - hipdnn_data_sdk::utilities::iterateAlongDimensions( - kernelSpatialDims, [&](const std::vector& kernelSpatialIndices) { - std::vector xSpatialIndices(static_cast(nSpatialDims)); - bool validPosition = true; - - // For each spatial dimension, calculate the corresponding x index - for(int64_t dim = 0; dim < nSpatialDims; ++dim) - { - auto dimIdx = static_cast(dim); - xSpatialIndices[dimIdx] - = (ySpatialIndices[dimIdx] * strides[dimIdx]) - + (kernelSpatialIndices[dimIdx] * dilations[dimIdx]) - - prePadding[dimIdx]; - - // In either case, this position does not exist in the logical x tensor. - // 1. (y_idx * stride) + (kernel_idx * dilation) - prePadding < 0 - // => (y_idx * stride) + (kernel_idx * dilation) < prePadding - // 2. (y_idx * stride) + (kernel_idx * dilation) - prePadding >= x_dim - // => (y_idx * stride) + (kernel_idx * dilation) >= x_dim + prePadding - // It is implicit in Case 2 that the position could be in the postPadding region. - if(xSpatialIndices[dimIdx] < 0 - || xSpatialIndices[dimIdx] >= xSpatialDims[dimIdx]) - { - validPosition = false; - break; - } - } - - if(validPosition) - { - // Input dims: [n, xChannel, ...] - // Thus, we index via global x channel index. - auto xFullIndices = hipdnn_data_sdk::utilities::buildTensorIndices( - nIdx, xChannel, xSpatialIndices); - - // Weight dims: [yChannels, - // xChannels/groupCount, ...] Thus, we index via flattened y channel index and - // group-offset x channel index (c). - const int64_t wIdx = (gIdx * yChannelsPerGroup) + kIdx; - auto wFullIndices = hipdnn_data_sdk::utilities::buildTensorIndices( - wIdx, c, kernelSpatialIndices); - - XDataType xVal = x.getHostValue(xFullIndices); - WDataType wVal = w.getHostValue(wFullIndices); - - accumulator = accumulator - + (static_cast(xVal) - * static_cast(wVal)); - } - }); + const XDataType* xChannel = xBatch + (c * xStrides[1]); + const WDataType* wChannel = wFilter + (c * wStrides[1]); + + for(const auto& tap : taps) + { + accumulator = accumulator + + (static_cast(xChannel[tap.sourceOffset]) + * static_cast(wChannel[tap.windowOffset])); + } } - const int64_t yChannel = (gIdx * yChannelsPerGroup) + kIdx; - auto yFullIndices - = hipdnn_data_sdk::utilities::buildTensorIndices(nIdx, yChannel, ySpatialIndices); + const int64_t yOffset = (nIdx * yStrides[0]) + (yChannel * yStrides[1]) + + hipdnn_test_sdk::detail::flatOffset( + ySpatialIndices, yStrides.data() + 2, nSpatialDims); - y.setHostValue(static_cast(accumulator), yFullIndices); + yBase[yOffset] = static_cast(accumulator); }; // Build dimensions for parallel iteration std::vector parallelDims = {nGroups, nBatch, yChannelsPerGroup}; parallelDims.insert(parallelDims.end(), ySpatialDims.begin(), ySpatialDims.end()); - auto parallelFunc - = hipdnn_test_sdk::detail::makeParallelTensorFunctor(convolutionFunc, parallelDims); + auto parallelFunc = hipdnn_test_sdk::detail::makeParallelTensorFunctorWithScratch< + hipdnn_test_sdk::detail::ConvolutionWindow>(convolutionFunc, parallelDims); parallelFunc(std::thread::hardware_concurrency()); y.memory().markHostModified(); @@ -190,7 +190,8 @@ class CpuFpReferenceConvolution const std::vector& prePadding, const std::vector& postPadding) { - validateInput(gradX, w, gradY, strides, dilations, prePadding, postPadding); + validateInput( + gradX, w, gradY, {"dx", "w", "dy"}, strides, dilations, prePadding, postPadding); // Extract dimensions - NC[spatial...] format for x/y, [G*K][C][spatial...] for w const auto& xDims = gradX.dims(); @@ -201,7 +202,7 @@ class CpuFpReferenceConvolution const int64_t totalOutputChannels = wDims[0]; // G * K (flattened) int64_t channelsPerGroup = wDims[1]; // C - int64_t nSpatialDims = static_cast(xDims.size()) - 2; + const std::size_t nSpatialDims = xDims.size() - 2; std::vector xSpatialDims(xDims.begin() + 2, xDims.end()); std::vector kernelSpatialDims(wDims.begin() + 2, wDims.end()); std::vector ySpatialDims(yDims.begin() + 2, yDims.end()); @@ -211,93 +212,89 @@ class CpuFpReferenceConvolution const int64_t nGroups = nInputChannels / channelsPerGroup; // G const int64_t yChannelsPerGroup = totalOutputChannels / nGroups; // K - // This lambda computes a single element of the x gradient tensor (dx) - auto convolutionFunc = [&](const std::vector& indices) { - auto gIdx = indices[0]; // group index - auto nIdx = indices[1]; // batch index - auto cIdx = indices[2]; // channel index within group + // Addressed through hoisted base pointers and strides; see fprop. + DxDataType* gradXBase = gradX.memory().hostData(); + const WDataType* wBase = w.memory().hostData(); + const DyDataType* gradYBase = gradY.memory().hostData(); - // Add 3 because [gIdx, nIdx, cIdx] are the first 3 elements - std::vector xSpatialIndices(indices.begin() + 3, indices.end()); + const auto& xStrides = gradX.strides(); + const auto& wStrides = w.strides(); + const auto& yStrides = gradY.strides(); - auto vAcc = static_cast(0.0f); + // This lambda computes a single element of the x gradient tensor (dx). `window` is + // per-thread scratch owned by the parallel functor. + auto convolutionFunc = [&](hipdnn_test_sdk::detail::ConvolutionWindow& window, + const std::vector& indices) { + const int64_t gIdx = indices[0]; // group index + const int64_t nIdx = indices[1]; // batch index + const int64_t cIdx = indices[2]; // channel index within group - // Iterate over each spatial position of the kernel for contributing y gradients - hipdnn_data_sdk::utilities::iterateAlongDimensions( - kernelSpatialDims, [&](const std::vector& kernelSpatialIndices) { - std::vector ySpatialIndices(static_cast(nSpatialDims)); - bool validPosition = true; - - // For each spatial dimension, calculate the corresponding y gradient index - for(int64_t dim = 0; dim < nSpatialDims; ++dim) + // Add 3 because [gIdx, nIdx, cIdx] are the first 3 elements + const int64_t* xSpatialIndices = indices.data() + 3; + + // Which kernel taps have a contributing y gradient depends only on the x + // spatial position, so resolve the whole window once instead of per y channel. + window.build( + nSpatialDims, + kernelSpatialDims.data(), + wStrides.data() + 2, + yStrides.data() + 2, + [&](size_t dim, int64_t kernelIndex) -> int64_t { + const int64_t tmp + = xSpatialIndices[dim] + prePadding[dim] - (kernelIndex * dilations[dim]); + + // Check if the current x position could have contributed to an y element. If the + // remainder is non-zero, this combination is not aligned with the stride, so it's not a valid + // mapping from the forward pass. + if(tmp % strides[dim] != 0) { - auto dimIdx = static_cast(dim); - const int64_t tmp = xSpatialIndices[dimIdx] + prePadding[dimIdx] - - (kernelSpatialIndices[dimIdx] * dilations[dimIdx]); - - // Check if the current x position could have contributed to an y element. If the - // remainder is non-zero, this combination is not aligned with the stride, so it's not a valid - // mapping from the forward pass. - if(tmp % strides[dimIdx] != 0) - { - validPosition = false; - break; - } - - ySpatialIndices[dimIdx] = tmp / strides[dimIdx]; - - // Check if position does not exist in the logical y tensor. - // 1. (x_idx + prePadding - (kernel_idx * dilation)) / stride < 0 - // => numerator < 0 => sampling from a location before the y tensor - // 2. (x_idx + prePadding - (kernel_idx * dilation)) / stride >= y_dim - // => x_idx + prePadding >= (y_dim * stride) + (kernel_idx * dilation) => beyond the y tensor - if(ySpatialIndices[dimIdx] < 0 - || ySpatialIndices[dimIdx] >= ySpatialDims[dimIdx]) - { - validPosition = false; - break; - } + return -1; } - if(validPosition) - { - // Iterate over each y channel in the group, as they all contribute to the x gradient - for(int64_t k = 0; k < yChannelsPerGroup; ++k) - { - auto yChannelIdx = (gIdx * yChannelsPerGroup) + k; - - auto gradOutputFullIndices - = hipdnn_data_sdk::utilities::buildTensorIndices( - nIdx, yChannelIdx, ySpatialIndices); - - auto wBatchIdx = yChannelIdx; - auto wChannelIdx = cIdx; - auto wFullIndices = hipdnn_data_sdk::utilities::buildTensorIndices( - wBatchIdx, wChannelIdx, kernelSpatialIndices); - - DyDataType vOut = gradY.getHostValue(gradOutputFullIndices); - WDataType vWei = w.getHostValue(wFullIndices); - - vAcc = vAcc - + (static_cast(vOut) - * static_cast(vWei)); - } - } + // Check if position does not exist in the logical y tensor. + // 1. (x_idx + prePadding - (kernel_idx * dilation)) / stride < 0 + // => numerator < 0 => sampling from a location before the y tensor + // 2. (x_idx + prePadding - (kernel_idx * dilation)) / stride >= y_dim + // => x_idx + prePadding >= (y_dim * stride) + (kernel_idx * dilation) => beyond the y tensor + const int64_t yIndex = tmp / strides[dim]; + return (yIndex < 0 || yIndex >= ySpatialDims[dim]) ? -1 : yIndex; }); + const int64_t gradYBatch + = (nIdx * yStrides[0]) + (gIdx * yChannelsPerGroup * yStrides[1]); + const int64_t wFilter = (gIdx * yChannelsPerGroup * wStrides[0]) + (cIdx * wStrides[1]); + + auto vAcc = static_cast(0.0f); + + for(const auto& tap : window.expand()) + { + // Iterate over each y channel in the group, as they all contribute to the x gradient + for(int64_t k = 0; k < yChannelsPerGroup; ++k) + { + const DyDataType vOut + = gradYBase[gradYBatch + (k * yStrides[1]) + tap.sourceOffset]; + const WDataType vWei = wBase[wFilter + (k * wStrides[0]) + tap.windowOffset]; + + vAcc = vAcc + + (static_cast(vOut) + * static_cast(vWei)); + } + } + const int64_t xChannelIdx = (gIdx * channelsPerGroup) + cIdx; - auto gradInputFullIndices = hipdnn_data_sdk::utilities::buildTensorIndices( - nIdx, xChannelIdx, xSpatialIndices); + const int64_t gradXOffset = (nIdx * xStrides[0]) + (xChannelIdx * xStrides[1]) + + hipdnn_test_sdk::detail::flatOffset( + xSpatialIndices, xStrides.data() + 2, nSpatialDims); - gradX.setHostValue(static_cast(vAcc), gradInputFullIndices); + gradXBase[gradXOffset] = static_cast(vAcc); }; // Build dimensions for parallel iteration std::vector parallelDims = {nGroups, nBatch, channelsPerGroup}; parallelDims.insert(parallelDims.end(), xSpatialDims.begin(), xSpatialDims.end()); - auto parallelFunc - = hipdnn_test_sdk::detail::makeParallelTensorFunctor(convolutionFunc, parallelDims); + auto parallelFunc = hipdnn_test_sdk::detail::makeParallelTensorFunctorWithScratch< + hipdnn_test_sdk::detail::ConvolutionWindow>(convolutionFunc, parallelDims); parallelFunc(std::thread::hardware_concurrency()); gradX.memory().markHostModified(); @@ -323,7 +320,8 @@ class CpuFpReferenceConvolution const std::vector& prePadding, const std::vector& postPadding) { - validateInput(x, gradW, gradY, strides, dilations, prePadding, postPadding); + validateInput( + x, gradW, gradY, {"x", "dw", "dy"}, strides, dilations, prePadding, postPadding); // Extract dimensions - NCHW format for x/y, [G*K][C][Y][X] for w (4D flattened) const auto& xDims = x.dims(); @@ -332,7 +330,7 @@ class CpuFpReferenceConvolution int64_t nBatch = yDims[0]; - int64_t nSpatialDims = static_cast(xDims.size()) - 2; + const std::size_t nSpatialDims = xDims.size() - 2; std::vector xSpatialDims(xDims.begin() + 2, xDims.end()); std::vector kernelSpatialDims(wDims.begin() + 2, wDims.end()); std::vector ySpatialDims(yDims.begin() + 2, yDims.end()); @@ -345,91 +343,92 @@ class CpuFpReferenceConvolution const int64_t nGroups = nInputChannels / channelsPerGroup; // G const int64_t yChannelsPerGroup = totalOutputChannels / nGroups; // K - auto convolutionFunc = [&](const std::vector& indices) { - auto gIdx = indices[0]; - auto kIdx = indices[1]; - auto cIdx = indices[2]; - - // Add 3 because [gIdx, nIdx, cIdx] are the first 3 elements - std::vector kernelSpatialIndices(indices.begin() + 3, indices.end()); + // Addressed through hoisted base pointers and strides; see fprop. + const XDataType* xBase = x.memory().hostData(); + DwDataType* gradWBase = gradW.memory().hostData(); + const DyDataType* gradYBase = gradY.memory().hostData(); + + const auto& xStrides = x.strides(); + const auto& wStrides = gradW.strides(); + const auto& yStrides = gradY.strides(); + + // `window` is per-thread scratch owned by the parallel functor. Its extent is the + // y gradient's spatial extent, so it is walked rather than expanded. + auto convolutionFunc = [&](hipdnn_test_sdk::detail::ConvolutionWindow& window, + const std::vector& indices) { + const int64_t gIdx = indices[0]; + const int64_t kIdx = indices[1]; + const int64_t cIdx = indices[2]; + + // Add 3 because [gIdx, kIdx, cIdx] are the first 3 elements + const int64_t* kernelSpatialIndices = indices.data() + 3; + + // Which y gradient positions sample a real x element depends only on the + // kernel spatial position, so resolve the whole window once instead of per batch. + window.build(nSpatialDims, + ySpatialDims.data(), + yStrides.data() + 2, + xStrides.data() + 2, + [&](size_t dim, int64_t yIndex) { + const int64_t xIndex = (yIndex * strides[dim]) + + (kernelSpatialIndices[dim] * dilations[dim]) + - prePadding[dim]; + + return (xIndex < 0 || xIndex >= xSpatialDims[dim]) ? -1 : xIndex; + }); + + const int64_t yChannelIdx = (gIdx * yChannelsPerGroup) + kIdx; + const int64_t xChannelIdx = (gIdx * channelsPerGroup) + cIdx; + const DyDataType* gradYChannel = gradYBase + (yChannelIdx * yStrides[1]); + const XDataType* xChannel = xBase + (xChannelIdx * xStrides[1]); auto vAcc = static_cast(0.0f); - // Iterate over each spatial position of the kernel for contributing y gradients - hipdnn_data_sdk::utilities::iterateAlongDimensions( - ySpatialDims, [&](const std::vector& ySpatialIndices) { - std::vector xSpatialIndices(static_cast(nSpatialDims)); - - bool validPosition = true; - - // For each spatial dimension, calculate the corresponding y gradient index - for(int64_t dim = 0; dim < nSpatialDims; ++dim) - { - auto dimIdx = static_cast(dim); - - const int64_t tmp = (ySpatialIndices[dimIdx] * strides[dimIdx]) - + (kernelSpatialIndices[dimIdx] * dilations[dimIdx]) - - prePadding[dimIdx]; - - xSpatialIndices[dimIdx] = tmp; - - if(xSpatialIndices[dimIdx] < 0 - || xSpatialIndices[dimIdx] >= xSpatialDims[dimIdx]) - { - validPosition = false; - break; - } - } - - if(validPosition) - { - for(int64_t n = 0; n < nBatch; ++n) - { - auto yChannelIdx = (gIdx * yChannelsPerGroup) + kIdx; - - auto gradOutputFullIndices - = hipdnn_data_sdk::utilities::buildTensorIndices( - n, yChannelIdx, ySpatialIndices); - - auto xChannelIdx = (gIdx * channelsPerGroup) + cIdx; - - auto xFullIndices = hipdnn_data_sdk::utilities::buildTensorIndices( - n, xChannelIdx, xSpatialIndices); - - DyDataType vOut = gradY.getHostValue(gradOutputFullIndices); - - XDataType vIn = x.getHostValue(xFullIndices); + window.forEachTap([&](const hipdnn_test_sdk::detail::ConvolutionWindow::Tap& tap) { + for(int64_t n = 0; n < nBatch; ++n) + { + const DyDataType vOut = gradYChannel[(n * yStrides[0]) + tap.windowOffset]; + const XDataType vIn = xChannel[(n * xStrides[0]) + tap.sourceOffset]; - vAcc = vAcc - + (static_cast(vOut) - * static_cast(vIn)); - } - } - }); + vAcc = vAcc + + (static_cast(vOut) + * static_cast(vIn)); + } + }); - auto wN = (gIdx * yChannelsPerGroup) + kIdx; - auto wFullIndices - = hipdnn_data_sdk::utilities::buildTensorIndices(wN, cIdx, kernelSpatialIndices); + const int64_t gradWOffset = (yChannelIdx * wStrides[0]) + (cIdx * wStrides[1]) + + hipdnn_test_sdk::detail::flatOffset(kernelSpatialIndices, + wStrides.data() + 2, + nSpatialDims); - gradW.setHostValue(static_cast(vAcc), wFullIndices); + gradWBase[gradWOffset] = static_cast(vAcc); }; // Build dimensions for parallel iteration std::vector parallelDims = {nGroups, yChannelsPerGroup, channelsPerGroup}; parallelDims.insert(parallelDims.end(), kernelSpatialDims.begin(), kernelSpatialDims.end()); - auto parallelFunc - = hipdnn_test_sdk::detail::makeParallelTensorFunctor(convolutionFunc, parallelDims); + auto parallelFunc = hipdnn_test_sdk::detail::makeParallelTensorFunctorWithScratch< + hipdnn_test_sdk::detail::ConvolutionWindow>(convolutionFunc, parallelDims); parallelFunc(std::thread::hardware_concurrency()); gradW.memory().markHostModified(); } private: + // Argument names reported by validation, in (x, w, y) role order. + struct ArgumentNames + { + const char* x; + const char* w; + const char* y; + }; + template static void validateInput(const hipdnn_data_sdk::utilities::TensorBase& x, const hipdnn_data_sdk::utilities::TensorBase& w, const hipdnn_data_sdk::utilities::TensorBase& y, + const ArgumentNames& names, const std::vector& strides, const std::vector& dilations, const std::vector& prePadding, @@ -438,9 +437,14 @@ class CpuFpReferenceConvolution if(x.dims().size() < 3) { throw std::invalid_argument( - "Input tensor must have at least 3 dimensions (N, C, spatial...)"); + std::string(PREFIX) + + "input tensor must have at least 3 dimensions (N, C, spatial...)"); } + hipdnn_test_sdk::detail::validateNoRaggedTensor(x, PREFIX, names.x); + hipdnn_test_sdk::detail::validateNoRaggedTensor(w, PREFIX, names.w); + hipdnn_test_sdk::detail::validateNoRaggedTensor(y, PREFIX, names.y); + hipdnn_test_sdk::utilities::validateConvolutionParams( x, w, y, strides, dilations, prePadding, postPadding); } diff --git a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceLayernorm.hpp b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceLayernorm.hpp index 57b227b15a61..b1ebb3669ba4 100644 --- a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceLayernorm.hpp +++ b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceLayernorm.hpp @@ -4,9 +4,18 @@ #pragma once #include "hipdnn_data_sdk/utilities/ShapeUtilities.hpp" +#include +#include +#include #include #include #include +#include +#include +#include +#include +#include +#include #include namespace hipdnn_test_sdk::utilities @@ -14,6 +23,8 @@ namespace hipdnn_test_sdk::utilities class CpuFpReferenceLayernorm { + static constexpr auto PREFIX = "CpuFpReferenceLayernorm: "; + public: // Layer normalization forward pass. // Normalizes over the last `normalizedDimCount` dimensions of the input tensor. @@ -36,8 +47,11 @@ class CpuFpReferenceLayernorm // - Avoids catastrophic cancellation (no E[x²] - E[x]² subtraction) // - Is numerically stable for arbitrary value ranges and element counts // - // Scale and bias, if provided, have shape matching the normalized dimensions. - // Mean and rstd outputs, if provided, have shape matching the batch dimensions. + // y has the shape of x. + // Scale and bias, if provided, share one shape: the normalized dimensions, optionally + // preceded by 1s. + // Mean and rstd outputs, if provided, share one shape: the batch dimensions, optionally + // followed by 1s, with at least one dimension. template * rstd = nullptr) { const auto& dims = x.dims(); - auto ndim = static_cast(dims.size()); - if(ndim < 1) - { - throw std::runtime_error("Layernorm fprop requires at least 1D tensor."); - } + validateNormalizedDimCount(dims, normalizedDimCount, "fprop"); - if(normalizedDimCount < 1 || normalizedDimCount > ndim) + if(std::any_of(dims.begin(), dims.end(), [](int64_t d) { return d <= 0; })) { - throw std::runtime_error( - "normalizedDimCount must be between 1 and the number of tensor dimensions."); + throw std::runtime_error(std::string(PREFIX) + + "fprop requires every dimension to be positive."); } - if(scale != nullptr && bias != nullptr && scale->dims().size() != bias->dims().size()) + if(y.dims() != dims) { - throw std::runtime_error("Scale and bias tensors must have the same rank."); + throw std::runtime_error(std::string(PREFIX) + + "fprop requires y to have the same shape as x."); } - - // Split dimensions into batch dims and normalized dims - std::vector batchDims; - std::vector normalizedDims; - if(mean != nullptr) - { - batchDims = mean->dims(); - } - else if(rstd != nullptr) + if(scale != nullptr && bias != nullptr && scale->dims() != bias->dims()) { - batchDims = rstd->dims(); + throw std::runtime_error(std::string(PREFIX) + + "fprop requires scale and bias to have the same shape."); } - else + if(mean != nullptr && rstd != nullptr && mean->dims() != rstd->dims()) { - batchDims - = std::vector(dims.begin(), dims.begin() + ndim - normalizedDimCount); - } - if(scale != nullptr) - { - normalizedDims = scale->dims(); - } - else if(bias != nullptr) - { - normalizedDims = bias->dims(); - } - else - { - normalizedDims - = std::vector(dims.begin() + ndim - normalizedDimCount, dims.end()); + throw std::runtime_error(std::string(PREFIX) + + "fprop requires mean and rstd to have the same shape."); } - for(auto d : normalizedDims) + hipdnn_test_sdk::detail::validateNoRaggedTensor(x, PREFIX, "x"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(y, PREFIX, "y"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(scale, PREFIX, "scale"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(bias, PREFIX, "bias"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(mean, PREFIX, "mean"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(rstd, PREFIX, "rstd"); + + const auto* affine = (scale != nullptr) ? scale : bias; + if(affine != nullptr) { - if(d <= 0) - { - throw std::runtime_error( - "Normalized dimensions must all be positive (no zero-size dimensions)."); - } + validateAffineShape(dims, affine->dims(), normalizedDimCount, "fprop"); } - for(auto d : batchDims) + const auto* stat = (mean != nullptr) ? mean : rstd; + if(stat != nullptr) { - if(d <= 0) - { - throw std::runtime_error( - "Batch dimensions must all be positive (no zero-size dimensions)."); - } + validateStatShape(dims, stat->dims(), normalizedDimCount, "fprop"); } - auto epsilonCompute = static_cast(epsilon); + const auto normCount = static_cast(normalizedDimCount); + const size_t batchDimCount = dims.size() - normCount; + const auto split = dims.begin() + static_cast(batchDimCount); + std::vector batchExtents(dims.begin(), split); + const std::vector normExtents(split, dims.end()); - // If batchDims is empty (entire tensor is normalized), use a single scalar iteration - if(batchDims.empty()) + // A whole-tensor normalization has a single batch position, indexed by no dimension. + if(batchExtents.empty()) { - batchDims.push_back(1); + batchExtents.push_back(1); } + const auto epsilonCompute = static_cast(epsilon); + + const XDataType* xBase = x.memory().hostData(); + YDataType* yBase = y.memory().hostData(); + const ScaleBiasDataType* scaleBase + = (scale != nullptr) ? scale->memory().hostData() : nullptr; + const ScaleBiasDataType* biasBase = (bias != nullptr) ? bias->memory().hostData() : nullptr; + MeanRstdDataType* meanBase = (mean != nullptr) ? mean->memory().hostData() : nullptr; + MeanRstdDataType* rstdBase = (rstd != nullptr) ? rstd->memory().hostData() : nullptr; + + const auto& xStrides = x.strides(); + const auto& yStrides = y.strides(); + + // Every batch position walks the normalized dims the same way, so that walk is a + // flat offset table built once. x and y cover the normalized dims with their + // trailing strides; scale and bias with their trailing strides too, since any + // leading dims they have are 1s. + const auto xNormOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + normExtents, xStrides.data() + batchDimCount); + const auto yNormOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + normExtents, yStrides.data() + batchDimCount); + const auto scaleOffsets = (scale != nullptr) + ? hipdnn_test_sdk::detail::buildDenseOffsets( + normExtents, trailingStrides(*scale, normCount)) + : std::vector{}; + const auto biasOffsets = (bias != nullptr) + ? hipdnn_test_sdk::detail::buildDenseOffsets( + normExtents, trailingStrides(*bias, normCount)) + : std::vector{}; + const int64_t* meanStrides = (mean != nullptr) ? mean->strides().data() : nullptr; + const int64_t* rstdStrides = (rstd != nullptr) ? rstd->strides().data() : nullptr; + + const auto normElementCount = xNormOffsets.size(); + auto layernormFpropFunc = [&](const std::vector& batchIndices) { + const int64_t xBatchOffset = hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), xStrides.data(), batchDimCount); + const int64_t yBatchOffset = hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), yStrides.data(), batchDimCount); + // Pass 1: Welford's online algorithm for mean and variance - int64_t count = 0; - auto batchMean = static_cast(0.0); - auto m2 = static_cast(0.0); - - hipdnn_data_sdk::utilities::iterateAlongDimensions( - normalizedDims, [&](const std::vector& normIndices) { - auto fullIndices - = buildFullIndices(batchIndices, normIndices, ndim, normalizedDimCount); - auto xVal = static_cast(x.getHostValue(fullIndices)); - - count++; - auto delta = xVal - batchMean; - batchMean += delta / static_cast(count); - auto delta2 = xVal - batchMean; - m2 += delta * delta2; - }); - - auto batchVariance = m2 / static_cast(count); - auto invStd = static_cast(1.0) - / hipdnn_data_sdk::types::sqrt(batchVariance + epsilonCompute); + const auto [batchMean, invStd] = welfordMeanAndRstd( + xBase + xBatchOffset, xNormOffsets, epsilonCompute); // Pass 2: normalize and apply scale/bias - hipdnn_data_sdk::utilities::iterateAlongDimensions( - normalizedDims, [&](const std::vector& normIndices) { - auto fullIndices - = buildFullIndices(batchIndices, normIndices, ndim, normalizedDimCount); - auto xVal = static_cast(x.getHostValue(fullIndices)); - auto xHat = (xVal - batchMean) * invStd; - - ComputeDataType yVal = xHat; - if(scale != nullptr) - { - yVal - = static_cast(scale->getHostValue(normIndices)) * yVal; - } - if(bias != nullptr) - { - yVal = yVal + static_cast(bias->getHostValue(normIndices)); - } - - y.setHostValue(static_cast(yVal), fullIndices); - }); + for(size_t element = 0; element < normElementCount; ++element) + { + auto xVal + = static_cast(xBase[xBatchOffset + xNormOffsets[element]]); + auto xHat = (xVal - batchMean) * invStd; + + ComputeDataType yVal = xHat; + if(scaleBase != nullptr) + { + yVal = static_cast(scaleBase[scaleOffsets[element]]) * yVal; + } + if(biasBase != nullptr) + { + yVal = yVal + static_cast(biasBase[biasOffsets[element]]); + } + + yBase[yBatchOffset + yNormOffsets[element]] = static_cast(yVal); + } // Save mean and rstd for this batch position if requested - if(mean != nullptr) + if(meanBase != nullptr) { - mean->setHostValue(static_cast(batchMean), batchIndices); + meanBase[hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), meanStrides, batchDimCount)] + = static_cast(batchMean); } - if(rstd != nullptr) + if(rstdBase != nullptr) { - rstd->setHostValue(static_cast(invStd), batchIndices); + rstdBase[hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), rstdStrides, batchDimCount)] + = static_cast(invStd); } }; // Parallelize over batch dimensions auto parallelFunc - = hipdnn_test_sdk::detail::makeParallelTensorFunctor(layernormFpropFunc, batchDims); + = hipdnn_test_sdk::detail::makeParallelTensorFunctor(layernormFpropFunc, batchExtents); parallelFunc(std::thread::hardware_concurrency()); y.memory().markHostModified(); @@ -222,8 +241,11 @@ class CpuFpReferenceLayernorm // dscale_n = dscale_sum_b // dbias_n = dbias_sum_b // - // Scale and bias have shape matching the normalized dimensions. - // Mean and rstd inputs, if provided, have shape matching the batch dimensions. + // x and dx have the shape of dy. + // Scale, dscale and dbias share one shape: the normalized dimensions, optionally + // preceded by 1s. + // Mean and rstd inputs, if provided, share one shape: the batch dimensions, optionally + // followed by 1s, with at least one dimension. template & dx, hipdnn_data_sdk::utilities::TensorBase& dscale, hipdnn_data_sdk::utilities::TensorBase& dbias, - [[maybe_unused]] const double epsilon, + const double epsilon, const hipdnn_data_sdk::utilities::TensorBase* mean, const hipdnn_data_sdk::utilities::TensorBase* rstd, const int64_t normalizedDimCount) { const auto& dims = dy.dims(); - auto ndim = static_cast(dims.size()); - if(ndim < 1) - { - throw std::runtime_error("Layernorm bprop requires at least 1D tensor."); - } + validateNormalizedDimCount(dims, normalizedDimCount, "bprop"); - if(normalizedDimCount < 1 || normalizedDimCount > ndim) + if(x.dims() != dims || dx.dims() != dims) { - throw std::runtime_error( - "normalizedDimCount must be between 1 and the number of tensor dimensions."); + throw std::runtime_error(std::string(PREFIX) + + "bprop requires x and dx to have the same shape as dy."); } - if(scale.dims() != dscale.dims() || scale.dims() != dbias.dims()) { throw std::runtime_error( - "Scale, dscale and dbias tensors must have the same dimensions."); + std::string(PREFIX) + + "bprop requires scale, dscale and dbias to have the same shape."); } - if((mean == nullptr) != (rstd == nullptr)) { throw std::runtime_error( - "Layernorm backward requires both mean and rstd to be provided, or neither."); + std::string(PREFIX) + + "bprop requires both mean and rstd to be provided, or neither."); } - if(mean != nullptr && mean->dims() != rstd->dims()) { - throw std::runtime_error("Mean and rstd tensors must have the same dimensions."); + throw std::runtime_error(std::string(PREFIX) + + "bprop requires mean and rstd to have the same shape."); } - // Split dimensions into batch dims and normalized dims - auto normalizedDims = scale.dims(); - const int64_t normalizedDimsSize = std::accumulate( - normalizedDims.begin(), normalizedDims.end(), int64_t{1}, std::multiplies{}); - std::vector batchDims; + hipdnn_test_sdk::detail::validateNoRaggedTensor(dy, PREFIX, "dy"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(x, PREFIX, "x"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(scale, PREFIX, "scale"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(dx, PREFIX, "dx"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(dscale, PREFIX, "dscale"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(dbias, PREFIX, "dbias"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(mean, PREFIX, "mean"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(rstd, PREFIX, "rstd"); + + validateAffineShape(dims, scale.dims(), normalizedDimCount, "bprop"); if(mean != nullptr) { - batchDims = mean->dims(); - } - else if(dims.size() == normalizedDims.size()) - { - batchDims = std::vector(dims.size(), 1); - for(size_t i = 0; i < dims.size(); ++i) - { - if(dims[i] != normalizedDims[i]) - { - batchDims[i] = dims[i]; - } - } - } - else - { - batchDims = std::vector(static_cast(ndim - normalizedDimCount), 1); - for(size_t i = 0; i < static_cast(ndim - normalizedDimCount); ++i) - { - batchDims[i] = dims[i]; - } - } - auto strideOrder = hipdnn_data_sdk::utilities::extractStrideOrder(x.strides()); - auto batchStrides = hipdnn_data_sdk::utilities::generateStrides(batchDims, strideOrder); - const int64_t batchDimsSize = std::accumulate( - batchDims.begin(), batchDims.end(), int64_t{1}, std::multiplies{}); - - std::vector tmpMean; - std::vector tmpRstd; - if(mean == nullptr || rstd == nullptr) - { - tmpMean = std::vector(static_cast(batchDimsSize)); - tmpRstd = std::vector(static_cast(batchDimsSize)); + validateStatShape(dims, mean->dims(), normalizedDimCount, "bprop"); } - // If batchDims is empty (entire tensor is normalized), use a single scalar iteration - if(batchDims.empty()) + const auto normCount = static_cast(normalizedDimCount); + const size_t batchDimCount = dims.size() - normCount; + const auto split = dims.begin() + static_cast(batchDimCount); + std::vector batchExtents(dims.begin(), split); + const std::vector normExtents(split, dims.end()); + + const int64_t normalizedDimsSize = std::accumulate( + normExtents.begin(), normExtents.end(), int64_t{1}, std::multiplies{}); + + // A whole-tensor normalization has a single batch position, indexed by no dimension. + if(batchExtents.empty()) { - batchDims.push_back(1); + batchExtents.push_back(1); } + const auto epsilonCompute = static_cast(epsilon); + + const DyDataType* dyBase = dy.memory().hostData(); + const DxDataType* xBase = x.memory().hostData(); + const ScaleBiasDataType* scaleBase = scale.memory().hostData(); + DxDataType* dxBase = dx.memory().hostData(); + ScaleBiasDataType* dscaleBase = dscale.memory().hostData(); + ScaleBiasDataType* dbiasBase = dbias.memory().hostData(); + const MeanRstdDataType* meanBase = (mean != nullptr) ? mean->memory().hostData() : nullptr; + const MeanRstdDataType* rstdBase = (rstd != nullptr) ? rstd->memory().hostData() : nullptr; + + const auto& dyStrides = dy.strides(); + const auto& xStrides = x.strides(); + const auto& dxStrides = dx.strides(); + const int64_t* meanStrides = (mean != nullptr) ? mean->strides().data() : nullptr; + const int64_t* rstdStrides = (rstd != nullptr) ? rstd->strides().data() : nullptr; + + // Each pass holds one index space fixed and walks the other, and the walk is identical + // every time, so both walks are flat offset tables built once. dy/x/dx cover the + // normalized dims with their trailing strides and the batch dims with their leading + // ones; scale covers the normalized dims with its trailing strides, since any leading + // dims it has are 1s. + const auto dyNormOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + normExtents, dyStrides.data() + batchDimCount); + const auto xNormOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + normExtents, xStrides.data() + batchDimCount); + const auto dxNormOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + normExtents, dxStrides.data() + batchDimCount); + const auto scaleNormOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + normExtents, trailingStrides(scale, normCount)); + + const auto dyBatchOffsets + = hipdnn_test_sdk::detail::buildDenseOffsets(batchExtents, dyStrides.data()); + const auto xBatchOffsets + = hipdnn_test_sdk::detail::buildDenseOffsets(batchExtents, xStrides.data()); + + const auto normElementCount = dyNormOffsets.size(); + const auto batchElementCount = dyBatchOffsets.size(); + + // The statistics of every batch position, row-major over batchExtents: read from + // mean/rstd when provided, recomputed from x otherwise. Pass 1 fills them and pass 2 + // reads them back in the order it walks the batch. + std::vector batchMean(batchElementCount); + std::vector batchRstd(batchElementCount); + const auto batchRowMajorStrides = hipdnn_data_sdk::utilities::generateStrides(batchExtents); + // Pass 1: backward values auto layernormBpropValuesFunc = [&](const std::vector& batchIndices) { + const int64_t dyBatchOffset = hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), dyStrides.data(), batchDimCount); + const int64_t xBatchOffset = hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), xStrides.data(), batchDimCount); + const int64_t dxBatchOffset = hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), dxStrides.data(), batchDimCount); + auto sumDyScaleX = static_cast(0.0); auto sumDyScale = static_cast(0.0); - hipdnn_data_sdk::utilities::iterateAlongDimensions( - normalizedDims, [&](const std::vector& normIndices) { - auto fullIndices - = buildFullIndices(batchIndices, normIndices, ndim, normalizedDimCount); - auto dyVal = static_cast(dy.getHostValue(fullIndices)); - auto scaleVal = static_cast(scale.getHostValue(normIndices)); - auto xVal = static_cast(x.getHostValue(fullIndices)); - - sumDyScaleX += dyVal * scaleVal * xVal; - sumDyScale += dyVal * scaleVal; - }); + for(size_t element = 0; element < normElementCount; ++element) + { + auto dyVal + = static_cast(dyBase[dyBatchOffset + dyNormOffsets[element]]); + auto scaleVal = static_cast(scaleBase[scaleNormOffsets[element]]); + auto xVal + = static_cast(xBase[xBatchOffset + xNormOffsets[element]]); + + sumDyScaleX += dyVal * scaleVal * xVal; + sumDyScale += dyVal * scaleVal; + } ComputeDataType meanVal; ComputeDataType rstdVal; - if(mean == nullptr || rstd == nullptr) + if(meanBase == nullptr) { - int64_t count = 0; - meanVal = static_cast(0.0); - auto m2 = static_cast(0.0); - - hipdnn_data_sdk::utilities::iterateAlongDimensions( - normalizedDims, [&](const std::vector& normIndices) { - auto fullIndices - = buildFullIndices(batchIndices, normIndices, ndim, normalizedDimCount); - auto xVal = static_cast(x.getHostValue(fullIndices)); - - count++; - auto delta = xVal - meanVal; - meanVal += delta / static_cast(count); - auto delta2 = xVal - meanVal; - m2 += delta * delta2; - }); - - auto batchVariance = m2 / static_cast(count); - rstdVal = static_cast(1.0) - / hipdnn_data_sdk::types::sqrt(batchVariance - + static_cast(epsilon)); - - auto idx = static_cast(std::inner_product( - batchIndices.begin(), batchIndices.end(), batchStrides.begin(), int64_t{0})); - tmpMean[idx] = meanVal; - tmpRstd[idx] = rstdVal; + std::tie(meanVal, rstdVal) = welfordMeanAndRstd( + xBase + xBatchOffset, xNormOffsets, epsilonCompute); } else { - meanVal = static_cast(mean->getHostValue(batchIndices)); - rstdVal = static_cast(rstd->getHostValue(batchIndices)); + meanVal = static_cast(meanBase[hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), meanStrides, batchDimCount)]); + rstdVal = static_cast(rstdBase[hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), rstdStrides, batchDimCount)]); } + const auto statIdx = static_cast(hipdnn_test_sdk::detail::flatOffset( + batchIndices.data(), batchRowMajorStrides.data(), batchIndices.size())); + batchMean[statIdx] = meanVal; + batchRstd[statIdx] = rstdVal; + auto a = rstdVal * rstdVal * rstdVal * (sumDyScaleX - sumDyScale * meanVal) / static_cast(normalizedDimsSize); auto b = rstdVal * sumDyScale / static_cast(normalizedDimsSize) - a * meanVal; - hipdnn_data_sdk::utilities::iterateAlongDimensions( - normalizedDims, [&](const std::vector& normIndices) { - auto fullIndices - = buildFullIndices(batchIndices, normIndices, ndim, normalizedDimCount); - auto dyVal = static_cast(dy.getHostValue(fullIndices)); - auto scaleVal = static_cast(scale.getHostValue(normIndices)); - auto xVal = static_cast(x.getHostValue(fullIndices)); - auto dxVal = rstdVal * dyVal * scaleVal - a * xVal - b; - dx.setHostValue(static_cast(dxVal), fullIndices); - }); + for(size_t element = 0; element < normElementCount; ++element) + { + auto dyVal + = static_cast(dyBase[dyBatchOffset + dyNormOffsets[element]]); + auto scaleVal = static_cast(scaleBase[scaleNormOffsets[element]]); + auto xVal + = static_cast(xBase[xBatchOffset + xNormOffsets[element]]); + auto dxVal = rstdVal * dyVal * scaleVal - a * xVal - b; + dxBase[dxBatchOffset + dxNormOffsets[element]] = static_cast(dxVal); + } }; // Pass 2: backward weights + const int64_t* dscaleStrides = trailingStrides(dscale, normCount); + const int64_t* dbiasStrides = trailingStrides(dbias, normCount); auto layernormBpropWeightsFunc = [&](const std::vector& normIndices) { + const int64_t dyNormOffset = hipdnn_test_sdk::detail::flatOffset( + normIndices.data(), dyStrides.data() + batchDimCount, normCount); + const int64_t xNormOffset = hipdnn_test_sdk::detail::flatOffset( + normIndices.data(), xStrides.data() + batchDimCount, normCount); + auto dscaleVal = static_cast(0.0); auto dbiasVal = static_cast(0.0); - hipdnn_data_sdk::utilities::iterateAlongDimensions( - batchDims, [&](const std::vector& batchIndices) { - auto fullIndices - = buildFullIndices(batchIndices, normIndices, ndim, normalizedDimCount); - auto dyVal = static_cast(dy.getHostValue(fullIndices)); - auto xVal = static_cast(x.getHostValue(fullIndices)); - ComputeDataType meanVal; - ComputeDataType rstdVal; - if(mean == nullptr || rstd == nullptr) - { - auto idx = static_cast(std::inner_product(batchIndices.begin(), - batchIndices.end(), - batchStrides.begin(), - int64_t{0})); - meanVal = tmpMean[idx]; - rstdVal = tmpRstd[idx]; - } - else - { - meanVal = static_cast(mean->getHostValue(batchIndices)); - rstdVal = static_cast(rstd->getHostValue(batchIndices)); - } - dscaleVal += dyVal * (xVal - meanVal) * rstdVal; - dbiasVal += dyVal; - }); - dscale.setHostValue(static_cast(dscaleVal), normIndices); - dbias.setHostValue(static_cast(dbiasVal), normIndices); + for(size_t element = 0; element < batchElementCount; ++element) + { + auto dyVal + = static_cast(dyBase[dyNormOffset + dyBatchOffsets[element]]); + auto xVal + = static_cast(xBase[xNormOffset + xBatchOffsets[element]]); + dscaleVal += dyVal * (xVal - batchMean[element]) * batchRstd[element]; + dbiasVal += dyVal; + } + + dscaleBase[hipdnn_test_sdk::detail::flatOffset( + normIndices.data(), dscaleStrides, normCount)] + = static_cast(dscaleVal); + dbiasBase[hipdnn_test_sdk::detail::flatOffset( + normIndices.data(), dbiasStrides, normCount)] + = static_cast(dbiasVal); }; - // Parallelize over batch dimensions + // Parallelize pass 1 over the batch and pass 2 over the normalized dimensions auto parallelValuesFunc = hipdnn_test_sdk::detail::makeParallelTensorFunctor( - layernormBpropValuesFunc, batchDims); + layernormBpropValuesFunc, batchExtents); parallelValuesFunc(std::thread::hardware_concurrency()); auto parallelWeightsFunc = hipdnn_test_sdk::detail::makeParallelTensorFunctor( - layernormBpropWeightsFunc, normalizedDims); + layernormBpropWeightsFunc, normExtents); parallelWeightsFunc(std::thread::hardware_concurrency()); dx.memory().markHostModified(); @@ -434,23 +467,96 @@ class CpuFpReferenceLayernorm } private: - // Build full N-dimensional indices from batch indices and normalized indices. - // For a tensor of ndim dimensions with the last normalizedDimCount being normalized: - // fullIndices = [batchIndices..., normIndices...] - // Handles the case where batchDims was padded with a leading 1 (scalar batch). - static std::vector buildFullIndices(const std::vector& batchIndices, - const std::vector& normIndices, - int64_t ndim, - int64_t normalizedDimCount) + static void validateNormalizedDimCount(const std::vector& dims, + int64_t normalizedDimCount, + const char* pass) + { + const auto ndim = static_cast(dims.size()); + if(ndim < 1) + { + throw std::runtime_error(std::string(PREFIX) + pass + " requires at least 1D tensor."); + } + if(normalizedDimCount < 1 || normalizedDimCount > ndim) + { + throw std::runtime_error( + std::string(PREFIX) + + "normalizedDimCount must be between 1 and the number of tensor dimensions."); + } + } + + // scale/bias shape: the trailing normalizedDimCount dims of x, optionally preceded by 1s. + static void validateAffineShape(const std::vector& xDims, + const std::vector& affineDims, + int64_t normalizedDimCount, + const char* pass) + { + const auto count = static_cast(normalizedDimCount); + if(static_cast(affineDims.size()) < count + || !std::equal(affineDims.end() - count, affineDims.end(), xDims.end() - count) + || std::any_of( + affineDims.begin(), affineDims.end() - count, [](int64_t d) { return d != 1; })) + { + throw std::runtime_error(std::string(PREFIX) + pass + + " requires scale/bias to have the trailing " + "normalizedDimCount dims of the input, optionally " + "preceded by 1s."); + } + } + + // mean/rstd shape: the leading batch dims of x, optionally followed by 1s. A rank-0 + // tensor holds no element, so the shape keeps at least one dimension even when every + // dimension of x is normalized. + static void validateStatShape(const std::vector& xDims, + const std::vector& statDims, + int64_t normalizedDimCount, + const char* pass) + { + const auto count = static_cast(xDims.size()) + - static_cast(normalizedDimCount); + if(statDims.empty() || static_cast(statDims.size()) < count + || !std::equal(xDims.begin(), xDims.begin() + count, statDims.begin()) + || std::any_of( + statDims.begin() + count, statDims.end(), [](int64_t d) { return d != 1; })) + { + throw std::runtime_error(std::string(PREFIX) + pass + + " requires mean/rstd to have the leading batch dims of " + "the input, optionally followed by 1s, and at least " + "one dimension."); + } + } + + // Strides of the last `count` dims of a validated scale-shaped tensor. + static const int64_t* trailingStrides(const hipdnn_data_sdk::utilities::ITensor& tensor, + size_t count) + { + return tensor.strides().data() + (tensor.strides().size() - count); + } + + // Welford's online mean and variance over one batch position's normalized elements. + // Returns {mean, 1 / sqrt(variance + epsilon)}. + template + static std::pair welfordMeanAndRstd( + const XDataType* xBatch, const std::vector& xNormOffsets, ComputeDataType epsilon) { - auto batchDimCount = ndim - normalizedDimCount; - std::vector fullIndices; - fullIndices.reserve(static_cast(ndim)); - fullIndices.insert( - fullIndices.end(), batchIndices.begin(), batchIndices.begin() + batchDimCount); - fullIndices.insert( - fullIndices.end(), normIndices.end() - normalizedDimCount, normIndices.end()); - return fullIndices; + int64_t count = 0; + auto mean = static_cast(0.0); + auto m2 = static_cast(0.0); + + for(const auto offset : xNormOffsets) + { + auto xVal = static_cast(xBatch[offset]); + + count++; + auto delta = xVal - mean; + mean += delta / static_cast(count); + auto delta2 = xVal - mean; + m2 += delta * delta2; + } + + auto variance = m2 / static_cast(count); + auto rstd + = static_cast(1.0) / hipdnn_data_sdk::types::sqrt(variance + epsilon); + return {mean, rstd}; } }; diff --git a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceRMSNorm.hpp b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceRMSNorm.hpp index df610e55c374..7a9f47461f5c 100644 --- a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceRMSNorm.hpp +++ b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/CpuFpReferenceRMSNorm.hpp @@ -4,13 +4,16 @@ #pragma once #include -#include #include #include +#include +#include #include +#include #include #include +#include #include #include @@ -19,6 +22,8 @@ namespace hipdnn_test_sdk::utilities class CpuFpReferenceRMSNorm { + static constexpr auto PREFIX = "CpuFpReferenceRMSNorm: "; + static void computeLeadingAndReductionDims(const std::vector& xDims, const std::vector& scaleDims, std::vector& leadingDims, @@ -37,8 +42,9 @@ class CpuFpReferenceRMSNorm if(reductionStart == rank) { // Validator should have rejected this; defensive guard for direct callers. - throw std::runtime_error("RMSNorm: scale has no trailing dims matching input — no " - "normalized axes can be derived."); + throw std::runtime_error(std::string(PREFIX) + + "scale has no trailing dims matching input — no " + "normalized axes can be derived."); } // Leading dims: [0, reductionStart) — batch + leading dims preserved through @@ -49,6 +55,15 @@ class CpuFpReferenceRMSNorm reductionDims.assign(xDims.begin() + splitOffset, xDims.end()); } + // invRms keeps the leading dims of x and collapses the reduction dims to 1. + static std::vector invRmsShape(const std::vector& xDims, + size_t reductionStart) + { + std::vector dims = xDims; + std::fill(dims.begin() + static_cast(reductionStart), dims.end(), 1); + return dims; + } + public: /// RMSNorm forward: y = x / RMS(x) * scale [+ bias] /// @@ -85,15 +100,31 @@ class CpuFpReferenceRMSNorm if(xDims.size() < 2) { - throw std::runtime_error("RMSNorm forward requires at least 2D input tensor (batch and " - "at least one feature dim)."); + throw std::runtime_error(std::string(PREFIX) + + "forward requires at least 2D input tensor (batch and " + "at least one feature dim)."); } if(scaleDims.size() != xDims.size()) { - throw std::runtime_error("RMSNorm forward requires scale rank to equal input rank."); + throw std::runtime_error(std::string(PREFIX) + + "forward requires scale rank to equal input rank."); + } + if(y.dims() != xDims) + { + throw std::runtime_error(std::string(PREFIX) + + "forward requires y to have the same shape as x."); + } + if(bias != nullptr && bias->dims() != scaleDims) + { + throw std::runtime_error(std::string(PREFIX) + + "forward requires bias to have the same shape as scale."); } - const size_t rank = xDims.size(); + hipdnn_test_sdk::detail::validateNoRaggedTensor(x, PREFIX, "x"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(scale, PREFIX, "scale"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(y, PREFIX, "y"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(invRms, PREFIX, "invRms"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(bias, PREFIX, "bias"); // Compute leading and reduction dims based on input and scale shapes std::vector leadingDims; @@ -101,70 +132,105 @@ class CpuFpReferenceRMSNorm computeLeadingAndReductionDims(xDims, scaleDims, leadingDims, reductionDims); const auto reductionStart = leadingDims.size(); + if(invRms != nullptr && invRms->dims() != invRmsShape(xDims, reductionStart)) + { + throw std::runtime_error(std::string(PREFIX) + + "forward requires invRms to have the input shape with " + "the normalized dims set to 1."); + } + const auto reductionCount = std::accumulate( reductionDims.begin(), reductionDims.end(), int64_t{1}, std::multiplies<>{}); const auto reductionCountCompute = static_cast(reductionCount); const auto epsilonCompute = static_cast(epsilon); + // Tensors are addressed through hoisted base pointers and strides, which the checks + // above have restricted to the dense layout. + const XDataType* xBase = x.memory().hostData(); + const ScaleDataType* scaleBase = scale.memory().hostData(); + YDataType* yBase = y.memory().hostData(); + ComputeDataType* invRmsBase = (invRms != nullptr) ? invRms->memory().hostData() : nullptr; + const ScaleDataType* biasBase = (bias != nullptr) ? bias->memory().hostData() : nullptr; + + const auto& xStrides = x.strides(); + const auto& yStrides = y.strides(); + const auto& scaleStrides = scale.strides(); + const int64_t* invRmsStrides = (invRms != nullptr) ? invRms->strides().data() : nullptr; + + // Reduction-region offset tables: one flat offset per reduction position, built + // once and shared by every leading position. scale/bias need only this table - + // their leading-region strides never contribute, since scaleIdx is zero there. + // Running example: input [N, C, H, W], scale [1, 1, H, W] + // leadingDims = [N, C] reductionDims = [H, W] redIdx = [h, w] + // xRedOffsets[i] / yRedOffsets[i] = flat offset of (h, w) in x's / y's H,W strides + // scaleRedOffsets[i] / biasRedOffsets[i] = flat offset of (h, w) in scale's / bias's + // H,W strides + const std::vector xRedOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + reductionDims, xStrides.data() + reductionStart); + const std::vector yRedOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + reductionDims, yStrides.data() + reductionStart); + const std::vector scaleRedOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + reductionDims, scaleStrides.data() + reductionStart); + std::vector biasRedOffsets; + if(bias != nullptr) + { + biasRedOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + reductionDims, bias->strides().data() + reductionStart); + } + const size_t redCount = xRedOffsets.size(); + // Compute RMS-normalized output for one leading position. // - // Running example: input [N, C, H, W], scale [1, 1, H, W] - // leadingDims = [N, C] reductionDims = [H, W] - // leadingIdx = [n, c] redIdx = [h, w] - // fullIdx = [n, c, h, w] (leadingIdx ++ redIdx) - // scaleIdx = [0, 0, h, w] (zero for leading region, redIdx for the rest) - // invRmsIdx = [n, c, 0, 0] (leadingIdx padded with zeros to rank) + // Running example (continued): leadingIdx = [n, c] + // xLeadOffset / yLeadOffset = flat offset of (n, c) in x's / y's N,C strides + // x's element at redIdx i = xBase[xLeadOffset + xRedOffsets[i]] + // scale's/bias's element = scaleBase[scaleRedOffsets[i]] (no leading component) + // invRms's element = invRmsBase[flat offset of (n, c) in invRms's N,C strides] auto rmsnormFwdFunc = [&](const std::vector& leadingIdx) { + const int64_t xLeadOffset = hipdnn_test_sdk::detail::flatOffset( + leadingIdx.data(), xStrides.data(), reductionStart); + // Pass 1: accumulate sum(x^2) over the reduction dims at this leading position. // invRms[n, c, 0, 0] = 1 / sqrt(mean(x^2 over [h, w]) + epsilon) auto sumSquares = static_cast(0.0); - hipdnn_data_sdk::utilities::iterateAlongDimensions( - reductionDims, [&](const std::vector& redIdx) { - std::vector fullIdx = leadingIdx; - fullIdx.insert(fullIdx.end(), redIdx.begin(), redIdx.end()); - const auto inVal = static_cast(x.getHostValue(fullIdx)); - sumSquares += inVal * inVal; - }); + for(size_t i = 0; i < redCount; ++i) + { + const auto inVal + = static_cast(xBase[xLeadOffset + xRedOffsets[i]]); + sumSquares += inVal * inVal; + } const auto meanSquares = sumSquares / reductionCountCompute; const auto invRmsValue = static_cast(1.0) / hipdnn_data_sdk::types::sqrt(meanSquares + epsilonCompute); - // Pass 2: write y = scale * x * invRms (+ bias), iterating the same reduction - // positions but now reading scale/bias at their leading-zeroed indices. + // Pass 2: write y = scale * x * invRms (+ bias), walking the same reduction + // offsets but now also reading scale/bias at their leading-zeroed positions. // y[n, c, h, w] = scale[0, 0, h, w] * x[n, c, h, w] * invRms[n, c, 0, 0] // (+ bias[0, 0, h, w]) - hipdnn_data_sdk::utilities::iterateAlongDimensions( - reductionDims, [&](const std::vector& redIdx) { - std::vector fullIdx = leadingIdx; - fullIdx.insert(fullIdx.end(), redIdx.begin(), redIdx.end()); - - // Scale's leading-region positions are 0 (broadcast); the reduction-region - // positions match fullIdx. For the running example: scaleIdx = [0, 0, h, w]. - std::vector scaleIdx(rank, 0); - for(size_t i = reductionStart; i < rank; ++i) - { - scaleIdx[i] = fullIdx[i]; - } - - const auto xVal = static_cast(x.getHostValue(fullIdx)); - const auto xNorm = xVal * invRmsValue; - ComputeDataType yVal - = static_cast(scale.getHostValue(scaleIdx)) * xNorm; - if(bias != nullptr) - { - yVal += static_cast(bias->getHostValue(scaleIdx)); - } - y.setHostValue(static_cast(yVal), fullIdx); - }); + const int64_t yLeadOffset = hipdnn_test_sdk::detail::flatOffset( + leadingIdx.data(), yStrides.data(), reductionStart); + for(size_t i = 0; i < redCount; ++i) + { + const auto xVal = static_cast(xBase[xLeadOffset + xRedOffsets[i]]); + const auto xNorm = xVal * invRmsValue; + ComputeDataType yVal + = static_cast(scaleBase[scaleRedOffsets[i]]) * xNorm; + if(biasBase != nullptr) + { + yVal += static_cast(biasBase[biasRedOffsets[i]]); + } + yBase[yLeadOffset + yRedOffsets[i]] = static_cast(yVal); + } // invRms keeps the leading dims and collapses the reduction dims to 1, so its - // index is leadingIdx padded with zeros. Running example: invRmsIdx = [n, c, 0, 0]. - if(invRms != nullptr) + // offset has only a leading component - no reduction-walk table needed. + // Running example: invRms[n, c, 0, 0]. + if(invRmsBase != nullptr) { - std::vector invRmsIdx = leadingIdx; - invRmsIdx.resize(rank, 0); - invRms->setHostValue(static_cast(invRmsValue), invRmsIdx); + const int64_t invRmsOffset = hipdnn_test_sdk::detail::flatOffset( + leadingIdx.data(), invRmsStrides, reductionStart); + invRmsBase[invRmsOffset] = static_cast(invRmsValue); } }; @@ -197,19 +263,34 @@ class CpuFpReferenceRMSNorm if(xDims.size() < 2) { - throw std::runtime_error( - "RMSNorm backward requires at least 2D input tensor (batch and " - "at least one feature dim)."); + throw std::runtime_error(std::string(PREFIX) + + "backward requires at least 2D input tensor (batch and " + "at least one feature dim)."); } - - if(dy.dims().size() != xDims.size() || scaleDims.size() != xDims.size() - || invRms.dims().size() != xDims.size()) + if(scaleDims.size() != xDims.size()) + { + throw std::runtime_error(std::string(PREFIX) + + "backward requires scale rank to equal input rank."); + } + if(dy.dims() != xDims || dx.dims() != xDims) { - throw std::runtime_error("RMSNorm backward requires dy, scale, and invRms to all have " - "the same rank as input."); + throw std::runtime_error(std::string(PREFIX) + + "backward requires dy and dx to have the same shape as x."); + } + if(dscale.dims() != scaleDims || (dbias != nullptr && dbias->dims() != scaleDims)) + { + throw std::runtime_error( + std::string(PREFIX) + + "backward requires dscale and dbias to have the same shape as scale."); } - const size_t rank = xDims.size(); + hipdnn_test_sdk::detail::validateNoRaggedTensor(dy, PREFIX, "dy"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(x, PREFIX, "x"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(scale, PREFIX, "scale"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(invRms, PREFIX, "invRms"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(dx, PREFIX, "dx"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(dscale, PREFIX, "dscale"); + hipdnn_test_sdk::detail::validateNoRaggedTensor(dbias, PREFIX, "dbias"); // Compute leading and reduction dims based on input and scale shapes std::vector leadingDims; @@ -217,42 +298,84 @@ class CpuFpReferenceRMSNorm computeLeadingAndReductionDims(xDims, scaleDims, leadingDims, reductionDims); const auto reductionStart = leadingDims.size(); + if(invRms.dims() != invRmsShape(xDims, reductionStart)) + { + throw std::runtime_error(std::string(PREFIX) + + "backward requires invRms to have the input shape with " + "the normalized dims set to 1."); + } + const auto reductionCount = std::accumulate( reductionDims.begin(), reductionDims.end(), int64_t{1}, std::multiplies<>{}); const auto reductionCountCompute = static_cast(reductionCount); + // Tensors are addressed through hoisted base pointers and strides; see forward(). + const DyDataType* dyBase = dy.memory().hostData(); + const XDataType* xBase = x.memory().hostData(); + const ScaleDataType* scaleBase = scale.memory().hostData(); + const ComputeDataType* invRmsBase = invRms.memory().hostData(); + DxDataType* dxBase = dx.memory().hostData(); + ScaleDataType* dscaleBase = dscale.memory().hostData(); + ScaleDataType* dbiasBase = (dbias != nullptr) ? dbias->memory().hostData() : nullptr; + + const auto& dyStrides = dy.strides(); + const auto& xStrides = x.strides(); + const auto& scaleStrides = scale.strides(); + const auto& invRmsStrides = invRms.strides(); + const auto& dxStrides = dx.strides(); + const auto& dscaleStrides = dscale.strides(); + const int64_t* dbiasStrides = (dbias != nullptr) ? dbias->strides().data() : nullptr; + + // Leading-region offset tables: one flat offset per leading position, built once + // and shared by every reduction position walked in this functor. Roles are + // inverted from forward()/rmsnormBwdDataFunc below - this functor is parallelized + // over reductionDims and walks leadingDims internally, so it is the leading walk + // (not the reduction walk) that benefits from a shared table. + // Running example: input [N, C, H, W], scale [1, 1, H, W] + // leadingDims = [N, C] reductionDims = [H, W] + // dyLeadOffsets[i] / xLeadOffsets[i] / invRmsLeadOffsets[i] = flat offset of the + // i-th (n, c) position in dy's / x's / invRms's N,C strides + const std::vector dyLeadOffsets + = hipdnn_test_sdk::detail::buildDenseOffsets(leadingDims, dyStrides.data()); + const std::vector xLeadOffsets + = hipdnn_test_sdk::detail::buildDenseOffsets(leadingDims, xStrides.data()); + const std::vector invRmsLeadOffsets + = hipdnn_test_sdk::detail::buildDenseOffsets(leadingDims, invRmsStrides.data()); + const size_t leadCount = dyLeadOffsets.size(); + + // redIdx = [h, w] is fixed for the whole functor call, so its contribution to the + // dy/x/dscale/dbias offsets (dyRedOffset, xRedOffset, dscaleOffset, dbiasOffset + // below) is a scalar computed once, not a table. auto rmsnormBwdWeightBiasFunc = [&](const std::vector& redIdx) { auto sumDScale = static_cast(0.0); auto sumDBias = static_cast(0.0); - // Compute the reduction-sum components of dscale and dbias for this reduction position. - hipdnn_data_sdk::utilities::iterateAlongDimensions( - leadingDims, [&](const std::vector& leadingIdx) { - std::vector fullIdx = leadingIdx; - fullIdx.insert(fullIdx.end(), redIdx.begin(), redIdx.end()); - - std::vector invRmsIdx = leadingIdx; - invRmsIdx.resize(rank, 0); + const int64_t dyRedOffset = hipdnn_test_sdk::detail::flatOffset( + redIdx.data(), dyStrides.data() + reductionStart, reductionDims.size()); + const int64_t xRedOffset = hipdnn_test_sdk::detail::flatOffset( + redIdx.data(), xStrides.data() + reductionStart, reductionDims.size()); - const auto pdy = static_cast(dy.getHostValue(fullIdx)); - const auto px = static_cast(x.getHostValue(fullIdx)); - const auto prstd = static_cast(invRms.getHostValue(invRmsIdx)); - - sumDScale += pdy * (px * prstd); - sumDBias += pdy; - }); - - // Write dscale and dbias for this reduction position. - std::vector scaleIdx(rank, 0); - for(size_t i = reductionStart; i < rank; ++i) + // Compute the reduction-sum components of dscale and dbias for this reduction position. + for(size_t i = 0; i < leadCount; ++i) { - scaleIdx[i] = redIdx[i - reductionStart]; + const auto pdy + = static_cast(dyBase[dyLeadOffsets[i] + dyRedOffset]); + const auto px = static_cast(xBase[xLeadOffsets[i] + xRedOffset]); + const auto prstd = static_cast(invRmsBase[invRmsLeadOffsets[i]]); + + sumDScale += pdy * (px * prstd); + sumDBias += pdy; } - dscale.setHostValue(static_cast(sumDScale), scaleIdx); - if(dbias != nullptr) + // Write dscale and dbias for this reduction position. + const int64_t dscaleOffset = hipdnn_test_sdk::detail::flatOffset( + redIdx.data(), dscaleStrides.data() + reductionStart, reductionDims.size()); + dscaleBase[dscaleOffset] = static_cast(sumDScale); + if(dbiasBase != nullptr) { - dbias->setHostValue(static_cast(sumDBias), scaleIdx); + const int64_t dbiasOffset = hipdnn_test_sdk::detail::flatOffset( + redIdx.data(), dbiasStrides + reductionStart, reductionDims.size()); + dbiasBase[dbiasOffset] = static_cast(sumDBias); } }; @@ -265,57 +388,60 @@ class CpuFpReferenceRMSNorm dbias->memory().markHostModified(); } + // Reduction-region offset tables for the data-gradient pass below, built once and + // shared by every leading position (mirrors forward()'s tables). + const std::vector dyRedOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + reductionDims, dyStrides.data() + reductionStart); + const std::vector xRedOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + reductionDims, xStrides.data() + reductionStart); + const std::vector scaleRedOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + reductionDims, scaleStrides.data() + reductionStart); + const std::vector dxRedOffsets = hipdnn_test_sdk::detail::buildDenseOffsets( + reductionDims, dxStrides.data() + reductionStart); + const size_t redCount = dyRedOffsets.size(); + auto rmsnormBwdDataFunc = [&](const std::vector& leadingIdx) { auto meanDyXW = static_cast(0.0); + const int64_t dyLeadOffset = hipdnn_test_sdk::detail::flatOffset( + leadingIdx.data(), dyStrides.data(), reductionStart); + const int64_t xLeadOffset = hipdnn_test_sdk::detail::flatOffset( + leadingIdx.data(), xStrides.data(), reductionStart); + // Compute the reduction-sum component of dx = (dy * scale - meanDyXW * x) * invRms, // where meanDyXW = mean(dy * scale * x) over the reduction dims at this leading position. - hipdnn_data_sdk::utilities::iterateAlongDimensions( - reductionDims, [&](const std::vector& redIdx) { - std::vector fullIdx = leadingIdx; - fullIdx.insert(fullIdx.end(), redIdx.begin(), redIdx.end()); - - std::vector scaleIdx(rank, 0); - for(size_t i = reductionStart; i < rank; ++i) - { - scaleIdx[i] = fullIdx[i]; - } - - const auto pdy = static_cast(dy.getHostValue(fullIdx)); - const auto px = static_cast(x.getHostValue(fullIdx)); - const auto pw = static_cast(scale.getHostValue(scaleIdx)); + for(size_t i = 0; i < redCount; ++i) + { + const auto pdy + = static_cast(dyBase[dyLeadOffset + dyRedOffsets[i]]); + const auto px = static_cast(xBase[xLeadOffset + xRedOffsets[i]]); + const auto pw = static_cast(scaleBase[scaleRedOffsets[i]]); - meanDyXW += pdy * pw * px; - }); + meanDyXW += pdy * pw * px; + } meanDyXW /= reductionCountCompute; // Get invRms for this leading position - std::vector invRmsIdx = leadingIdx; - invRmsIdx.resize(rank, 0); - const auto prstd = static_cast(invRms.getHostValue(invRmsIdx)); + const int64_t invRmsOffset = hipdnn_test_sdk::detail::flatOffset( + leadingIdx.data(), invRmsStrides.data(), reductionStart); + const auto prstd = static_cast(invRmsBase[invRmsOffset]); const auto invRmsCube = prstd * prstd * prstd; // Compute dx = (dy * scale - meanDyXW * x) * invRms - hipdnn_data_sdk::utilities::iterateAlongDimensions( - reductionDims, [&](const std::vector& redIdx) { - std::vector fullIdx = leadingIdx; - fullIdx.insert(fullIdx.end(), redIdx.begin(), redIdx.end()); - - std::vector scaleIdx(rank, 0); - for(size_t i = reductionStart; i < rank; ++i) - { - scaleIdx[i] = fullIdx[i]; - } - - const auto pdy = static_cast(dy.getHostValue(fullIdx)); - const auto px = static_cast(x.getHostValue(fullIdx)); - const auto pw = static_cast(scale.getHostValue(scaleIdx)); + const int64_t dxLeadOffset = hipdnn_test_sdk::detail::flatOffset( + leadingIdx.data(), dxStrides.data(), reductionStart); + for(size_t i = 0; i < redCount; ++i) + { + const auto pdy + = static_cast(dyBase[dyLeadOffset + dyRedOffsets[i]]); + const auto px = static_cast(xBase[xLeadOffset + xRedOffsets[i]]); + const auto pw = static_cast(scaleBase[scaleRedOffsets[i]]); - const auto dxVal = (pdy * pw * prstd) - (meanDyXW * px * invRmsCube); + const auto dxVal = (pdy * pw * prstd) - (meanDyXW * px * invRmsCube); - dx.setHostValue(static_cast(dxVal), fullIdx); - }); + dxBase[dxLeadOffset + dxRedOffsets[i]] = static_cast(dxVal); + } }; auto parallelDataFunc diff --git a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/detail/CpuFpReferenceUtilities.hpp b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/detail/CpuFpReferenceUtilities.hpp index 5832c56fb66a..e88c765f4001 100644 --- a/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/detail/CpuFpReferenceUtilities.hpp +++ b/projects/hipdnn/test_sdk/include/hipdnn_test_sdk/utilities/detail/CpuFpReferenceUtilities.hpp @@ -151,17 +151,18 @@ auto callFuncUnpackArgs(F f, T args) return callFuncUnpackArgsImpl(f, args, std::make_index_sequence{}); } -template -struct ParallelTensorFunctorDynamic +/** + * @brief Row-major decomposition of a flat work index, plus the split of that work + * across threads. Shared by the parallel tensor functors below. + */ +struct ParallelTensorRange { - F func; std::vector lengths; std::vector strides; std::size_t totalElements{1}; - ParallelTensorFunctorDynamic(F f, const std::vector& dimensions) - : func(f) - , lengths(dimensions.begin(), dimensions.end()) + explicit ParallelTensorRange(const std::vector& dimensions) + : lengths(dimensions.begin(), dimensions.end()) , strides(dimensions.size()) { if(lengths.empty()) @@ -175,20 +176,24 @@ struct ParallelTensorFunctorDynamic totalElements = strides[0] * lengths[0]; } - std::vector getNdIndices(std::size_t i) const + void fillNdIndices(std::size_t i, std::vector& indices) const { - std::vector indices(lengths.size()); + indices.resize(lengths.size()); for(std::size_t idim = 0; idim < lengths.size(); ++idim) { indices[idim] = static_cast(i / strides[idim]); i -= static_cast(indices[idim]) * strides[idim]; } - - return indices; } - void operator()(std::size_t numThreads = 1) const + /** + * @brief Runs `body(workBegin, workEnd)` once per thread over disjoint work ranges. + * + * Every thread is joined before this returns, so `body` may capture by reference. + */ + template + void runChunked(std::size_t numThreads, const Body& body) const { if(totalElements == 0) { @@ -205,36 +210,344 @@ struct ParallelTensorFunctorDynamic const std::size_t workBegin = threadIdx * workPerThread; const std::size_t workEnd = std::min((threadIdx + 1) * workPerThread, totalElements); - auto threadFunc = [=, *this] { - for(std::size_t workIdx = workBegin; workIdx < workEnd; ++workIdx) + threads[threadIdx] + = JoinableThread([&body, workBegin, workEnd] { body(workBegin, workEnd); }); + } + } + + /** + * @brief Calls `func(scratch, indices)` for every position of the index space, in + * parallel. + * + * Each worker thread owns one default-constructed `Scratch` and one index buffer for + * its whole work range. fillNdIndices rewrites the buffer before every call, so the + * callee may modify it. A callee returning `bool` stops its own thread early when it + * returns false. + */ + template + void forEachIndex(std::size_t numThreads, const F& func) const + { + static_assert(std::is_default_constructible_v, + "Scratch must be default constructible; one is created per worker thread"); + + runChunked(numThreads, [this, &func](std::size_t workBegin, std::size_t workEnd) { + Scratch scratch; + std::vector indices; + + for(std::size_t workIdx = workBegin; workIdx < workEnd; ++workIdx) + { + fillNdIndices(workIdx, indices); + + // Probed with the call-site argument types - non-const lvalues. Probing an + // rvalue or a `const&` misses a callee taking `std::vector&`, which + // would then bind fine at the call below and have its bool silently dropped. + if constexpr(std::is_invocable_r_v&>) { - if constexpr(std::is_invocable_r_v>) - { - if(!func(getNdIndices(workIdx))) - { - return; - } - } - else + if(!func(scratch, indices)) { - func(getNdIndices(workIdx)); + return; } } - }; - threads[threadIdx] = JoinableThread(threadFunc); + else + { + func(scratch, indices); + } + } + }); + } +}; + +/// Scratch for callees that need none. +struct NoScratch +{ +}; + +/** + * @brief Runs a functor over every position of an index space, in parallel. + * + * The callee is invoked as `func(indices)`. A functor returning `bool` stops its own + * thread early when it returns false. + */ +template +struct ParallelTensorFunctorDynamic : ParallelTensorRange +{ + F func; + + ParallelTensorFunctorDynamic(F f, const std::vector& dimensions) + : ParallelTensorRange(dimensions) + , func(f) + { + } + + void operator()(std::size_t numThreads = 1) const + { + forEachIndex( + numThreads, [this](NoScratch&, std::vector& indices) -> decltype(auto) { + return func(indices); + }); + } +}; + +/** + * @brief Runs a functor over every position of an index space, in parallel, handing each + * worker thread its own `Scratch`. + * + * The scratch is constructed once per thread and reused for every work item that thread + * handles: the lifetime of a buffer the callee rebuilds per work item but must not + * reallocate per work item, such as a ConvolutionWindow. + * + * The callee is invoked as `func(scratch, indices)`. + */ +template +struct ParallelTensorFunctorWithScratch : ParallelTensorRange +{ + F func; + + ParallelTensorFunctorWithScratch(F f, const std::vector& dimensions) + : ParallelTensorRange(dimensions) + , func(f) + { + } + + void operator()(std::size_t numThreads = 1) const + { + forEachIndex(numThreads, func); + } +}; + +/** + * @brief The valid taps of a convolution window, as flat offset pairs. + * + * Each spatial dimension maps a window index to a source index independently, and + * independently decides that the tap has no source element - it lands in padding, or + * (for dgrad) is not stride-aligned. `build()` therefore resolves the valid taps of each + * dimension separately, and the window is their row-major cartesian product. A consumer + * that revisits the window `expand()`s the product into a list; one whose window spans a + * whole tensor walks it with `forEachTap()` instead, which holds no product. + * + * Contracts: + * - Not thread safe. One instance per concurrent user; obtain one from + * ParallelTensorFunctorWithScratch rather than sharing or making it `static`. + * - Buffers are reused across rebuilds, so a reused instance stops allocating once it + * has seen its largest window. + * - `build()` invalidates the reference returned by any prior `expand()`. + * - `expand()` and `forEachTap()` both produce the taps in row-major window order. + * Convolution accumulates in tap order, so this ordering is part of the numerical + * contract, not an implementation detail. + */ +class ConvolutionWindow +{ +public: + struct Tap + { + int64_t windowOffset; ///< flat offset into the tensor being walked + int64_t sourceOffset; ///< flat offset into the tensor being sampled + }; + + /** + * @brief Resolves the valid taps of every spatial dimension for one window position. + * + * @param nDims Number of spatial dimensions. + * @param extents Window extent per spatial dimension. + * @param windowStrides Strides of the walked tensor, spatial dimensions only. + * @param sourceStrides Strides of the sampled tensor, spatial dimensions only. + * @param mapIndex `(dim, windowIndex) -> sourceIndex`, negative when the tap + * has no source element. + */ + template + void build(std::size_t nDims, + const int64_t* extents, + const int64_t* windowStrides, + const int64_t* sourceStrides, + MapIndex&& mapIndex) + { + _dimTaps.clear(); + _dimEnds.clear(); + + for(std::size_t dim = 0; dim < nDims; ++dim) + { + for(int64_t windowIndex = 0; windowIndex < extents[dim]; ++windowIndex) + { + const int64_t sourceIndex = mapIndex(dim, windowIndex); + if(sourceIndex >= 0) + { + _dimTaps.push_back( + Tap{windowIndex * windowStrides[dim], sourceIndex * sourceStrides[dim]}); + } + } + _dimEnds.push_back(_dimTaps.size()); + } + } + + /** + * @brief The whole window as a list, one Tap per valid window position. + * + * The list grows with the window, so use it where the window is the kernel's extent. + */ + const std::vector& expand() + { + _taps.assign(1, Tap{0, 0}); + + for(std::size_t dim = 0; dim < _dimEnds.size(); ++dim) + { + // Prefixes outer, this dimension inner: keeps the product in row-major order. + _scratch.clear(); + for(const auto& prefix : _taps) + { + for(std::size_t i = dimBegin(dim); i < _dimEnds[dim]; ++i) + { + _scratch.push_back(Tap{prefix.windowOffset + _dimTaps[i].windowOffset, + prefix.sourceOffset + _dimTaps[i].sourceOffset}); + } + } + _taps.swap(_scratch); } + + return _taps; + } + + /** + * @brief Calls `visit(tap)` for every tap, in the same order as `expand()`. + * + * Holds one cursor per dimension rather than the product, so memory stays bounded by + * the sum of the extents where the window spans a whole tensor. + */ + template + void forEachTap(Visit&& visit) + { + const std::size_t nDims = _dimEnds.size(); + if(nDims == 0) + { + visit(Tap{0, 0}); + return; + } + for(std::size_t dim = 0; dim < nDims; ++dim) + { + if(dimBegin(dim) == _dimEnds[dim]) + { + return; + } + } + + // An odometer over the outer dimensions; the innermost is walked directly. + const std::size_t inner = nDims - 1; + _cursor.resize(inner); + for(std::size_t dim = 0; dim < inner; ++dim) + { + _cursor[dim] = dimBegin(dim); + } + + do + { + Tap prefix{0, 0}; + for(std::size_t dim = 0; dim < inner; ++dim) + { + prefix.windowOffset += _dimTaps[_cursor[dim]].windowOffset; + prefix.sourceOffset += _dimTaps[_cursor[dim]].sourceOffset; + } + + for(std::size_t i = dimBegin(inner); i < _dimEnds[inner]; ++i) + { + visit(Tap{prefix.windowOffset + _dimTaps[i].windowOffset, + prefix.sourceOffset + _dimTaps[i].sourceOffset}); + } + } while(advanceCursor()); + } + +private: + std::size_t dimBegin(std::size_t dim) const + { + return dim == 0 ? 0 : _dimEnds[dim - 1]; + } + + // Steps the outer-dimension odometer, last dimension fastest. Returns false once + // every combination has been visited. + bool advanceCursor() + { + for(std::size_t dim = _cursor.size(); dim-- > 0;) + { + if(++_cursor[dim] < _dimEnds[dim]) + { + return true; + } + _cursor[dim] = dimBegin(dim); + } + + return false; } + + std::vector _dimTaps; ///< valid taps of every dimension, dimension-major + std::vector _dimEnds; ///< end of each dimension's run in _dimTaps + std::vector _cursor; ///< forEachTap position in each outer dimension + std::vector _taps; + std::vector _scratch; }; +/** + * @brief Row-major flat offsets for a dense walk of `extents` against `strides`. + * + * Entry i is the offset of the i-th row-major position of the walk. `strides` must have + * `extents.size()` entries. A zero stride is meaningful, and is how a broadcast axis + * contributes nothing to the address. + */ +inline std::vector buildDenseOffsets(const std::vector& extents, + const int64_t* strides) +{ + std::vector offsets{0}; + std::vector scratch; + + for(std::size_t dim = 0; dim < extents.size(); ++dim) + { + scratch.clear(); + scratch.reserve(offsets.size() * static_cast(extents[dim])); + + // Prefixes outer, this dimension inner: keeps the product in row-major order. + for(const auto prefix : offsets) + { + for(int64_t index = 0; index < extents[dim]; ++index) + { + scratch.push_back(prefix + (index * strides[dim])); + } + } + + offsets.swap(scratch); + } + + return offsets; +} + +/// Flat offset of the first `count` entries of `indices` against `strides`. +inline int64_t flatOffset(const int64_t* indices, const int64_t* strides, std::size_t count) +{ + int64_t offset = 0; + for(std::size_t dim = 0; dim < count; ++dim) + { + offset += indices[dim] * strides[dim]; + } + + return offset; +} + template auto makeParallelTensorFunctor(F f, const std::vector& dimensions) { return ParallelTensorFunctorDynamic(f, dimensions); } +/// Companion to makeParallelTensorFunctor for callees that need per-thread scratch. +/// `Scratch` is explicit; `F` is deduced. The callee takes `(Scratch&, indices)`. +template +auto makeParallelTensorFunctorWithScratch(F f, const std::vector& dimensions) +{ + return ParallelTensorFunctorWithScratch(f, dimensions); +} + /** * @brief Reject a ragged tensor with a message identifying which argument it was. * + * The CPU references address memory through hoisted base pointers and strides, which + * is the dense layout only; a ragged tensor rebases every batch at its own offset. + * * @param tensor The tensor to check. * @param errorPrefix Prefix identifying the calling CPU reference (e.g. "MyOp: "). * @param name The argument name to report if the tensor is ragged. @@ -249,4 +562,15 @@ inline void validateNoRaggedTensor(const hipdnn_data_sdk::utilities::ITensor& te } } +/// Overload for an optional argument; a null tensor passes. +inline void validateNoRaggedTensor(const hipdnn_data_sdk::utilities::ITensor* tensor, + const std::string& errorPrefix, + const char* name) +{ + if(tensor != nullptr) + { + validateNoRaggedTensor(*tensor, errorPrefix, name); + } +} + } // namespace hipdnn_test_sdk::detail diff --git a/projects/hipdnn/test_sdk/tests/CMakeLists.txt b/projects/hipdnn/test_sdk/tests/CMakeLists.txt index a9f3a3dd08f2..9f83ee3c520f 100644 --- a/projects/hipdnn/test_sdk/tests/CMakeLists.txt +++ b/projects/hipdnn/test_sdk/tests/CMakeLists.txt @@ -30,6 +30,7 @@ add_executable( utilities/TestCpuFpReferenceMoeGroupedMatmul.cpp utilities/TestCpuFpReferenceMoeGroupedMatmulBwd.cpp utilities/TestCpuFpReferenceReduction.cpp + utilities/TestCpuFpReferenceInputValidation.cpp utilities/TestCpuFpReferenceMiopenRmsValidation.cpp utilities/TestCpuFpReferenceUtilities.cpp utilities/TestCpuFpReferenceValidation.cpp diff --git a/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceInputValidation.cpp b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceInputValidation.cpp new file mode 100644 index 000000000000..55b0b6cac6e3 --- /dev/null +++ b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceInputValidation.cpp @@ -0,0 +1,484 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +// The convolution and normalization references address memory through hoisted base +// pointers and strides, so every tensor they touch must be dense and must have the shape +// its offsets are computed against. A ragged tensor rebases each batch at +// ragged_offset[b] inside getIndexImpl, which flat stride arithmetic never calls; a +// mis-shaped tensor is read or written out of bounds. Both are rejected up front. +// +// The CPU graph executor already validates graphs one level up (see +// cpu_graph_executor/TestCpuReferenceRaggedRejection.cpp). These tests cover the direct +// callers the executor does not sit in front of - the samples and the gpu-ref fixtures. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +using namespace hipdnn_test_sdk::utilities; +using namespace hipdnn_data_sdk::utilities; + +namespace +{ + +// Canonical BSHD-packed ragged geometry, mirroring the data SDK's ragged tests: +// dims [B, S_max, H, D] with batch 0 holding 2 sequence rows and batch 1 holding 3. +// Read as NCHW this is a legal convolution/normalization input, so the references get +// past their shape checks and reach the ragged guard under test. +const std::vector RAGGED_DIMS = {2, 3, 2, 2}; +const std::vector RAGGED_STRIDES = {12, 4, 2, 1}; +const std::vector RAGGED_OFFSETS = {0, 8, 20}; +constexpr int RAGGED_SEQ_AXIS = 1; + +std::shared_ptr makeOffsetAux() +{ + auto aux = std::make_shared>( + std::vector{static_cast(RAGGED_OFFSETS.size()), 1, 1, 1}); + for(size_t i = 0; i < RAGGED_OFFSETS.size(); ++i) + { + aux->setHostValue( + static_cast(RAGGED_OFFSETS[i]), static_cast(i), 0, 0, 0); + } + return aux; +} + +RaggedTensor makeRaggedTensor(const std::shared_ptr& aux) +{ + return {RAGGED_DIMS, RAGGED_STRIDES, RAGGED_SEQ_AXIS, aux}; +} + +// Asserts the call was rejected for the reason under test rather than by some unrelated +// check that happens to throw from the same function. +template +void expectRejectedWith(Callable&& call, const std::string& needle) +{ + try + { + call(); + ADD_FAILURE() << "expected a rejection mentioning \"" << needle << "\""; + } + catch(const std::runtime_error& error) + { + const std::string message = error.what(); + EXPECT_NE(message.find(needle), std::string::npos) + << "rejected, but not for the expected reason: " << message; + } +} + +} // namespace + +// ============================================================================ +// Ragged rejection +// ============================================================================ + +TEST(TestCpuFpReferenceRaggedRejection, ConvolutionFpropRejectsRaggedInput) +{ + auto aux = makeOffsetAux(); + auto x = makeRaggedTensor(aux); + Tensor w({1, 3, 1, 1}); + Tensor y({2, 1, 2, 2}); + + expectRejectedWith( + [&] { + CpuFpReferenceConvolution::fprop( + x, w, y, {1, 1}, {1, 1}, {0, 0}); + }, + "ragged x tensor"); +} + +TEST(TestCpuFpReferenceRaggedRejection, ConvolutionDgradRejectsRaggedGradient) +{ + auto aux = makeOffsetAux(); + auto gradX = makeRaggedTensor(aux); + Tensor w({1, 3, 1, 1}); + Tensor gradY({2, 1, 2, 2}); + + expectRejectedWith( + [&] { + CpuFpReferenceConvolution::dgrad( + gradX, w, gradY, {1, 1}, {1, 1}, {0, 0}); + }, + "ragged dx tensor"); +} + +TEST(TestCpuFpReferenceRaggedRejection, ConvolutionWgradRejectsRaggedInput) +{ + auto aux = makeOffsetAux(); + auto x = makeRaggedTensor(aux); + Tensor gradW({1, 3, 1, 1}); + Tensor gradY({2, 1, 2, 2}); + + expectRejectedWith( + [&] { + CpuFpReferenceConvolution::wgrad( + x, gradW, gradY, {1, 1}, {1, 1}, {0, 0}); + }, + "ragged x tensor"); +} + +TEST(TestCpuFpReferenceRaggedRejection, ConvolutionFpropRejectsRaggedOutput) +{ + // The output is written through the same hoisted pointer, so it needs the guard too. + auto aux = makeOffsetAux(); + Tensor x({2, 3, 2, 2}); + Tensor w({3, 3, 1, 1}); + auto y = makeRaggedTensor(aux); + + expectRejectedWith( + [&] { + CpuFpReferenceConvolution::fprop( + x, w, y, {1, 1}, {1, 1}, {0, 0}); + }, + "ragged y tensor"); +} + +TEST(TestCpuFpReferenceRaggedRejection, LayernormFpropRejectsRaggedInput) +{ + auto aux = makeOffsetAux(); + auto x = makeRaggedTensor(aux); + Tensor y(RAGGED_DIMS); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::fprop( + x, nullptr, nullptr, y, 1e-5, 1); + }, + "ragged x tensor"); +} + +TEST(TestCpuFpReferenceRaggedRejection, LayernormBpropRejectsRaggedGradient) +{ + auto aux = makeOffsetAux(); + auto dx = makeRaggedTensor(aux); + Tensor dy(RAGGED_DIMS); + Tensor x(RAGGED_DIMS); + Tensor scale({1, 1, 1, 2}); + Tensor dscale({1, 1, 1, 2}); + Tensor dbias({1, 1, 1, 2}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::bprop( + dy, x, scale, dx, dscale, dbias, 1e-5, nullptr, nullptr, 1); + }, + "ragged dx tensor"); +} + +TEST(TestCpuFpReferenceRaggedRejection, RmsNormForwardRejectsRaggedInput) +{ + auto aux = makeOffsetAux(); + auto x = makeRaggedTensor(aux); + Tensor scale({1, 1, 2, 2}); + Tensor y(RAGGED_DIMS); + + expectRejectedWith( + [&] { CpuFpReferenceRMSNorm::forward(x, scale, y, 1e-5); }, + "ragged x tensor"); +} + +TEST(TestCpuFpReferenceRaggedRejection, RmsNormBackwardRejectsRaggedGradient) +{ + auto aux = makeOffsetAux(); + auto dx = makeRaggedTensor(aux); + Tensor dy(RAGGED_DIMS); + Tensor x(RAGGED_DIMS); + Tensor scale({1, 1, 2, 2}); + Tensor invRms({2, 3, 1, 1}); + Tensor dscale({1, 1, 2, 2}); + + expectRejectedWith( + [&] { + CpuFpReferenceRMSNorm::backward( + dy, x, scale, invRms, dx, dscale); + }, + "ragged dx tensor"); +} + +// ============================================================================ +// Shape validation: layernorm +// ============================================================================ + +TEST(TestCpuFpReferenceShapeValidation, LayernormFpropRejectsOutputNotShapedLikeInput) +{ + Tensor x({2, 3, 4}); + Tensor y({2, 3, 5}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::fprop( + x, nullptr, nullptr, y, 1e-5, 1); + }, + "y to have the same shape as x"); +} + +TEST(TestCpuFpReferenceShapeValidation, LayernormFpropRejectsBiasNotShapedLikeScale) +{ + // bias is walked with scale's offsets, so a shorter bias is read past its end. + Tensor x({2, 4}); + Tensor scale({4}); + Tensor bias({3}); + Tensor y({2, 4}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::fprop( + x, &scale, &bias, y, 1e-5, 1); + }, + "scale and bias to have the same shape"); +} + +TEST(TestCpuFpReferenceShapeValidation, LayernormFpropRejectsRstdNotShapedLikeMean) +{ + // The batch walk is taken from mean, so a lower-rank rstd is written through strides + // it does not have. + Tensor x({2, 3, 4}); + Tensor y({2, 3, 4}); + Tensor mean({2, 3}); + Tensor rstd({2}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::fprop( + x, nullptr, nullptr, y, 1e-5, 1, &mean, &rstd); + }, + "mean and rstd to have the same shape"); +} + +TEST(TestCpuFpReferenceShapeValidation, LayernormFpropRejectsScaleRankBelowNormalizedDimCount) +{ + Tensor x({2, 3, 4}); + Tensor scale({4}); + Tensor bias({4}); + Tensor y({2, 3, 4}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::fprop( + x, &scale, &bias, y, 1e-5, 2); + }, + "trailing normalizedDimCount dims"); +} + +TEST(TestCpuFpReferenceShapeValidation, LayernormFpropRejectsScaleNotMatchingNormalizedDims) +{ + Tensor x({2, 4}); + Tensor scale({3}); + Tensor y({2, 4}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::fprop( + x, &scale, nullptr, y, 1e-5, 1); + }, + "trailing normalizedDimCount dims"); +} + +TEST(TestCpuFpReferenceShapeValidation, LayernormFpropRejectsStatsNotMatchingBatchDims) +{ + Tensor x({2, 3, 4}); + Tensor y({2, 3, 4}); + Tensor mean({2, 4}); + Tensor rstd({2, 4}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::fprop( + x, nullptr, nullptr, y, 1e-5, 1, &mean, &rstd); + }, + "leading batch dims"); +} + +TEST(TestCpuFpReferenceShapeValidation, LayernormFpropRejectsRankZeroStats) +{ + // A rank-0 tensor holds no element, so a whole-tensor normalization still needs a + // one-element mean/rstd to write its single batch position to. + Tensor x({4}); + Tensor scale({4}); + Tensor bias({4}); + Tensor y({4}); + Tensor mean({}); + Tensor rstd({}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::fprop( + x, &scale, &bias, y, 1e-5, 1, &mean, &rstd); + }, + "at least one dimension"); +} + +TEST(TestCpuFpReferenceShapeValidation, LayernormBpropRejectsInputNotShapedLikeGradient) +{ + Tensor dy({2, 3, 4}); + Tensor x({2, 3, 4}); + Tensor dx({2, 4, 4}); + Tensor scale({4}); + Tensor dscale({4}); + Tensor dbias({4}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::bprop( + dy, x, scale, dx, dscale, dbias, 1e-5, nullptr, nullptr, 1); + }, + "x and dx to have the same shape as dy"); +} + +TEST(TestCpuFpReferenceShapeValidation, LayernormBpropRejectsScaleRankBelowNormalizedDimCount) +{ + Tensor dy({2, 3, 4}); + Tensor x({2, 3, 4}); + Tensor dx({2, 3, 4}); + Tensor scale({4}); + Tensor dscale({4}); + Tensor dbias({4}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::bprop( + dy, x, scale, dx, dscale, dbias, 1e-5, nullptr, nullptr, 2); + }, + "trailing normalizedDimCount dims"); +} + +TEST(TestCpuFpReferenceShapeValidation, LayernormBpropRejectsRankZeroStats) +{ + Tensor dy({4}); + Tensor x({4}); + Tensor dx({4}); + Tensor scale({4}); + Tensor dscale({4}); + Tensor dbias({4}); + Tensor mean({}); + Tensor rstd({}); + + expectRejectedWith( + [&] { + CpuFpReferenceLayernorm::bprop( + dy, x, scale, dx, dscale, dbias, 1e-5, &mean, &rstd, 1); + }, + "at least one dimension"); +} + +// ============================================================================ +// Shape validation: RMSNorm +// ============================================================================ + +TEST(TestCpuFpReferenceShapeValidation, RmsNormForwardRejectsOutputNotShapedLikeInput) +{ + Tensor x({2, 3, 4}); + Tensor scale({1, 3, 4}); + Tensor y({2, 3, 5}); + + expectRejectedWith( + [&] { CpuFpReferenceRMSNorm::forward(x, scale, y, 1e-5); }, + "y to have the same shape as x"); +} + +TEST(TestCpuFpReferenceShapeValidation, RmsNormForwardRejectsBiasNotShapedLikeScale) +{ + Tensor x({2, 3, 4}); + Tensor scale({1, 3, 4}); + Tensor y({2, 3, 4}); + Tensor bias({1, 1, 4}); + Tensor* noInvRms = nullptr; + + expectRejectedWith( + [&] { + CpuFpReferenceRMSNorm::forward( + x, scale, y, 1e-5, noInvRms, &bias); + }, + "bias to have the same shape as scale"); +} + +TEST(TestCpuFpReferenceShapeValidation, RmsNormForwardRejectsInvRmsNotCollapsingNormalizedDims) +{ + // scale [1, 1, 4] normalizes over the last dim only, so invRms is [2, 3, 1]; the + // [2, 1, 1] of a wider normalization is written past its end. + Tensor x({2, 3, 4}); + Tensor scale({1, 1, 4}); + Tensor y({2, 3, 4}); + Tensor invRms({2, 1, 1}); + + expectRejectedWith( + [&] { + CpuFpReferenceRMSNorm::forward(x, scale, y, 1e-5, &invRms); + }, + "invRms to have the input shape"); +} + +TEST(TestCpuFpReferenceShapeValidation, RmsNormBackwardRejectsGradientNotShapedLikeInput) +{ + Tensor dy({2, 3, 4}); + Tensor x({2, 3, 4}); + Tensor dx({2, 3, 5}); + Tensor scale({1, 3, 4}); + Tensor invRms({2, 1, 1}); + Tensor dscale({1, 3, 4}); + + expectRejectedWith( + [&] { + CpuFpReferenceRMSNorm::backward( + dy, x, scale, invRms, dx, dscale); + }, + "dy and dx to have the same shape as x"); +} + +TEST(TestCpuFpReferenceShapeValidation, RmsNormBackwardRejectsWeightGradientNotShapedLikeScale) +{ + Tensor dy({2, 3, 4}); + Tensor x({2, 3, 4}); + Tensor dx({2, 3, 4}); + Tensor scale({1, 3, 4}); + Tensor invRms({2, 1, 1}); + Tensor dscale({3, 4}); + + expectRejectedWith( + [&] { + CpuFpReferenceRMSNorm::backward( + dy, x, scale, invRms, dx, dscale); + }, + "dscale and dbias to have the same shape as scale"); +} + +TEST(TestCpuFpReferenceShapeValidation, RmsNormBackwardRejectsBiasGradientNotShapedLikeScale) +{ + Tensor dy({2, 3, 4}); + Tensor x({2, 3, 4}); + Tensor dx({2, 3, 4}); + Tensor scale({1, 3, 4}); + Tensor invRms({2, 1, 1}); + Tensor dscale({1, 3, 4}); + Tensor dbias({1, 1, 4}); + + expectRejectedWith( + [&] { + CpuFpReferenceRMSNorm::backward( + dy, x, scale, invRms, dx, dscale, &dbias); + }, + "dscale and dbias to have the same shape as scale"); +} + +TEST(TestCpuFpReferenceShapeValidation, RmsNormBackwardRejectsInvRmsNotCollapsingNormalizedDims) +{ + Tensor dy({2, 3, 4}); + Tensor x({2, 3, 4}); + Tensor dx({2, 3, 4}); + Tensor scale({1, 1, 4}); + Tensor invRms({2, 1, 1}); + Tensor dscale({1, 1, 4}); + + expectRejectedWith( + [&] { + CpuFpReferenceRMSNorm::backward( + dy, x, scale, invRms, dx, dscale); + }, + "invRms to have the input shape"); +} diff --git a/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceLayernorm.cpp b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceLayernorm.cpp index 1c2eb08ffe36..4ba12a254fec 100644 --- a/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceLayernorm.cpp +++ b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceLayernorm.cpp @@ -1292,7 +1292,7 @@ TEST(TestCpuFpReferenceLayernormFp64, Fprop5DNormalizeLast2) } // ============================================================================ -// buildFullIndices extra tests +// Scale/bias and mean/rstd padded with 1s to the input rank // ============================================================================ TEST(TestCpuFpReferenceLayernormFp64, Fprop5DNormalizeLast2OnePadded) diff --git a/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceLayernormBackward.cpp b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceLayernormBackward.cpp index 2b05bef1e479..f9998b104ef3 100644 --- a/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceLayernormBackward.cpp +++ b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceLayernormBackward.cpp @@ -1732,3 +1732,91 @@ TEST(TestCpuFpReferenceLayernormBackwardFp64, Bprop5DNormalizeLast2) } } } + +// ============================================================================ +// Whole-tensor normalization with recomputed statistics +// ============================================================================ + +// The intersection of two paths that are each covered separately elsewhere: a +// whole-tensor normalization (every dimension is normalized, so the batch walk is a single +// position indexed by no dimension) and absent mean/rstd (so bprop recomputes the +// statistics with Welford and carries them from pass 1 to pass 2 itself). +// +// The recomputed statistics have to agree with supplied ones, so the two paths are run +// against each other rather than against hand-computed values. +TEST(TestCpuFpReferenceLayernormBackwardFp32, BpropWholeTensorNormalizationRecomputesStats) +{ + // Rank-1 input with a rank-2 scale: normalizedDimCount covers every input dimension, + // which leaves no batch dimension. + const std::vector ioDims{4}; + const std::vector paramDims{1, 4}; + constexpr int64_t NORMALIZED_DIM_COUNT = 1; + constexpr double EPSILON = 1e-5; + + Tensor x(ioDims); + Tensor dy(ioDims); + Tensor scale(paramDims); + Tensor bias(paramDims); + + const std::vector xValues{1.0f, 2.0f, 4.0f, 8.0f}; + const std::vector dyValues{0.5f, -1.5f, 2.0f, 0.25f}; + const std::vector scaleValues{1.5f, -0.5f, 2.0f, 0.75f}; + for(int64_t i = 0; i < ioDims[0]; ++i) + { + x.setHostValue(xValues[static_cast(i)], i); + dy.setHostValue(dyValues[static_cast(i)], i); + scale.setHostValue(scaleValues[static_cast(i)], 0, i); + bias.setHostValue(0.0f, 0, i); + } + + // Produce the statistics the recompute path has to reproduce. + Tensor y(ioDims); + Tensor mean({1}); + Tensor rstd({1}); + CpuFpReferenceLayernorm::fprop( + x, &scale, &bias, y, EPSILON, NORMALIZED_DIM_COUNT, &mean, &rstd); + + Tensor dxSupplied(ioDims); + Tensor dscaleSupplied(paramDims); + Tensor dbiasSupplied(paramDims); + CpuFpReferenceLayernorm::bprop(dy, + x, + scale, + dxSupplied, + dscaleSupplied, + dbiasSupplied, + EPSILON, + &mean, + &rstd, + NORMALIZED_DIM_COUNT); + + // The path under test: no statistics, so bprop recomputes them for an empty batch + // index space. + Tensor dxRecomputed(ioDims); + Tensor dscaleRecomputed(paramDims); + Tensor dbiasRecomputed(paramDims); + CpuFpReferenceLayernorm::bprop(dy, + x, + scale, + dxRecomputed, + dscaleRecomputed, + dbiasRecomputed, + EPSILON, + nullptr, + nullptr, + NORMALIZED_DIM_COUNT); + + for(int64_t i = 0; i < ioDims[0]; ++i) + { + EXPECT_FLOAT_EQ(dxRecomputed.getHostValue(i), dxSupplied.getHostValue(i)) + << "dx element " << i; + EXPECT_FLOAT_EQ(dscaleRecomputed.getHostValue(0, i), dscaleSupplied.getHostValue(0, i)) + << "dscale element " << i; + EXPECT_FLOAT_EQ(dbiasRecomputed.getHostValue(0, i), dbiasSupplied.getHostValue(0, i)) + << "dbias element " << i; + } + + // Guard against both paths being uniformly zero, which would satisfy the comparison + // above without exercising anything. + EXPECT_NE(dxRecomputed.getHostValue(0), 0.0f); +} diff --git a/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceUtilities.cpp b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceUtilities.cpp index 4f7e029ef41b..aa004e7e25be 100644 --- a/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceUtilities.cpp +++ b/projects/hipdnn/test_sdk/tests/utilities/TestCpuFpReferenceUtilities.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -33,15 +34,18 @@ TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamic1DIndexCalculati auto functor = makeParallelTensorFunctor( [](const std::vector& indices) { (void)indices; }, std::vector{10}); - auto indices0 = functor.getNdIndices(0); + std::vector indices0; + functor.fillNdIndices(0, indices0); EXPECT_EQ(indices0.size(), 1); EXPECT_EQ(indices0[0], 0); - auto indices5 = functor.getNdIndices(5); + std::vector indices5; + functor.fillNdIndices(5, indices5); EXPECT_EQ(indices5.size(), 1); EXPECT_EQ(indices5[0], 5); - auto indices9 = functor.getNdIndices(9); + std::vector indices9; + functor.fillNdIndices(9, indices9); EXPECT_EQ(indices9.size(), 1); EXPECT_EQ(indices9[0], 9); } @@ -51,24 +55,29 @@ TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamic2DIndexCalculati auto functor = makeParallelTensorFunctor( [](const std::vector& indices) { (void)indices; }, std::vector{3, 4}); - auto indices0 = functor.getNdIndices(0); // Should be (0, 0) + std::vector indices0; + functor.fillNdIndices(0, indices0); // Should be (0, 0) EXPECT_EQ(indices0.size(), 2); EXPECT_EQ(indices0[0], 0); EXPECT_EQ(indices0[1], 0); - auto indices3 = functor.getNdIndices(3); // Should be (0, 3) + std::vector indices3; + functor.fillNdIndices(3, indices3); // Should be (0, 3) EXPECT_EQ(indices3[0], 0); EXPECT_EQ(indices3[1], 3); - auto indices4 = functor.getNdIndices(4); // Should be (1, 0) + std::vector indices4; + functor.fillNdIndices(4, indices4); // Should be (1, 0) EXPECT_EQ(indices4[0], 1); EXPECT_EQ(indices4[1], 0); - auto indices7 = functor.getNdIndices(7); // Should be (1, 3) + std::vector indices7; + functor.fillNdIndices(7, indices7); // Should be (1, 3) EXPECT_EQ(indices7[0], 1); EXPECT_EQ(indices7[1], 3); - auto indices11 = functor.getNdIndices(11); // Should be (2, 3) + std::vector indices11; + functor.fillNdIndices(11, indices11); // Should be (2, 3) EXPECT_EQ(indices11[0], 2); EXPECT_EQ(indices11[1], 3); } @@ -78,18 +87,21 @@ TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamic3DIndexCalculati auto functor = makeParallelTensorFunctor( [](const std::vector& indices) { (void)indices; }, std::vector{2, 3, 4}); - auto indices0 = functor.getNdIndices(0); // Should be (0, 0, 0) + std::vector indices0; + functor.fillNdIndices(0, indices0); // Should be (0, 0, 0) EXPECT_EQ(indices0.size(), 3); EXPECT_EQ(indices0[0], 0); EXPECT_EQ(indices0[1], 0); EXPECT_EQ(indices0[2], 0); - auto indices12 = functor.getNdIndices(12); // Should be (1, 0, 0) + std::vector indices12; + functor.fillNdIndices(12, indices12); // Should be (1, 0, 0) EXPECT_EQ(indices12[0], 1); EXPECT_EQ(indices12[1], 0); EXPECT_EQ(indices12[2], 0); - auto indices23 = functor.getNdIndices(23); // Should be (1, 2, 3) + std::vector indices23; + functor.fillNdIndices(23, indices23); // Should be (1, 2, 3) EXPECT_EQ(indices23[0], 1); EXPECT_EQ(indices23[1], 2); EXPECT_EQ(indices23[2], 3); @@ -101,26 +113,43 @@ TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamic4DIndexCalculati = makeParallelTensorFunctor([](const std::vector& indices) { (void)indices; }, std::vector{2, 2, 2, 2}); - auto indices0 = functor.getNdIndices(0); // Should be (0, 0, 0, 0) + std::vector indices0; + functor.fillNdIndices(0, indices0); // Should be (0, 0, 0, 0) EXPECT_EQ(indices0.size(), 4); EXPECT_EQ(indices0[0], 0); EXPECT_EQ(indices0[1], 0); EXPECT_EQ(indices0[2], 0); EXPECT_EQ(indices0[3], 0); - auto indices8 = functor.getNdIndices(8); // Should be (1, 0, 0, 0) + std::vector indices8; + functor.fillNdIndices(8, indices8); // Should be (1, 0, 0, 0) EXPECT_EQ(indices8[0], 1); EXPECT_EQ(indices8[1], 0); EXPECT_EQ(indices8[2], 0); EXPECT_EQ(indices8[3], 0); - auto indices15 = functor.getNdIndices(15); // Should be (1, 1, 1, 1) + std::vector indices15; + functor.fillNdIndices(15, indices15); // Should be (1, 1, 1, 1) EXPECT_EQ(indices15[0], 1); EXPECT_EQ(indices15[1], 1); EXPECT_EQ(indices15[2], 1); EXPECT_EQ(indices15[3], 1); } +TEST_F(TestCpuFpReferenceUtilities, FillNdIndicesOverwritesAReusedBuffer) +{ + // Each worker refills one buffer for its whole range, so a fill must size the buffer + // to the rank and overwrite every entry, whatever an earlier item or callee left in it. + const ParallelTensorRange range(std::vector{3, 4}); + + std::vector indices{9, 9, 9, 9, 9}; + range.fillNdIndices(7, indices); + EXPECT_EQ(indices, (std::vector{1, 3})); + + range.fillNdIndices(4, indices); + EXPECT_EQ(indices, (std::vector{1, 0})); +} + TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamicSingleThreadExecution) { std::atomic sum{0}; @@ -326,7 +355,8 @@ TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamic5DIndexCalculati = makeParallelTensorFunctor([](const std::vector& indices) { (void)indices; }, std::vector{2, 2, 2, 2, 2}); - auto indices0 = functor.getNdIndices(0); // Should be (0, 0, 0, 0, 0) + std::vector indices0; + functor.fillNdIndices(0, indices0); // Should be (0, 0, 0, 0, 0) EXPECT_EQ(indices0.size(), 5); EXPECT_EQ(indices0[0], 0); EXPECT_EQ(indices0[1], 0); @@ -334,14 +364,16 @@ TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamic5DIndexCalculati EXPECT_EQ(indices0[3], 0); EXPECT_EQ(indices0[4], 0); - auto indices16 = functor.getNdIndices(16); // Should be (1, 0, 0, 0, 0) + std::vector indices16; + functor.fillNdIndices(16, indices16); // Should be (1, 0, 0, 0, 0) EXPECT_EQ(indices16[0], 1); EXPECT_EQ(indices16[1], 0); EXPECT_EQ(indices16[2], 0); EXPECT_EQ(indices16[3], 0); EXPECT_EQ(indices16[4], 0); - auto indices31 = functor.getNdIndices(31); // Should be (1, 1, 1, 1, 1) + std::vector indices31; + functor.fillNdIndices(31, indices31); // Should be (1, 1, 1, 1, 1) EXPECT_EQ(indices31[0], 1); EXPECT_EQ(indices31[1], 1); EXPECT_EQ(indices31[2], 1); @@ -418,13 +450,15 @@ TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamicSingleElementDim EXPECT_EQ(count.load(), 15); // 1 * 5 * 1 * 3 = 15 // Verify index calculation for this shape - auto indices0 = functor.getNdIndices(0); // Should be (0, 0, 0, 0) + std::vector indices0; + functor.fillNdIndices(0, indices0); // Should be (0, 0, 0, 0) EXPECT_EQ(indices0[0], 0); EXPECT_EQ(indices0[1], 0); EXPECT_EQ(indices0[2], 0); EXPECT_EQ(indices0[3], 0); - auto indices7 = functor.getNdIndices(7); // Should be (0, 2, 0, 1) + std::vector indices7; + functor.fillNdIndices(7, indices7); // Should be (0, 2, 0, 1) EXPECT_EQ(indices7[0], 0); EXPECT_EQ(indices7[1], 2); EXPECT_EQ(indices7[2], 0); @@ -457,3 +491,409 @@ TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamicThreadSafety) << elementCounts[i].load() << " times instead of 1"; } } + +// ============================================================================ +// buildDenseOffsets / flatOffset +// ============================================================================ + +TEST_F(TestCpuFpReferenceUtilities, BuildDenseOffsetsIsRowMajor) +{ + // The normalization references accumulate in walk order, so the order is part of + // their numerical contract: last dimension varies fastest. + const std::vector strides{10, 1}; + const auto offsets = buildDenseOffsets({2, 3}, strides.data()); + + EXPECT_EQ(offsets, (std::vector{0, 1, 2, 10, 11, 12})); +} + +TEST_F(TestCpuFpReferenceUtilities, BuildDenseOffsetsHonoursNonPackedStrides) +{ + // A permuted (channels-last) stride set must yield the addresses the strides + // describe, not the packed ones the extents imply. + const std::vector strides{1, 4}; + const auto offsets = buildDenseOffsets({2, 3}, strides.data()); + + EXPECT_EQ(offsets, (std::vector{0, 4, 8, 1, 5, 9})); +} + +TEST_F(TestCpuFpReferenceUtilities, BuildDenseOffsetsTreatsZeroStrideAsBroadcast) +{ + // A zero stride is how an axis the walk does not address contributes nothing: the + // walk still visits every position, but they alias onto the same offsets. + const std::vector strides{0, 1}; + const auto offsets = buildDenseOffsets({3, 2}, strides.data()); + + EXPECT_EQ(offsets, (std::vector{0, 1, 0, 1, 0, 1})); +} + +TEST_F(TestCpuFpReferenceUtilities, BuildDenseOffsetsOfNoExtentsIsOnePosition) +{ + // The scalar walk is one position at offset 0, not an empty table - a reference + // looping over an empty table would silently skip its only element. + const std::vector strides{}; + const auto offsets = buildDenseOffsets({}, strides.data()); + + EXPECT_EQ(offsets, (std::vector{0})); +} + +TEST_F(TestCpuFpReferenceUtilities, FlatOffsetReducesOnlyTheLeadingIndices) +{ + // Each normalization pass holds one index space fixed and walks the other, so it + // reduces a prefix of the index vector, never all of it. + const std::vector indices{2, 3, 4}; + const std::vector strides{100, 10, 1}; + + EXPECT_EQ(flatOffset(indices.data(), strides.data(), 0), 0); + EXPECT_EQ(flatOffset(indices.data(), strides.data(), 1), 200); + EXPECT_EQ(flatOffset(indices.data(), strides.data(), 2), 230); + EXPECT_EQ(flatOffset(indices.data(), strides.data(), 3), 234); +} + +// ============================================================================ +// ConvolutionWindow +// ============================================================================ + +TEST_F(TestCpuFpReferenceUtilities, ConvolutionWindowEmitsTapsInRowMajorWindowOrder) +{ + // Convolution accumulates in tap order, so this ordering is part of its numerical + // contract: last window dimension varies fastest. + const std::vector extents{2, 3}; + const std::vector windowStrides{3, 1}; + const std::vector sourceStrides{100, 10}; + + ConvolutionWindow window; + window.build(extents.size(), + extents.data(), + windowStrides.data(), + sourceStrides.data(), + [](size_t, int64_t index) { return index; }); + + const std::vector expectedWindow{0, 1, 2, 3, 4, 5}; + const std::vector expectedSource{0, 10, 20, 100, 110, 120}; + + const auto& taps = window.expand(); + ASSERT_EQ(taps.size(), expectedWindow.size()); + for(size_t i = 0; i < expectedWindow.size(); ++i) + { + EXPECT_EQ(taps[i].windowOffset, expectedWindow[i]) << "tap " << i; + EXPECT_EQ(taps[i].sourceOffset, expectedSource[i]) << "tap " << i; + } +} + +TEST_F(TestCpuFpReferenceUtilities, ConvolutionWindowForEachTapMatchesExpand) +{ + // wgrad walks its window with forEachTap where fprop and dgrad expand theirs; both + // accumulate in tap order, so the walk must visit exactly the expanded taps in the + // same order. Three dimensions with taps dropped from each exercise every odometer + // carry. + const std::vector extents{3, 4, 5}; + const std::vector windowStrides{20, 5, 1}; + const std::vector sourceStrides{1000, 100, 10}; + + ConvolutionWindow window; + window.build(extents.size(), + extents.data(), + windowStrides.data(), + sourceStrides.data(), + [](size_t dim, int64_t index) -> int64_t { + return (index == static_cast(dim)) ? -1 : index; + }); + + const auto expanded = window.expand(); + ASSERT_EQ(expanded.size(), 2u * 3u * 4u); + + std::vector walked; + window.forEachTap([&walked](const ConvolutionWindow::Tap& tap) { walked.push_back(tap); }); + + ASSERT_EQ(walked.size(), expanded.size()); + for(size_t i = 0; i < expanded.size(); ++i) + { + EXPECT_EQ(walked[i].windowOffset, expanded[i].windowOffset) << "tap " << i; + EXPECT_EQ(walked[i].sourceOffset, expanded[i].sourceOffset) << "tap " << i; + } +} + +TEST_F(TestCpuFpReferenceUtilities, ConvolutionWindowDropsTapsWithNoSourceElement) +{ + // A tap in the padding region has no source element. It must disappear while the + // surviving taps keep the window offsets they would have had. + const std::vector extents{3}; + const std::vector strides{1}; + + ConvolutionWindow window; + window.build( + extents.size(), extents.data(), strides.data(), strides.data(), [](size_t, int64_t index) { + return index == 0 ? -1 : index - 1; + }); + + const auto& taps = window.expand(); + ASSERT_EQ(taps.size(), 2u); + EXPECT_EQ(taps[0].windowOffset, 1); + EXPECT_EQ(taps[0].sourceOffset, 0); + EXPECT_EQ(taps[1].windowOffset, 2); + EXPECT_EQ(taps[1].sourceOffset, 1); +} + +TEST_F(TestCpuFpReferenceUtilities, ConvolutionWindowFactorsValidityPerDimension) +{ + // The whole optimization rests on validity factoring per dimension: dropping one + // index of the outer dimension must remove that entire slice of the product, not a + // single tap. + const std::vector extents{2, 2}; + const std::vector windowStrides{2, 1}; + const std::vector sourceStrides{10, 1}; + + ConvolutionWindow window; + window.build(extents.size(), + extents.data(), + windowStrides.data(), + sourceStrides.data(), + [](size_t dim, int64_t index) { return (dim == 0 && index == 1) ? -1 : index; }); + + const auto& taps = window.expand(); + ASSERT_EQ(taps.size(), 2u); + EXPECT_EQ(taps[0].windowOffset, 0); + EXPECT_EQ(taps[1].windowOffset, 1); +} + +TEST_F(TestCpuFpReferenceUtilities, ConvolutionWindowIsEmptyWhenNoTapHasASource) +{ + // A fully padded window must accumulate nothing rather than address anything. + const std::vector extents{2, 2}; + const std::vector strides{2, 1}; + + ConvolutionWindow window; + window.build( + extents.size(), extents.data(), strides.data(), strides.data(), [](size_t, int64_t) { + return -1; + }); + + EXPECT_TRUE(window.expand().empty()); + + int visits = 0; + window.forEachTap([&visits](const ConvolutionWindow::Tap&) { visits++; }); + EXPECT_EQ(visits, 0); +} + +TEST_F(TestCpuFpReferenceUtilities, ConvolutionWindowRebuildReplacesPreviousTaps) +{ + // One window is reused for every output position, so a rebuild must replace the + // previous position's taps rather than append to or leak them. + const std::vector extents{4}; + const std::vector strides{1}; + + ConvolutionWindow window; + window.build( + extents.size(), extents.data(), strides.data(), strides.data(), [](size_t, int64_t index) { + return index; + }); + ASSERT_EQ(window.expand().size(), 4u); + + window.build( + extents.size(), extents.data(), strides.data(), strides.data(), [](size_t, int64_t index) { + return index < 2 ? index : -1; + }); + + const auto& taps = window.expand(); + ASSERT_EQ(taps.size(), 2u); + EXPECT_EQ(taps[0].windowOffset, 0); + EXPECT_EQ(taps[1].windowOffset, 1); +} + +// ============================================================================ +// ParallelTensorFunctorWithScratch +// ============================================================================ + +namespace +{ + +// Counts its own construction so a test can tell per-thread scratch from per-work-item +// scratch, which is the entire difference between this functor and the plain one. +struct CountingScratch +{ + static inline std::atomic constructions{0}; + + CountingScratch() + { + constructions++; + } + + int itemsSeen{0}; +}; + +} // namespace + +TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorWithScratchConstructsOnePerThread) +{ + // The point of the scratch functor: construction is paid once per worker thread, not + // once per work item. A per-item scratch would report 64 constructions here. + CountingScratch::constructions = 0; + + constexpr int64_t ELEMENT_COUNT = 64; + constexpr std::size_t THREAD_COUNT = 4; + + std::atomic visits{0}; + auto functor = makeParallelTensorFunctorWithScratch( + [&visits](CountingScratch& scratch, const std::vector& indices) { + (void)indices; + scratch.itemsSeen++; + visits++; + }, + std::vector{ELEMENT_COUNT}); + functor(THREAD_COUNT); + + EXPECT_EQ(visits.load(), static_cast(ELEMENT_COUNT)); + EXPECT_EQ(CountingScratch::constructions.load(), static_cast(THREAD_COUNT)); +} + +TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorWithScratchReusesScratchAcrossWorkItems) +{ + // Single thread, so every work item must land on the same instance. A per-item + // scratch would leave itemsSeen at 1 and defeat the buffer reuse this exists for. + CountingScratch::constructions = 0; + + constexpr int64_t ELEMENT_COUNT = 10; + int maxItemsSeen = 0; + + auto functor = makeParallelTensorFunctorWithScratch( + [&maxItemsSeen](CountingScratch& scratch, const std::vector& indices) { + (void)indices; + scratch.itemsSeen++; + maxItemsSeen = std::max(maxItemsSeen, scratch.itemsSeen); + }, + std::vector{ELEMENT_COUNT}); + functor(1); + + EXPECT_EQ(CountingScratch::constructions.load(), 1); + EXPECT_EQ(maxItemsSeen, static_cast(ELEMENT_COUNT)); +} + +TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorWithScratchIsolatesThreads) +{ + // Threads must not share one scratch; sharing would be a data race in every caller. + // Every worker parks on its first item until all of them have arrived, so all their + // scratches are alive at once and distinct scratches must have distinct addresses. + constexpr int64_t ELEMENT_COUNT = 1000; + constexpr std::size_t THREAD_COUNT = 4; + + std::mutex mutex; + std::condition_variable arrived; + std::size_t arrivals = 0; + std::set addresses; + + auto functor = makeParallelTensorFunctorWithScratch( + [&](CountingScratch& scratch, const std::vector& indices) { + (void)indices; + if(scratch.itemsSeen++ > 0) + { + return; + } + + std::unique_lock lock(mutex); + addresses.insert(&scratch); + ++arrivals; + arrived.notify_all(); + // Bounded so a worker that never starts fails the test instead of hanging it. + arrived.wait_for( + lock, std::chrono::seconds(30), [&arrivals] { return arrivals == THREAD_COUNT; }); + }, + std::vector{ELEMENT_COUNT}); + functor(THREAD_COUNT); + + EXPECT_EQ(arrivals, THREAD_COUNT); + EXPECT_EQ(addresses.size(), THREAD_COUNT); +} + +TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorWithScratchVisitsEveryIndexOnce) +{ + // Same coverage contract as the plain functor: the scratch parameter must not + // perturb how the index space is decomposed or divided. + constexpr int64_t HEIGHT = 5; + constexpr int64_t WIDTH = 7; + + std::vector> visitCounts(static_cast(HEIGHT * WIDTH)); + for(auto& count : visitCounts) + { + count = 0; + } + + auto functor = makeParallelTensorFunctorWithScratch( + [&visitCounts](CountingScratch& scratch, const std::vector& indices) { + (void)scratch; + visitCounts[static_cast((indices[0] * WIDTH) + indices[1])]++; + }, + std::vector{HEIGHT, WIDTH}); + functor(3); + + for(size_t i = 0; i < visitCounts.size(); ++i) + { + EXPECT_EQ(visitCounts[i].load(), 1) + << "index " << i << " was visited the wrong number of times"; + } +} + +TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorWithScratchDoesNothingForAnEmptyRange) +{ + // A zero-sized dimension must spawn no threads and construct no scratch. + CountingScratch::constructions = 0; + std::atomic visits{0}; + + auto functor = makeParallelTensorFunctorWithScratch( + [&visits](CountingScratch& scratch, const std::vector& indices) { + (void)scratch; + (void)indices; + visits++; + }, + std::vector{0}); + functor(4); + + EXPECT_EQ(visits.load(), 0); + EXPECT_EQ(CountingScratch::constructions.load(), 0); +} + +// A bool-returning functor taking a non-const lvalue reference. It binds fine at the call +// site either way, so an early-exit probe that models the argument as an rvalue or as a +// `const&` silently drops the bool and runs the whole range. +namespace +{ + +struct MutableIndexEarlyExit +{ + std::atomic* visits; + + bool operator()(std::vector& indices) const + { + (*visits)++; + return indices[0] < 3; + } +}; + +} // namespace + +TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorDynamicStopsOnFalseForMutableIndices) +{ + std::atomic visits{0}; + + auto functor + = makeParallelTensorFunctor(MutableIndexEarlyExit{&visits}, std::vector{100}); + functor(1); + + // Indices 0, 1 and 2 continue; index 3 returns false and stops the thread. + EXPECT_EQ(visits.load(), 4); +} + +TEST_F(TestCpuFpReferenceUtilities, ParallelTensorFunctorWithScratchStopsOnFalseForMutableIndices) +{ + std::atomic visits{0}; + + auto functor = makeParallelTensorFunctorWithScratch( + [&visits](CountingScratch& scratch, std::vector& indices) { + (void)scratch; + visits++; + return indices[0] < 3; + }, + std::vector{100}); + functor(1); + + EXPECT_EQ(visits.load(), 4); +}