fix(conv): require large-tensor CK instance for grouped wrw on >INT_MAX element count

projects/miopen/CHANGELOG.md

@@ -3,6 +3,11 @@
 
 Full documentation for MIOpen is available [here](https://rocm.docs.amd.com/projects/MIOpen/en/latest/)
 
+## MIOpen 3.6.0 for ROCm 10.1.0
+
+### Resolved Issues
+* [Conv] Fixed silently incorrect results from the grouped backward-weights CK xdlops solver when a tensor's element extent exceeds INT_MAX but its individual lengths and strides still fit int32; such problems now use a large-tensor (int64) CK instance instead of overflowing int32 indexing.
+
 ## MIOpen 3.6.0 for ROCm 10.0.0
 ### Added
 * [Conv] Added gfx950 (MI350X/MI355X) 7x7 depthwise forward and backward-data convolution support (fp16/bf16), fixing a slow fallback-to-naive-kernel regression in ConvNeXt-style depthwise convolutions.

projects/miopen/src/ck_impl/implicitgemm_ck_util.hpp

@@ -6,6 +6,8 @@
 #include <miopen/solver/implicitgemm_ck_util_common.hpp>
 #include <miopen/kernel_tuning_mode.hpp>
 
+#include <limits>
+
 #if MIOPEN_BACKEND_HIP && MIOPEN_USE_COMPOSABLEKERNEL
 #include <ck/utility/data_type.hpp>
 #include <ck/utility/numeric_limits.hpp>
@@ -167,7 +169,8 @@ typename ConvPtrsType::iterator FindConvPtrByID(ConvPtrsType& conv_ptrs,
     });
 }
 
-// Returns true when any tensor stride/dim exceeds INT_MAX, so the problem
+// Returns true when any tensor stride/dim exceeds INT_MAX, or when a grouped
+// backward-weights tensor's addressable element extent exceeds INT_MAX, so the problem
 // must be dispatched to a CK instance whose argument types are int64 end-to-end.
 //
 // CK exposes both index_t (int32) and long_index_t (int64) MakeArgumentPointer
@@ -179,7 +182,29 @@ typename ConvPtrsType::iterator FindConvPtrByID(ConvPtrsType& conv_ptrs,
 template <typename ProblemDescriptionType>
 inline bool RequiresLargeTensorCKInstance(const ProblemDescriptionType& problem)
 {
-    return !problem.AllTensorsDimsFitIntoInt();
+    if(!problem.AllTensorsDimsFitIntoInt())
+        return true;
+
+    // AllTensorsDimsFitIntoInt() only checks that each individual length/stride fits in
+    // int32; it does NOT bound the maximum flat element offset the kernel indexes. A grouped
+    // backward-weights tensor whose element extent exceeds INT_MAX overflows a non-large CK
+    // instance's 32-bit element indexing and returns silently wrong results, so require a
+    // large-tensor (int64) instance for those too. Use GetElementSpace() (the stride-aware
+    // max addressable offset, as GetNumBytes() does), not GetElementSize() (the packed
+    // element count): for a non-packed tensor the element space is larger and is the true
+    // bound on the int32 offset that overflows. Scoped to the backward-weights direction:
+    // forward and backward-data verify correctly on the int32 instances and must keep them
+    // (the large-tensor instances are measurably slower). This gate is rank-agnostic, so it
+    // covers both 2D (ConvHipImplicitGemmGroupWrwXdlops) and 3D
+    // (ConvHipImplicitGemm3DGroupWrwXdlops) grouped backward-weights.
+    if(problem.IsDirectionBackwardWrW())
+    {
+        constexpr std::size_t max_int32 = static_cast<std::size_t>(std::numeric_limits<int>::max());
+        return problem.GetIn().GetElementSpace() > max_int32 ||
+               problem.GetOut().GetElementSpace() > max_int32 ||
+               problem.GetWeights().GetElementSpace() > max_int32;
+    }
+    return false;
 }
 
 // Large-tensor xdl impls embed "Large_Tensor" in their GetTypeString() (see

projects/miopen/test/gtest/test_categories.yaml

@@ -160,6 +160,7 @@ test_categories:
       - "Full/GPU_UnitTestConvSolverImplicitGemm3DGroupBwdXdlops_LargeStride_*"
       - "Full/GPU_UnitTestConvSolverImplicitGemmGroupWrwXdlops_LargeStride_*"
       - "Full/GPU_UnitTestConvSolverImplicitGemm3DGroupWrwXdlops_LargeStride_*"
+      - "Full/GPU_UnitTestConvSolverImplicitGemmGroupWrwXdlops_ElementCountOverflow_*"
       # Large BN tests - full suite only
       - "Full/GPU_BNCKFWDTrainLarge2D_*"
       - "Full/GPU_BNOCLFWDTrainLarge2D_*"
@@ -257,6 +258,9 @@ test_categories:
       # share the same dispatch code path and stay in full. Keep BFP16 in
       # comprehensive (nightly) for coverage.
       - "Full/GPU_UnitTestConvSolverImplicitGemm*Xdlops_LargeStride_FP*"
+      # Heavyweight (~60s) large-tensor wrw test; full-only like _LargeStride.
+      # FP16-only (no BFP16 variant to keep in nightly), so exclude the family.
+      - "Full/GPU_UnitTestConvSolverImplicitGemmGroupWrwXdlops_ElementCountOverflow_*"
       # Only exclude explicitly disabled tests
       - "*DISABLED*"
 

projects/miopen/test/gtest/unit_conv_solver_ConvHipImplicitGemmGroupWrwXdlops_ElementCountOverflow.cpp

@@ -0,0 +1,153 @@
+// Copyright © Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier:  MIT
+//
+// Numerical 2D WRW test for ConvHipImplicitGemmGroupWrwXdlops on a shape whose flattened
+// input element count exceeds INT_MAX while every individual tensor length/stride still
+// fits int32. This complements unit_conv_solver_ConvHipImplicitGemmGroupWrwXdlops_LargeStride.cpp
+// (which covers an out-of-range *stride*): RequiresLargeTensorCKInstance() must also catch
+// tensors whose per-dimension strides are all small but whose total element count overflows
+// a non-large CK instance's int32 element indexing (see implicitgemm_ck_util.hpp).
+//
+// Shape: x = (9, 256, 1024, 1024), w = (64, 256, 3, 3), group=1, pad=0, stride=1, dilation=1.
+//   element count of x = 9 * 256 * 1024 * 1024 = 2.416 B (> INT_MAX = 2.147 B).
+//   element count of y (dy, 3x3/pad0/stride1 => Ho=Wo=1022) = 9 * 64 * 1022 * 1022 = 602 M
+//   (comfortably under INT_MAX). Max per-dim stride (NHWC, N-dim) = C*H*W ~= 268.4 M, also
+//   comfortably under INT_MAX. So this shape trips the *element-count* gate specifically,
+//   not the (already-covered) per-stride gate: a 1x1-filter variant with the same X tensor
+//   does not reproduce the bug on this CK build (Ho=Wo=H=W means the tuner favors different
+//   instances), but the 3x3 filter here routes to an int32-indexing CK instance whose flat
+//   element offset overflows once selected, producing a large silent RMS error.
+//
+// FP16 footprint: x ~= 4.83 GB, y ~= 1.2 GB. The full test allocates several such tensors
+// (X, W, Y on device plus host-side reference), so heavyweight instances are gated at
+// runtime by an explicit memory estimate, mirroring the LargeStride test.
+//
+// Heavyweight and slow: the large-tensor (int64) CK instance required post-fix is
+// measurably slower than the int32 instance it replaces on this shape. Excluded from the
+// standard (per-PR) test category via test_categories.yaml.
+
+#include <algorithm>
+#include <cstddef>
+#include <utility>
+
+#include "get_handle.hpp"
+#include "unit_conv_solver_group_xdlops.hpp"
+
+namespace {
+
+using TestCase     = miopen::unit_tests::GroupXdlopsNumericData;
+using TestDataType = miopen::unit_tests::TestDataType;
+
+template <TestDataType type>
+std::vector<TestCase> GetElementCountOverflowWrwTestCases()
+{
+    return {
+        // clang-format off
+        TestCase{{9, 256, 1024, 1024}, {64, 256, 3, 3}, {0, 0}, {1, 1}, {1, 1}, 1, false, false},
+        // clang-format on
+    };
+}
+
+template <TestDataType type>
+miopen::unit_tests::UnitTestConvSolverParams GetTestParams()
+{
+#if MIOPEN_BACKEND_HIP && MIOPEN_USE_COMPOSABLEKERNEL
+    // Restricted to gfx90A, gfx94X, and gfx950: covered by CI and manually
+    // qualified for the large-tensor kernel-launch path on this shape.
+    Gpu supportedDevices = Gpu::gfx90A | Gpu::gfx94X | Gpu::gfx950;
+#else
+    Gpu supportedDevices = Gpu::None;
+#endif
+    miopen::unit_tests::UnitTestConvSolverParams p(supportedDevices);
+    p.Tunable(5);
+    p.UsesCKDynamicLib();
+    p.UseGpuRef();
+    if constexpr(type == TestDataType::FP32)
+    {
+        // WRW reduces over N*H*W per output element (~2.4M FMAs per output pixel-channel
+        // reduction group at this shape). RMS error scales with reduction size, so a much
+        // larger tolerance bump than a small-reduction Fwd test is required.
+        p.SetTolerance(supportedDevices, miopenFloat, 80.0f);
+    }
+    return p;
+}
+
+// Conservative working-set estimate for the configured Wrw test. Sums the workspace
+// (queried from the solver), the X/W/Y device tensors, and 4x the largest tensor for the
+// host-side input/weights/output/reference allocations. Adds headroom for runtime/library
+// reservations, allocator fragmentation, and (on consumer cards) the display compositor --
+// using max(+1 GiB, +10%) to cover both the absolute and the proportional components.
+struct MemoryEstimate
+{
+    std::size_t required;
+    std::size_t available;
+};
+
+template <miopenDataType_t datatype>
+MemoryEstimate EstimateRequiredMemoryWrw(TestCase tc,
+                                         miopenTensorLayout_t layout,
+                                         const miopen::solver::conv::ConvSolverInterface& solver)
+{
+    auto conv_case = miopen::unit_tests::GetConvTestForGroupXdlops<datatype>(layout, std::move(tc));
+    const auto x_desc = conv_case.GetXTensorDescriptor();
+    const auto w_desc = conv_case.GetWTensorDescriptor();
+    const auto y_desc =
+        conv_case.GetConv().GetForwardOutputTensor(x_desc, w_desc, conv_case.GetYDataType());
+
+    auto&& handle      = get_handle();
+    const auto problem = miopen::conv::ProblemDescription(
+        y_desc, w_desc, x_desc, conv_case.GetConv(), miopen::conv::Direction::BackwardWeights);
+    auto ctx = miopen::ExecutionContext{&handle};
+    problem.SetupFloats(ctx);
+    problem.SetupComputeType(ctx);
+
+    const std::size_t ws_size =
+        solver.MayNeedWorkspace() ? solver.GetWorkspaceSize(ctx, problem) : 0;
+    const std::size_t x_bytes = x_desc.GetNumBytes();
+    const std::size_t y_bytes = y_desc.GetNumBytes();
+    const std::size_t w_size  = w_desc.GetNumBytes();
+    const std::size_t h_bytes = std::max(x_bytes, y_bytes);
+
+    const std::size_t raw_mem      = ws_size + x_bytes + y_bytes + w_size + 4 * h_bytes;
+    const std::size_t headroom     = std::max<std::size_t>(1ULL << 30, raw_mem / 10);
+    const std::size_t required_mem = raw_mem + headroom;
+    const std::size_t device_mem   = handle.GetGlobalMemorySize();
+
+    return {required_mem, device_mem};
+}
+
+} // namespace
+
+#define SKIP_IF_INSUFFICIENT_DEVICE_MEMORY(datatype, solver_expr)                           \
+    do                                                                                      \
+    {                                                                                       \
+        miopen::unit_tests::UnitTestConvSolverParams _params;                               \
+        miopenTensorLayout_t _layout;                                                       \
+        TestCase _tc;                                                                       \
+        std::tie(_params, _layout, _tc) = this->GetParam();                                 \
+        const auto _mem = EstimateRequiredMemoryWrw<datatype>(_tc, _layout, (solver_expr)); \
+        if(_mem.available < _mem.required)                                                  \
+        {                                                                                   \
+            GTEST_SKIP() << "Insufficient device memory: need " << _mem.required            \
+                         << " bytes, device has " << _mem.available;                        \
+        }                                                                                   \
+    } while(0)
+
+using GPU_UnitTestConvSolverImplicitGemmGroupWrwXdlops_ElementCountOverflow_FP16 =
+    miopen::unit_tests::UnitTestConvSolverGroupXDlops<miopen::conv::Direction::BackwardWeights,
+                                                      miopenHalf>;
+
+TEST_P(GPU_UnitTestConvSolverImplicitGemmGroupWrwXdlops_ElementCountOverflow_FP16,
+       ConvHipImplicitGemmGroupWrwXdlops)
+{
+    const auto solver = miopen::solver::conv::ConvHipImplicitGemmGroupWrwXdlops{};
+    SKIP_IF_INSUFFICIENT_DEVICE_MEMORY(miopenHalf, solver);
+    this->RunTest(solver);
+};
+
+INSTANTIATE_TEST_SUITE_P(
+    Full,
+    GPU_UnitTestConvSolverImplicitGemmGroupWrwXdlops_ElementCountOverflow_FP16,
+    testing::Combine(testing::Values(GetTestParams<TestDataType::FP16>()),
+                     testing::Values(miopenTensorNHWC),
+                     testing::ValuesIn(GetElementCountOverflowWrwTestCases<TestDataType::FP16>())));

projects/miopen/test/gtest/unit_implicitgemm_ck_util.cpp

@@ -26,6 +26,8 @@
 
 #include <miopen/conv/problem_description.hpp>
 #include <miopen/conv/solvers.hpp>
+#include <miopen/convolution.hpp>
+#include <miopen/tensor.hpp>
 #include "../../src/ck_impl/implicitgemm_ck_util.hpp"
 #include <gtest/gtest.h>
 #include <iostream>
@@ -416,3 +418,127 @@ TEST_F(CPU_SplitKGenericDirectTest_NONE, MultipleKernelsWithDifferentSplitKSuppo
 }
 
 } // namespace unit_implicitgemm_ck_util_test
+
+// ---------------------------------------------------------------------------
+// Regression: RequiresLargeTensorCKInstance must force a large-tensor (int64) CK
+// instance for a grouped backward-weights problem whose flattened element count
+// exceeds INT_MAX even though every individual length/stride still fits int32.
+// Without the element-count gate the loader keeps such a problem eligible for a
+// fast int32-indexing instance whose flat element offset overflows, silently
+// returning wrong results. The gate is scoped to backward-weights: forward and
+// backward-data verify correctly on the int32 instances. Shape is synthetic.
+// ---------------------------------------------------------------------------
+namespace {
+
+miopen::conv::ProblemDescription MakeElemCountConvProblem(miopen::conv::Direction dir,
+                                                          std::size_t n,
+                                                          std::size_t c,
+                                                          std::size_t h,
+                                                          std::size_t w,
+                                                          std::size_t k)
+{
+    const miopen::TensorDescriptor x{miopenHalf, miopenTensorNHWC, {n, c, h, w}};
+    const miopen::TensorDescriptor weights{miopenHalf, miopenTensorNHWC, {k, c, 3, 3}};
+    const miopen::ConvolutionDescriptor conv{2,
+                                             miopenConvolution,
+                                             miopenPaddingDefault,
+                                             {1, 1}, // pads (SAME for 3x3)
+                                             {1, 1}, // strides
+                                             {1, 1}, // dilations
+                                             {0, 0}, // trans output pads
+                                             1,      // group count
+                                             1.0f};  // lowp quant
+    const miopen::TensorDescriptor y = conv.GetForwardOutputTensor(x, weights, miopenHalf);
+
+    // ProblemDescription ctor: (in = x for fwd / y for backward*, weights,
+    //                           out = y for fwd / x for backward*, conv, direction)
+    if(dir == miopen::conv::Direction::Forward)
+        return {x, weights, y, conv, dir};
+    return {y, weights, x, conv, dir};
+}
+
+// 3D (NDHWC) sibling of MakeElemCountConvProblem. The element-count gate is
+// rank-agnostic, so it must fire for 3D grouped wrw (ConvHipImplicitGemm3DGroupWrwXdlops)
+// on the same >INT_MAX-element-extent condition as the 2D case.
+miopen::conv::ProblemDescription MakeElemCount3DConvProblem(miopen::conv::Direction dir,
+                                                            std::size_t n,
+                                                            std::size_t c,
+                                                            std::size_t d,
+                                                            std::size_t h,
+                                                            std::size_t w,
+                                                            std::size_t k)
+{
+    const miopen::TensorDescriptor x{miopenHalf, miopenTensorNDHWC, {n, c, d, h, w}};
+    const miopen::TensorDescriptor weights{miopenHalf, miopenTensorNDHWC, {k, c, 3, 3, 3}};
+    const miopen::ConvolutionDescriptor conv{3,
+                                             miopenConvolution,
+                                             miopenPaddingDefault,
+                                             {1, 1, 1}, // pads (SAME for 3x3x3)
+                                             {1, 1, 1}, // strides
+                                             {1, 1, 1}, // dilations
+                                             {0, 0, 0}, // trans output pads
+                                             1,         // group count
+                                             1.0f};     // lowp quant
+    const miopen::TensorDescriptor y = conv.GetForwardOutputTensor(x, weights, miopenHalf);
+
+    if(dir == miopen::conv::Direction::Forward)
+        return {x, weights, y, conv, dir};
+    return {y, weights, x, conv, dir};
+}
+
+} // namespace
+
+TEST(CPU_UnitTestImplicitGemmCKUtilLargeTensor_NONE, WrwElementCountOverflowGate)
+{
+    using miopen::conv::Direction;
+    using miopen::solver::RequiresLargeTensorCKInstance;
+
+    // N*C*H*W = 256*1024*128*128 = 4.29e9 > INT_MAX, yet every length and stride
+    // (max stride = H*W*C = 16.78M) fits int32, so AllTensorsDimsFitIntoInt() is true
+    // and only the element-count gate can flag it.
+    constexpr std::size_t n = 256, c = 1024, hw = 128, k = 1024;
+
+    EXPECT_TRUE(RequiresLargeTensorCKInstance(
+        MakeElemCountConvProblem(Direction::BackwardWeights, n, c, hw, hw, k)))
+        << "grouped wrw with >INT_MAX element count must require a large-tensor instance";
+
+    EXPECT_FALSE(RequiresLargeTensorCKInstance(
+        MakeElemCountConvProblem(Direction::Forward, n, c, hw, hw, k)))
+        << "forward must stay on the (faster) int32 instances";
+
+    EXPECT_FALSE(RequiresLargeTensorCKInstance(
+        MakeElemCountConvProblem(Direction::BackwardData, n, c, hw, hw, k)))
+        << "backward-data must stay on the (faster) int32 instances";
+
+    // A small wrw problem (element count well within int32) stays on the int32 path.
+    EXPECT_FALSE(RequiresLargeTensorCKInstance(
+        MakeElemCountConvProblem(Direction::BackwardWeights, 1, 64, 32, 32, 64)));
+}
+
+TEST(CPU_UnitTestImplicitGemmCKUtilLargeTensor_NONE, Wrw3DElementCountOverflowGate)
+{
+    using miopen::conv::Direction;
+    using miopen::solver::RequiresLargeTensorCKInstance;
+
+    // N*C*D*H*W = 256*512*32*32*32 = 4.29e9 > INT_MAX, yet every length and stride
+    // (max stride = C*D*H*W = 16.78M) fits int32, so AllTensorsDimsFitIntoInt() is true
+    // and only the element-count gate can flag it. Mirrors WrwElementCountOverflowGate
+    // for the 3D (NDHWC) grouped wrw path.
+    constexpr std::size_t n = 256, c = 512, d = 32, hw = 32, k = 512;
+
+    EXPECT_TRUE(RequiresLargeTensorCKInstance(
+        MakeElemCount3DConvProblem(Direction::BackwardWeights, n, c, d, hw, hw, k)))
+        << "3D grouped wrw with >INT_MAX element count must require a large-tensor instance";
+
+    EXPECT_FALSE(RequiresLargeTensorCKInstance(
+        MakeElemCount3DConvProblem(Direction::Forward, n, c, d, hw, hw, k)))
+        << "3D forward must stay on the (faster) int32 instances";
+
+    EXPECT_FALSE(RequiresLargeTensorCKInstance(
+        MakeElemCount3DConvProblem(Direction::BackwardData, n, c, d, hw, hw, k)))
+        << "3D backward-data must stay on the (faster) int32 instances";
+
+    // A small 3D wrw problem (element count well within int32) stays on the int32 path.
+    EXPECT_FALSE(RequiresLargeTensorCKInstance(
+        MakeElemCount3DConvProblem(Direction::BackwardWeights, 1, 64, 8, 16, 16, 64)));
+}