fix(miopen): reject reduce tensors that do not fit the int32 legacy kernels

projects/miopen/CHANGELOG.md

@@ -5,13 +5,26 @@ Full documentation for MIOpen is available [here](https://rocm.docs.amd.com/proj
 
 ## MIOpen 3.6.2 for ROCm 10.2.0
 
+### Resolved Issues
+* [Reduce] Fixed `miopenReduceTensor` silently truncating tensor lengths, strides, element counts
+  and byte spans above INT32_MAX, and single reductions of about 2^30 or more fp16, bf16 or int8
+  elements, which produced incorrect results. Such tensors are now rejected with
+  `miopenStatusBadParm`.
+
 ### Added
 * [Conv] Added gfx950 depthwise backward-weights (fp16/bf16) and gfx1250 depthwise
   kernels to the `ConvHipConv` solver (hipconv v0.3.1).
 * [Conv] Added gfx950 patch-embedding kernels (fp16/bf16), gfx950 TF32 kernels for dense
   and depthwise forward/backward-data and dense backward-weights, and gfx1250 depthwise
   backward-weights kernels (fp16/bf16) to the `ConvHipConv` solver (hipconv v0.3.2).
 
+### Changed
+* [Reduce] `miopenGetReductionWorkspaceSize` and `miopenGetReductionIndicesSize` now reject tensors
+  of rank greater than 6, matching `miopenReduceTensor`.
+* [Reduce] Invalid arguments to the reduction API (mismatched descriptors, unsupported types,
+  undersized workspace or indices buffers) now return `miopenStatusBadParm` instead of
+  `miopenStatusUnknownError`.
+
 ### Removed
 * [Conv] Removed gfx803 convolution solver `ConvBinWinogradRxSFused` and its kernel sources.
 * [Conv] Removed the gfx803 code paths and kernel sources from `ConvBinWinograd3x3U` and `ConvBinWinogradRxS`.

projects/miopen/src/reducetensor.cpp

@@ -31,7 +31,9 @@
 #include <miopen/reducetensor.hpp>
 #include <miopen/solver/legacy_ck_common.hpp>
 
+#include <algorithm>
 #include <cassert>
+#include <limits>
 #include <sstream>
 #include <string>
 #include <utility>
@@ -200,18 +202,6 @@ struct ReductionKernelConfigurator
     };
 };
 
-inline int GetIndicesTypeSize(miopenIndicesType_t t)
-{
-    switch(t)
-    {
-    case MIOPEN_32BIT_INDICES: return (4);
-    case MIOPEN_64BIT_INDICES: return (8);
-    case MIOPEN_16BIT_INDICES: return (2);
-    case MIOPEN_8BIT_INDICES: return (1);
-    }
-    MIOPEN_THROW("Unknown data type");
-}
-
 inline int GetDataTypeSize(miopenDataType_t t)
 {
     switch(t)
@@ -228,11 +218,25 @@ inline int GetDataTypeSize(miopenDataType_t t)
     case miopenInt64: break;
     };
 
-    MIOPEN_THROW("Only float, half, double, bfloat16, int8 data types are supported.");
+    MIOPEN_THROW(miopenStatusBadParm,
+                 "Only float, half, double, bfloat16, int8 data types are supported.");
 };
 
 }; // end of namespace detail
 
+namespace {
+
+int ToIntOrThrow(std::size_t value, const char* what)
+{
+    if(value > std::numeric_limits<int>::max())
+        MIOPEN_THROW(miopenStatusBadParm,
+                     std::string{"Reduction "} + what + " " + std::to_string(value) +
+                         " exceeds INT32_MAX, which is not supported.");
+    return static_cast<int>(value);
+}
+
+} // namespace
+
 namespace detailDynamic {
 
 static ck::DataTypeEnum_t mapDataTypeId(miopenDataType_t t)
@@ -253,7 +257,7 @@ static ck::DataTypeEnum_t mapDataTypeId(miopenDataType_t t)
     case miopenInt64: break;
     };
 
-    MIOPEN_THROW("Only float, half, double data type is supported.");
+    MIOPEN_THROW(miopenStatusBadParm, "Only float, half, double data type is supported.");
 };
 
 static ck::ReduceTensorOp_t mapReduceOpId(miopenReduceTensorOp_t t)
@@ -272,7 +276,7 @@ static ck::ReduceTensorOp_t mapReduceOpId(miopenReduceTensorOp_t t)
     case MIOPEN_REDUCE_TENSOR_NORM2: return ReduceTensorOp_t::NORM2;
     };
 
-    MIOPEN_THROW("Operation is not supported");
+    MIOPEN_THROW(miopenStatusBadParm, "Operation is not supported");
 };
 
 static std::string get_network_config_string_from_type_enums(miopenDataType_t TSrc,
@@ -374,6 +378,7 @@ static std::pair<bool, bool> get_padding_need(ReductionMethod_t reduceImpl,
         bool dst_need_padding = false;
         int copySliceLen;
         int reduceSizePerBlock;
+        size_t paddedReduceLen;
 
         switch(reduceImpl)
         {
@@ -387,8 +392,10 @@ static std::pair<bool, bool> get_padding_need(ReductionMethod_t reduceImpl,
         case Reduce_DirectWarpWise:
             copySliceLen = warpSize * tunable->GredAccessesPerThreadInWarp;
             src_need_padding =
-                (invariantLen < GridSize * BlockSize / warpSize || toReduceLen % copySliceLen > 0);
-            dst_need_padding = (invariantLen < GridSize * BlockSize / warpSize);
+                (invariantLen < static_cast<size_t>(GridSize) * BlockSize / warpSize ||
+                 toReduceLen % copySliceLen > 0);
+            dst_need_padding =
+                (invariantLen < static_cast<size_t>(GridSize) * BlockSize / warpSize);
             return std::make_pair(src_need_padding, dst_need_padding);
 
         case Reduce_BlockWise:
@@ -402,8 +409,11 @@ static std::pair<bool, bool> get_padding_need(ReductionMethod_t reduceImpl,
                 (((toReduceLen + BlkGroupSize - 1) / BlkGroupSize + copySliceLen - 1) /
                  copySliceLen) *
                 copySliceLen;
-            src_need_padding =
-                (toReduceLen < static_cast<size_t>(reduceSizePerBlock) * BlkGroupSize);
+            // the kernel computes the padded length as int; with a single output the number of
+            // blocks is not capped, so it can exceed INT32_MAX even when the tensor itself fits
+            paddedReduceLen = static_cast<size_t>(reduceSizePerBlock) * BlkGroupSize;
+            ToIntOrThrow(paddedReduceLen, "padded reduced length");
+            src_need_padding = (toReduceLen < paddedReduceLen);
             return std::make_pair(src_need_padding, dst_need_padding);
         };
 
@@ -447,33 +457,84 @@ ReduceTensorDescriptor::ReduceTensorDescriptor(miopenReduceTensorOp_t reduceTens
 {
     if(reduceTensorIndices == MIOPEN_REDUCE_TENSOR_FLATTENED_INDICES &&
        reduceTensorIndicesType != MIOPEN_32BIT_INDICES)
-        MIOPEN_THROW("Only int32 type is supported for ReduceTensor indices.");
+        MIOPEN_THROW(miopenStatusBadParm, "Only int32 type is supported for ReduceTensor indices.");
 };
 
-// This is WS requirement of the dynamic reduction.
-// We must enforce it especially when reduction is used internally.
-constexpr std::size_t workspaceAlignRequirementBytes = 64;
+namespace {
 
-// return the size of the workspace in bytes, so that the workspace buffer can be prepared by the
-// user
-std::size_t ReduceTensorDescriptor::GetWorkspaceSize(const Handle& handle,
-                                                     const TensorDescriptor& inDesc,
-                                                     const TensorDescriptor& outDesc) const
+void ValidateReduceDescriptors(const TensorDescriptor& inDesc, const TensorDescriptor& outDesc)
 {
+    // the legacy reduction kernels take at most 6 dimensions
+    constexpr std::size_t maxReduceRank = 6;
+
     const auto& inDescLengths  = inDesc.GetLengths();
     const auto& outDescLengths = outDesc.GetLengths();
 
     if(inDescLengths.size() != outDescLengths.size())
-        MIOPEN_THROW("The number of dimensions of the input and output tensor should match.");
+        MIOPEN_THROW(miopenStatusBadParm,
+                     "The number of dimensions of the input and output tensor should match.");
 
-    for(int i = 0; i < inDescLengths.size(); i++)
+    if(inDescLengths.size() > maxReduceRank)
+        MIOPEN_THROW(miopenStatusBadParm,
+                     "Invalid TensorDescriptor, at most number of dimensions of 6 is supported.");
+
+    for(std::size_t i = 0; i < inDescLengths.size(); i++)
     {
         if(outDescLengths[i] != 1 && outDescLengths[i] != inDescLengths[i])
         {
-            MIOPEN_THROW("The length of the output tensor dimension should either be 1 or be equal "
+            MIOPEN_THROW(miopenStatusBadParm,
+                         "The length of the output tensor dimension should either be 1 or be equal "
                          "to the length of the corresponding dimension of the input tensor.");
         }
-    };
+    }
+}
+
+// The kernels receive lengths and strides as int and address global memory through 32-bit byte
+// offsets (legacy CK buffer addressing), so both have to fit into int. The stride of a length-1
+// dimension never contributes to an offset, so it is not restricted.
+void ValidateReduceTensorFitsIntoInt(const TensorDescriptor& desc, const char* name)
+{
+    const auto& lengths = desc.GetLengths();
+    const auto& strides = desc.GetStrides();
+
+    // offset of the last element; the terms are capped, so the sum cannot wrap (a capped term
+    // already makes the span exceed INT32_MAX bytes)
+    std::size_t lastOffset = 0;
+    for(std::size_t i = 0; i < lengths.size(); i++)
+    {
+        if(lengths[i] == 1)
+            continue;
+
+        if(lengths[i] > std::numeric_limits<int>::max() ||
+           strides[i] > std::numeric_limits<int>::max())
+            MIOPEN_THROW(miopenStatusBadParm,
+                         std::string{"Reduction "} + name +
+                             " tensor has a length or stride exceeding INT32_MAX, which is not "
+                             "supported.");
+
+        lastOffset +=
+            std::min<std::size_t>((lengths[i] - 1) * strides[i], std::numeric_limits<int>::max());
+    }
+
+    if((lastOffset + 1) * detail::GetDataTypeSize(desc.GetType()) > std::numeric_limits<int>::max())
+        MIOPEN_THROW(miopenStatusBadParm,
+                     std::string{"Reduction "} + name +
+                         " tensor spans more than INT32_MAX bytes, which is not supported.");
+}
+
+} // namespace
+
+// This is WS requirement of the dynamic reduction.
+// We must enforce it especially when reduction is used internally.
+constexpr std::size_t workspaceAlignRequirementBytes = 64;
+
+// return the size of the workspace in bytes, so that the workspace buffer can be prepared by the
+// user
+std::size_t ReduceTensorDescriptor::GetWorkspaceSize(const Handle& handle,
+                                                     const TensorDescriptor& inDesc,
+                                                     const TensorDescriptor& outDesc) const
+{
+    ValidateReduceDescriptors(inDesc, outDesc);
 
     auto invariantLength = outDesc.GetElementSize();
     auto toReduceLength  = inDesc.GetElementSize() / invariantLength;
@@ -508,20 +569,7 @@ std::size_t ReduceTensorDescriptor::GetWorkspaceSize(const Handle& handle,
 std::size_t ReduceTensorDescriptor::GetIndicesSize(const TensorDescriptor& inDesc,
                                                    const TensorDescriptor& outDesc) const
 {
-    const auto& inDescLengths  = inDesc.GetLengths();
-    const auto& outDescLengths = outDesc.GetLengths();
-
-    if(inDescLengths.size() != outDescLengths.size())
-        MIOPEN_THROW("The number of dimensions of the input and output tensor should match.");
-
-    for(int i = 0; i < inDescLengths.size(); i++)
-    {
-        if(outDescLengths[i] != 1 && outDescLengths[i] != inDescLengths[i])
-        {
-            MIOPEN_THROW("The length of the output tensor dimension should either be 1 or be equal "
-                         "to the length of the corresponding dimension of the input tensor.");
-        }
-    };
+    ValidateReduceDescriptors(inDesc, outDesc);
 
     auto reduceIndicesOpt = this->reduceTensorIndices_;
     auto reduceOp         = this->reduceTensorOp_;
@@ -571,32 +619,24 @@ void ReduceTensorDescriptor::ReduceTensor(const Handle& handle,
         (reduceOp == MIOPEN_REDUCE_TENSOR_MIN || reduceOp == MIOPEN_REDUCE_TENSOR_MAX ||
          reduceOp == MIOPEN_REDUCE_TENSOR_AMAX);
 
-    if(inDescLengths.size() > 6)
-        MIOPEN_THROW("Invalid TensorDescriptor, at most number of dimensions of 6 is supported.");
+    ValidateReduceDescriptors(aDesc, cDesc);
 
     if(need_indices && (reduceIndicesType != MIOPEN_32BIT_INDICES))
-        MIOPEN_THROW("Only int32 type can be used for ReduceTensor indices.");
-
-    if(inDescLengths.size() != outDescLengths.size())
-        MIOPEN_THROW("The number of dimensions of the input and output tensor should match.");
+        MIOPEN_THROW(miopenStatusBadParm, "Only int32 type can be used for ReduceTensor indices.");
 
-    for(int i = 0; i < inDescLengths.size(); i++)
-    {
-        if(outDescLengths[i] != 1 && outDescLengths[i] != inDescLengths[i])
-        {
-            MIOPEN_THROW("The length of the output tensor dimension should either be 1 or be equal "
-                         "to the length of the corresponding dimension of the input tensor.");
-        }
-    };
+    // GetWorkspaceSize() deliberately does not check this: it is also used just to size buffers
+    // for reductions that are never run (e.g. RNN inference)
+    ValidateReduceTensorFitsIntoInt(aDesc, "input");
+    ValidateReduceTensorFitsIntoInt(cDesc, "output");
 
     std::size_t ws_sizeInBytes      = this->GetWorkspaceSize(handle, aDesc, cDesc);
     std::size_t indices_sizeInBytes = this->GetIndicesSize(aDesc, cDesc);
 
     if(ws_sizeInBytes > workspaceSizeInBytes)
-        MIOPEN_THROW("The workspace size allocated is not enough!");
+        MIOPEN_THROW(miopenStatusBadParm, "The workspace size allocated is not enough!");
 
     if(indices_sizeInBytes > indicesSizeInBytes)
-        MIOPEN_THROW("The indices size allocated is not enough!");
+        MIOPEN_THROW(miopenStatusBadParm, "The indices size allocated is not enough!");
 
     // invariantLength and toReduceLength are used to determine the kernel configuration
     const auto invariantLength = cDesc.GetElementSize();
@@ -614,7 +654,11 @@ void ReduceTensorDescriptor::ReduceTensor(const Handle& handle,
 
     const ReductionMethod_t reduceImpl =
         configurator.getReductionMethod(invariantLength, toReduceLength);
-    const int gridSize = configurator.getGridSize(invariantLength, toReduceLength);
+    // element counts are only bounded by the byte span checks above for non-overlapping tensors
+    ToIntOrThrow(invariantLength, "output element count");
+    const int origReduceLen = ToIntOrThrow(toReduceLength, "reduced element count");
+    const int gridSize =
+        ToIntOrThrow(configurator.getGridSize(invariantLength, toReduceLength), "grid size");
     const int blkGroupSize =
         (reduceImpl == Reduce_MultiBlock) ? static_cast<int>(gridSize / invariantLength) : 0;
 
@@ -633,7 +677,8 @@ void ReduceTensorDescriptor::ReduceTensor(const Handle& handle,
 
     if(toReduceDims.empty())
     {
-        MIOPEN_THROW("Invalid TensorDescriptor, at least one dimension of the input tensor should "
+        MIOPEN_THROW(miopenStatusBadParm,
+                     "Invalid TensorDescriptor, at least one dimension of the input tensor should "
                      "be reduced.");
     }
 
@@ -647,13 +692,12 @@ void ReduceTensorDescriptor::ReduceTensor(const Handle& handle,
                          : *reinterpret_cast<const float*>(beta);
 
     { // use dynamic reduction
-        const int origReduceLen = toReduceLength;
-
         int p_inLengths[6]  = {0};
         int p_inStrides[6]  = {0};
         int p_outLengths[6] = {0};
         int p_outStrides[6] = {0};
 
+        // the casts below are safe, see ValidateReduceTensorFitsIntoInt()
         int pos = 0;
         for(int i = 0; i < outDescLengths.size(); i++)
         {
@@ -674,7 +718,8 @@ void ReduceTensorDescriptor::ReduceTensor(const Handle& handle,
             if(outDescLengths[i] == 1)
             {
                 p_inLengths[pos] = static_cast<int>(inDescLengths[i]);
-                p_inStrides[pos] = static_cast<int>(inDescStrides[i]);
+                // a length-1 stride is unused and may not fit into int
+                p_inStrides[pos] = inDescLengths[i] == 1 ? 1 : static_cast<int>(inDescStrides[i]);
                 pos++;
             };
         };
@@ -831,8 +876,8 @@ void ReduceTensorDescriptor::ReduceTensor(const Handle& handle,
         if(useTwoCalls)
         {
             const auto toReduceLength_2 = blkGroupSize;
-            const int gridSize_2 =
-                static_cast<int>(configurator.getGridSize_2(invariantLength, toReduceLength_2));
+            const int gridSize_2        = ToIntOrThrow(
+                configurator.getGridSize_2(invariantLength, toReduceLength_2), "grid size");
             const std::vector<size_t> vgd2_2 = {
                 static_cast<size_t>(gridSize_2) * tunable->BlockSize, size_t{1}, size_t{1}};
             const auto reduceImpl2  = configurator.GetReductionMethod_2(toReduceLength_2);

projects/miopen/src/reducetensor_api.cpp

@@ -61,7 +61,7 @@ static void LogCmdRedux(const miopen::ReduceTensorDescriptor reduceTensorDesc,
         std::transform(inLens.begin(),
                        inLens.end(),
                        std::back_inserter(strInLens),
-                       [](const int x) { return std::to_string(x); });
+                       [](std::size_t x) { return std::to_string(x); });
         ss << " -D " << miopen::JoinStrings(strInLens, ",");
         ss << " -I " << reduceTensorDesc.reduceTensorIndices_;
         ss << " -N " << reduceTensorDesc.reduceTensorNanOpt_;