Address review: rank guards, overflow-safe ceil-div, weighted grad test

Author: vraspar

onnxruntime/contrib_ops/cpu/quantization/matmul_bnb4.cc

@@ -58,8 +58,8 @@ Status MatMulBnb4::Compute(OpKernelContext* ctx) const {
                            "K * N would overflow int64: K=", K_, ", N=", N_);
   }
   const int64_t numel = K_ * N_;
-  const int64_t expected_b_quant_size = (numel + 1) / 2;
-  const int64_t expected_absmax_size = (numel + block_size_ - 1) / block_size_;
+  const int64_t expected_b_quant_size = ((numel - 1) / 2) + 1;
+  const int64_t expected_absmax_size = ((numel - 1) / block_size_) + 1;
 
   if (b_quant->Shape().Size() < expected_b_quant_size) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,

onnxruntime/contrib_ops/cuda/quantization/matmul_bnb4.cc

@@ -59,8 +59,8 @@ Status MatMulBnb4<T>::ComputeInternal(OpKernelContext* ctx) const {
                            "K * N would overflow int64: K=", K_, ", N=", N_);
   }
   const int64_t numel = K_ * N_;
-  const int64_t expected_b_quant_size = (numel + 1) / 2;
-  const int64_t expected_absmax_size = (numel + block_size_ - 1) / block_size_;
+  const int64_t expected_b_quant_size = ((numel - 1) / 2) + 1;
+  const int64_t expected_absmax_size = ((numel - 1) / block_size_) + 1;
 
   if (b_quant->Shape().Size() < expected_b_quant_size) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,

orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc

@@ -76,5 +76,24 @@ TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLarge) {
   test.Run(OpTester::ExpectResult::kExpectFailure, "out of range");
 }
 
+TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLargeWithWeights) {
+  OpTester test("SoftmaxCrossEntropyLossGrad", 1, onnxruntime::kMSDomain);
+  test.AddAttribute("reduction", std::string("mean"));
+  test.AddAttribute("ignore_index", static_cast<int64_t>(-1));
+
+  std::vector<float> dY_data = {1.0f};
+  std::vector<float> log_prob_data(3 * 5, -1.6094f);
+  std::vector<int64_t> index_data = {0, 5, 2};  // 5 is out of range [0, 5)
+  std::vector<float> weight_data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+
+  test.AddInput<float>("dY", {}, dY_data);
+  test.AddInput<float>("log_prob", {3, 5}, log_prob_data);
+  test.AddInput<int64_t>("index", {3}, index_data);
+  test.AddInput<float>("weight", {5}, weight_data);
+  test.AddOutput<float>("dX", {3, 5}, std::vector<float>(15, 0.0f));
+
+  test.Run(OpTester::ExpectResult::kExpectFailure, "out of range");
+}
+
 }  // namespace test
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc

@@ -56,15 +56,17 @@ void GetNDCFromLogitAndLabelShape(const TensorShape& logit_shape, const TensorSh
 void VerifyLogitWeightAndLabelShape(const TensorShape& logit_shape,
                                     const TensorShape& label_shape,
                                     const TensorShape* weight_shape) {
+  const size_t label_dims = label_shape.NumDimensions();
+  ORT_ENFORCE(label_dims >= 1, "label must be at least 1-D.");
+  ORT_ENFORCE(logit_shape.NumDimensions() >= 2, "logit must be at least 2-D.");
+  ORT_ENFORCE(logit_shape.NumDimensions() == label_dims + 1,
+              "logit_shape must be (1 + label_shape)");
+
   ORT_ENFORCE(nullptr == weight_shape || 1 == weight_shape->NumDimensions(), "Weights tensor is not 1-D.");
   ORT_ENFORCE(nullptr == weight_shape || (*weight_shape)[0] == logit_shape[1],
               "Weight tensor size (", (weight_shape ? (*weight_shape)[0] : 0),
               ") must equal the number of classes (", logit_shape[1], ")");
 
-  const size_t label_dims = label_shape.NumDimensions();
-  ORT_ENFORCE(logit_shape.NumDimensions() == label_dims + 1,
-              "logit_shape must be (1 + label_shape)");
-
   ORT_ENFORCE(label_shape[0] == logit_shape[0], "The shape of logit and label does not match");
 
   if (label_dims >= 2) {