feat(hipdnn): Implement a container like ProblemDescription to pack problem parameters

dnn-providers/hip-kernel-provider/src/engines/hip_mlops_engine/plans/RMSnorm/RMSnormBwdPlan.cpp

@@ -95,10 +95,10 @@ void RMSnormBwdPlan::compile([[maybe_unused]] const IKernelCompiler& kernelCompi
     }
 
     // Get problem dimensions
-    const unsigned normalizeDim = getNormalizeDim(_params.x()->dims(), _params.scale()->dims());
-    const int64_t stride = getStride(_params.x(), normalizeDim);
-    const int64_t outerSize = getOuterSize(_params.x()->dims(), normalizeDim, stride);
-    const int64_t innerSize = getInnerSize(_params.x()->dims(), normalizeDim);
+    const ProblemDescription problem(_params.x(), _params.scale(), Direction::BACKWARD);
+    const int64_t stride = problem.stride();
+    const int64_t outerSize = problem.outerSize();
+    const int64_t innerSize = problem.innerSize();
     if(outerSize * stride >= UINT32_MAX)
     {
         throw hipdnn_plugin_sdk::HipdnnPluginException(HIPDNN_PLUGIN_STATUS_BAD_PARAM,

dnn-providers/hip-kernel-provider/src/engines/hip_mlops_engine/plans/RMSnorm/RMSnormCommon.hpp

@@ -9,61 +9,128 @@
 namespace hip_kernel_provider::rmsnorm
 {
 
-inline int64_t
-    getOuterSize(const flatbuffers::Vector<int64_t>* xDims, unsigned normalizeDim, int64_t stride)
+enum class Direction
 {
-    int64_t outerSize = 1;
-    for(unsigned i = 0; i < normalizeDim; ++i)
+    FORWARD,
+    BACKWARD
+};
+
+// Infer the outer and inner normalization sizes.
+// 1) Work out the normalization dimension, as the index of the first dimension in
+//    scale tensor that is not 1.
+// 2) Work out the outerSize as the size of input dimensions for which scale is 1.
+//    We will have a work-group for each of these dimensions.
+//    When stride is not 1, we are in a channel-last layout and we ignore the
+//    channel dimension when calculating the outer size.
+// 3) Work out the innerSize as the size of the input dimensions for which scale is not 1.
+//
+// For an input of [N, C, H, W] with scale [1, C, H, W] this will give: a normalization
+// dimension of 1, outerSize of N, and innerSize of CxHxW.
+// The kernel will therefore consist of N workgroups with each workgroup normalizing over
+// CxHxW elements using a fixed number of threads.
+// For an input of [N, H, W, C] with scale [1, H, W, 1] this will give: a normalization
+// dimension of 2, outerSize of N, stride of C, and innerSize of HxW.
+// The kernel will therefore consist of NxC workgroups with each workgroup normalizing
+// over HxW elements using a fixed number of threads.
+class ProblemDescription
+{
+    static unsigned getNormalizeDim(const flatbuffers::Vector<int64_t>* xDims,
+                                    const flatbuffers::Vector<int64_t>* scaleDims)
+    {
+        const std::vector<int64_t> xDimsVec(xDims->begin(), xDims->end());
+        const std::vector<int64_t> scaleDimsVec(scaleDims->begin(), scaleDims->end());
+
+        // Find number of trailing dims where scaleDims[i] == inputDims[i]
+        const auto [scaleMismatch, _] = std::mismatch(
+            scaleDimsVec.rbegin(), scaleDimsVec.rend(), xDimsVec.rbegin(), xDimsVec.rend());
+        const auto matchCount
+            = static_cast<size_t>(std::distance(scaleDimsVec.rbegin(), scaleMismatch));
+
+        // Scale must have at least one normalization axis, so account for the case where
+        // input has a single batch and scale matches exactly.
+        const auto normalizeDim
+            = (matchCount == scaleDimsVec.size()) ? 1 : scaleDimsVec.size() - matchCount;
+        return static_cast<unsigned>(normalizeDim);
+    }
+
+    static int64_t getStride(const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes* x,
+                             unsigned normalizeDim)
     {
-        // Add channel size only if there is no stride
-        if(i == 1 && stride != 1)
+        int64_t stride = 1;
+        const auto isLayoutNHWC = core::utils::isChannelLastLayout(x);
+        if(normalizeDim > 1 && isLayoutNHWC)
         {
-            continue;
+            stride = static_cast<int64_t>(x->dims()->Get(1));
         }
-        outerSize *= static_cast<int64_t>(xDims->Get(i));
+        return stride;
     }
-    return outerSize;
-}
 
-inline int64_t getInnerSize(const flatbuffers::Vector<int64_t>* xDims, unsigned normalizeDim)
-{
-    int64_t innerSize = 1;
-    for(unsigned i = normalizeDim; i < xDims->size(); ++i)
+    static int64_t getOuterSize(const flatbuffers::Vector<int64_t>* xDims,
+                                unsigned normalizeDim,
+                                int64_t stride)
     {
-        innerSize *= xDims->Get(i);
+        int64_t outerSize = 1;
+        for(unsigned i = 0; i < normalizeDim; ++i)
+        {
+            // Add channel size only if there is no stride
+            if(i == 1 && stride != 1)
+            {
+                continue;
+            }
+            outerSize *= static_cast<int64_t>(xDims->Get(i));
+        }
+        return outerSize;
     }
-    return innerSize;
-}
 
-inline int64_t getStride(const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes* x,
-                         unsigned normalizeDim)
-{
-    int64_t stride = 1;
-    auto isLayoutNHWC = core::utils::isChannelLastLayout(x);
-    if(normalizeDim > 1 && isLayoutNHWC)
+    static int64_t getInnerSize(const flatbuffers::Vector<int64_t>* xDims, unsigned normalizeDim)
     {
-        stride = static_cast<int64_t>(x->dims()->Get(1));
+        int64_t innerSize = 1;
+        for(unsigned i = normalizeDim; i < xDims->size(); ++i)
+        {
+            innerSize *= xDims->Get(i);
+        }
+        return innerSize;
     }
-    return stride;
-}
 
-inline unsigned getNormalizeDim(const flatbuffers::Vector<int64_t>* xDims,
-                                const flatbuffers::Vector<int64_t>* scaleDims)
-{
-    const std::vector<int64_t> xDimsVec(xDims->begin(), xDims->end());
-    const std::vector<int64_t> scaleDimsVec(scaleDims->begin(), scaleDims->end());
+public:
+    ProblemDescription(const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes* x,
+                       const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes* scale,
+                       Direction direction)
+        : _direction(direction)
+        , _normalizeDim(getNormalizeDim(x->dims(), scale->dims()))
+        , _stride(getStride(x, _normalizeDim))
+        , _outerSize(getOuterSize(x->dims(), _normalizeDim, _stride))
+        , _innerSize(getInnerSize(x->dims(), _normalizeDim))
+    {
+    }
 
-    // Find number of trailing dims where scaleDims[i] == inputDims[i]
-    const auto [scaleMismatch, _] = std::mismatch(
-        scaleDimsVec.rbegin(), scaleDimsVec.rend(), xDimsVec.rbegin(), xDimsVec.rend());
-    const auto matchCount
-        = static_cast<size_t>(std::distance(scaleDimsVec.rbegin(), scaleMismatch));
+    Direction direction() const
+    {
+        return _direction;
+    }
+    unsigned normalizeDim() const
+    {
+        return _normalizeDim;
+    }
+    int64_t outerSize() const
+    {
+        return _outerSize;
+    }
+    int64_t innerSize() const
+    {
+        return _innerSize;
+    }
+    int64_t stride() const
+    {
+        return _stride;
+    }
 
-    // Scale must have at least one normalization axis, so account for the case where input has a single batch and scale
-    // matches exactly.
-    const auto normalizeDim
-        = (matchCount == scaleDimsVec.size()) ? 1 : scaleDimsVec.size() - matchCount;
-    return static_cast<unsigned>(normalizeDim);
-}
+private:
+    Direction _direction;
+    unsigned _normalizeDim;
+    int64_t _stride;
+    int64_t _outerSize;
+    int64_t _innerSize;
+};
 
 }

dnn-providers/hip-kernel-provider/src/engines/hip_mlops_engine/plans/RMSnorm/RMSnormFwdPlan.cpp

@@ -96,27 +96,10 @@ void RMSnormFwdPlan::compile(const IKernelCompiler& kernelCompiler,
                                                            + std::to_string(xRank));
     }
 
-    // Infer the outer and inner normalization sizes.
-    // 1) Work out the normalization dimension, as the index of the first dimension in
-    //    scale tensor that is not 1.
-    // 2) Work out the outerSize as the size of input dimensions for which scale is 1.
-    //    We will have a work-group for each of these dimensions.
-    //    When stride is not 1, we are in a channel-last layout and we ignore the
-    //    channel dimension when calculating the outer size.
-    // 3) Work out the innerSize as the size of the intput dimensions for which scale nor 1.
-    //
-    // For an input of [N, C, H, W] with scale [1, C, H, W] this will give: a normalization
-    // dimension of 1, outerSize of N, and innerSize of CxHxW.
-    // The kernel will therefore consist of N workgroups with each workgroup normalizing over
-    // CxHxW elements using a fixed number of threads.
-    // For an input of [N, H, W, C] with scale [1, H, W, 1] this will give: a normalization
-    // dimension of 2, outerSize of N, stride of C, and innerSize of HxW.
-    // The kernel will therefore consist of NxC workgroups with each workgroup normalizing
-    // over HxW elements using a fixed number of threads.
-    const unsigned normalizeDim = getNormalizeDim(_params.x()->dims(), _params.scale()->dims());
-    const int64_t stride = getStride(_params.x(), normalizeDim);
-    const int64_t outerSize = getOuterSize(_params.x()->dims(), normalizeDim, stride);
-    const int64_t innerSize = getInnerSize(_params.x()->dims(), normalizeDim);
+    const ProblemDescription problem(_params.x(), _params.scale(), Direction::FORWARD);
+    const int64_t stride = problem.stride();
+    const int64_t outerSize = problem.outerSize();
+    const int64_t innerSize = problem.innerSize();
     if(outerSize * stride >= UINT32_MAX)
     {
         throw hipdnn_plugin_sdk::HipdnnPluginException(HIPDNN_PLUGIN_STATUS_BAD_PARAM,