diff --git a/projects/miopen/src/ck_impl/implicitgemm_ck_util.hpp b/projects/miopen/src/ck_impl/implicitgemm_ck_util.hpp index d98d72124913..966e832179b2 100644 --- a/projects/miopen/src/ck_impl/implicitgemm_ck_util.hpp +++ b/projects/miopen/src/ck_impl/implicitgemm_ck_util.hpp @@ -169,42 +169,24 @@ typename ConvPtrsType::iterator FindConvPtrByID(ConvPtrsType& conv_ptrs, }); } -// 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 -// overloads, but most non-large-tensor impls override the int64 overload by -// silently narrowing back to int32 (e.g. -// device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp:2090-2117). Without -// this filter the loader would happily pick a narrowing instance for an -// out-of-range stride and corrupt the kernel arguments. +// Non-large-tensor CK instances silently narrow CK's int64 MakeArgumentPointer overload +// back to int32 and form byte offsets in 32-bit arithmetic, so any tensor whose byte span +// exceeds INT_MAX overflows and returns silently wrong results -- even when every individual +// dim/stride and the raw element count fit in int32. Must use GetNumBytes() (stride-aware +// byte span), not GetElementSpace()/GetElementSize() (element counts): a tensor can have a +// large element count but a sub-2GiB byte span (no risk), or the reverse -- a sub-INT_MAX +// element count with a >2GiB byte span, which does overflow. Applies to every direction: +// forward and backward-data hit the same 32-bit byte-offset overflow as backward-weights. template inline bool RequiresLargeTensorCKInstance(const ProblemDescriptionType& problem) { 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::numeric_limits::max()); - return problem.GetIn().GetElementSpace() > max_int32 || - problem.GetOut().GetElementSpace() > max_int32 || - problem.GetWeights().GetElementSpace() > max_int32; - } - return false; + constexpr std::size_t max_int32 = static_cast(std::numeric_limits::max()); + return problem.GetIn().GetNumBytes() > max_int32 || + problem.GetOut().GetNumBytes() > max_int32 || + problem.GetWeights().GetNumBytes() > max_int32; } // Large-tensor xdl impls embed "Large_Tensor" in their GetTypeString() (see diff --git a/projects/miopen/src/kernels/gfx942130.HIP.fdb.txt.bz2 b/projects/miopen/src/kernels/gfx942130.HIP.fdb.txt.bz2 index 9d78b290ecf7..e0cc09d53c98 100644 Binary files a/projects/miopen/src/kernels/gfx942130.HIP.fdb.txt.bz2 and b/projects/miopen/src/kernels/gfx942130.HIP.fdb.txt.bz2 differ diff --git a/projects/miopen/src/kernels/gfx942130.db.txt.bz2 b/projects/miopen/src/kernels/gfx942130.db.txt.bz2 index 4e1adfdd7f4f..6dc412a40163 100644 Binary files a/projects/miopen/src/kernels/gfx942130.db.txt.bz2 and b/projects/miopen/src/kernels/gfx942130.db.txt.bz2 differ diff --git a/projects/miopen/src/kernels/gfx942e4.HIP.fdb.txt.bz2 b/projects/miopen/src/kernels/gfx942e4.HIP.fdb.txt.bz2 index 79ceb76095a7..26b168bf899a 100644 Binary files a/projects/miopen/src/kernels/gfx942e4.HIP.fdb.txt.bz2 and b/projects/miopen/src/kernels/gfx942e4.HIP.fdb.txt.bz2 differ diff --git a/projects/miopen/src/kernels/gfx942e4.db.txt.bz2 b/projects/miopen/src/kernels/gfx942e4.db.txt.bz2 index 0849207a85c3..288ec510b0df 100644 Binary files a/projects/miopen/src/kernels/gfx942e4.db.txt.bz2 and b/projects/miopen/src/kernels/gfx942e4.db.txt.bz2 differ diff --git a/projects/miopen/src/kernels/gfx950100.HIP.fdb.txt.bz2 b/projects/miopen/src/kernels/gfx950100.HIP.fdb.txt.bz2 index 441f99d9fc9f..3958b4053955 100644 Binary files a/projects/miopen/src/kernels/gfx950100.HIP.fdb.txt.bz2 and b/projects/miopen/src/kernels/gfx950100.HIP.fdb.txt.bz2 differ diff --git a/projects/miopen/src/kernels/gfx950100.db.txt.bz2 b/projects/miopen/src/kernels/gfx950100.db.txt.bz2 index f9929fee188b..49497ab8f4ad 100644 Binary files a/projects/miopen/src/kernels/gfx950100.db.txt.bz2 and b/projects/miopen/src/kernels/gfx950100.db.txt.bz2 differ diff --git a/projects/miopen/test/gtest/unit_implicitgemm_ck_util.cpp b/projects/miopen/test/gtest/unit_implicitgemm_ck_util.cpp index d35304f83fd8..ede150935906 100644 --- a/projects/miopen/test/gtest/unit_implicitgemm_ck_util.cpp +++ b/projects/miopen/test/gtest/unit_implicitgemm_ck_util.cpp @@ -420,13 +420,9 @@ 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. +// Regression: the gate keys on byte span, not element count, and applies to every +// direction (previously backward-weights only). Shapes below are chosen so every +// individual length/stride still fits int32 -- only the byte-span check can flag them. // --------------------------------------------------------------------------- namespace { @@ -486,6 +482,35 @@ miopen::conv::ProblemDescription MakeElemCount3DConvProblem(miopen::conv::Direct return {y, weights, x, conv, dir}; } +// dtype-parameterized sibling of MakeElemCountConvProblem: fp32's larger byte-per-element +// ratio lets a shape cross the byte-span boundary while its element count stays well +// under INT_MAX, isolating the byte-span check from the element-count one. +miopen::conv::ProblemDescription MakeByteSpanConvProblem(miopenDataType_t dtype, + 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{dtype, miopenTensorNHWC, {n, c, h, w}}; + const miopen::TensorDescriptor weights{dtype, 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, dtype); + + if(dir == miopen::conv::Direction::Forward) + return {x, weights, y, conv, dir}; + return {y, weights, x, conv, dir}; +} + } // namespace TEST(CPU_UnitTestImplicitGemmCKUtilLargeTensor_NONE, WrwElementCountOverflowGate) @@ -493,22 +518,21 @@ 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. + // N*C*H*W = 256*1024*128*128 = 4.29e9 elements; fp16 byte span ~8.58 GB, over INT_MAX + // bytes despite every stride (max = H*W*C = 16.78M) fitting int32. 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"; + << "grouped wrw with >INT_MAX byte span must require a large-tensor instance"; - EXPECT_FALSE(RequiresLargeTensorCKInstance( + EXPECT_TRUE(RequiresLargeTensorCKInstance( MakeElemCountConvProblem(Direction::Forward, n, c, hw, hw, k))) - << "forward must stay on the (faster) int32 instances"; + << "forward with >INT_MAX byte span must require a large-tensor instance"; - EXPECT_FALSE(RequiresLargeTensorCKInstance( + EXPECT_TRUE(RequiresLargeTensorCKInstance( MakeElemCountConvProblem(Direction::BackwardData, n, c, hw, hw, k))) - << "backward-data must stay on the (faster) int32 instances"; + << "backward-data with >INT_MAX byte span must require a large-tensor instance"; // A small wrw problem (element count well within int32) stays on the int32 path. EXPECT_FALSE(RequiresLargeTensorCKInstance( @@ -520,25 +544,63 @@ TEST(CPU_UnitTestImplicitGemmCKUtilLargeTensor_NONE, Wrw3DElementCountOverflowGa 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. + // Mirrors WrwElementCountOverflowGate for the 3D (NDHWC) path: N*C*D*H*W = + // 256*512*32*32*32 = 4.29e9 elements, fp16 byte span ~8.58 GB, over INT_MAX bytes + // despite every stride (max = C*D*H*W = 16.78M) fitting int32. 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"; + << "3D grouped wrw with >INT_MAX byte span must require a large-tensor instance"; - EXPECT_FALSE(RequiresLargeTensorCKInstance( + EXPECT_TRUE(RequiresLargeTensorCKInstance( MakeElemCount3DConvProblem(Direction::Forward, n, c, d, hw, hw, k))) - << "3D forward must stay on the (faster) int32 instances"; + << "3D forward with >INT_MAX byte span must require a large-tensor instance"; - EXPECT_FALSE(RequiresLargeTensorCKInstance( + EXPECT_TRUE(RequiresLargeTensorCKInstance( MakeElemCount3DConvProblem(Direction::BackwardData, n, c, d, hw, hw, k))) - << "3D backward-data must stay on the (faster) int32 instances"; + << "3D backward-data with >INT_MAX byte span must require a large-tensor instance"; // 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))); } + +TEST(CPU_UnitTestImplicitGemmCKUtilLargeTensor_NONE, ByteSpanOverflowGate) +{ + using miopen::conv::Direction; + using miopen::solver::RequiresLargeTensorCKInstance; + + // fp32, c=512, h=w=32, k=512, 3x3/pad1/stride1 (=> Ho=Wo=32, Y = {n,512,32,32}); element + // count stays under INT_MAX on both sides of the boundary below, so only the byte span + // (not GetElementSpace()) distinguishes them -- this is the case the original bug missed. + // X and Y have equal span here (k=c, Ho=h, Wo=w): + // n=1024: elem = 536,870,912 (< INT_MAX); bytes = 2,147,483,648 (INT_MAX + 1) -> TRIGGER + // n=1023: elem = 536,346,624 (< INT_MAX); bytes = 2,145,386,496 (< INT_MAX) -> no trigger + constexpr std::size_t c = 512, hw = 32, k = 512; + constexpr std::size_t n_over = 1024; // byte span 2,147,483,648 > INT_MAX + constexpr std::size_t n_under = 1023; // byte span 2,145,386,496 < INT_MAX + + // Over the 2 GiB byte boundary: every direction must require a large-tensor instance. + EXPECT_TRUE(RequiresLargeTensorCKInstance( + MakeByteSpanConvProblem(miopenFloat, Direction::Forward, n_over, c, hw, hw, k))) + << "fwd: >INT_MAX byte span (element count still < INT_MAX) must require large-tensor"; + EXPECT_TRUE(RequiresLargeTensorCKInstance( + MakeByteSpanConvProblem(miopenFloat, Direction::BackwardData, n_over, c, hw, hw, k))) + << "bwd-data: >INT_MAX byte span must require large-tensor"; + EXPECT_TRUE(RequiresLargeTensorCKInstance( + MakeByteSpanConvProblem(miopenFloat, Direction::BackwardWeights, n_over, c, hw, hw, k))) + << "wrw: >INT_MAX byte span must require large-tensor"; + + // One N below the boundary: byte span < INT_MAX, element count < INT_MAX, all strides + // fit int32 -> every direction must stay on the int32 path. + EXPECT_FALSE(RequiresLargeTensorCKInstance( + MakeByteSpanConvProblem(miopenFloat, Direction::Forward, n_under, c, hw, hw, k))) + << "fwd: byte span just under INT_MAX must stay on the int32 instances"; + EXPECT_FALSE(RequiresLargeTensorCKInstance( + MakeByteSpanConvProblem(miopenFloat, Direction::BackwardData, n_under, c, hw, hw, k))) + << "bwd-data: byte span just under INT_MAX must stay on the int32 instances"; + EXPECT_FALSE(RequiresLargeTensorCKInstance( + MakeByteSpanConvProblem(miopenFloat, Direction::BackwardWeights, n_under, c, hw, hw, k))) + << "wrw: byte span just under INT_MAX must stay on the int32 instances"; +}