Albert/fix minimum and div core dump (#283)

* Fix(fix_minimum_and_div_core_dump)

Author: real-zhangzhe

musa_ext/kernels/math/musa_minimum_op.cc

@@ -54,7 +54,6 @@ using namespace tensorflow;
       ::tensorflow::musa::MusaMinimumOp<type>);
 
 REGISTER_MUSA_MIN(float);
-REGISTER_MUSA_MIN(double);
 REGISTER_MUSA_MIN(int32);
 REGISTER_MUSA_MIN(int64);
 REGISTER_MUSA_MIN(Eigen::half);  // FP16

musa_ext/kernels/math/musa_multiply_op.cc

@@ -318,6 +318,7 @@ class MusaMultiplyOp : public MusaOpKernel {
       MusaMultiplyOp<TYPE>);
 
 REGISTER_MUSA_MULTIPLY(float);
+REGISTER_MUSA_MULTIPLY(double);
 REGISTER_MUSA_MULTIPLY(Eigen::half);
 REGISTER_MUSA_MULTIPLY(bfloat16);
 REGISTER_MUSA_MULTIPLY(int32);

musa_ext/kernels/math/musa_realdiv_op.cc

@@ -56,7 +56,7 @@ class MusaRealDivOp : public MusaOpKernel {
     mTensor t_out = CreateMTensor(*out, format_);
 
     ::musa::dnn::Binary op;
-    op.SetMode(::musa::dnn::Binary::Mode::DIV);
+    op.SetMode(::musa::dnn::Binary::Mode::TRUEDIV);
 
     auto status = op.Run(handle, t_out, t_dividend, t_divisor);
     OP_REQUIRES(ctx, status == ::musa::dnn::Status::SUCCESS,

musa_ext/kernels/utils_op.cc

@@ -74,16 +74,23 @@ mTensor CreateMTensor(const Tensor& t, mFormat format) {
   // Reuse TensorFlow's shape storage directly instead of copying dims into a
   // temporary vector. For small elementwise ops this shaves a bit of host-side
   // wrapper overhead.
-  const int64_t* dims =
-      reinterpret_cast<const int64_t*>(dims_raw.data());
+  const int64_t* dims = reinterpret_cast<const int64_t*>(dims_raw.data());
 
   if (rank >= 4) {
     rst.SetFormat(format);
   } else {
     rst.SetFormat(mFormat::NCHW);
   }
 
-  rst.SetNdInfo(rank, dims);
+  // muDNN does not accept rank-0 (scalar) tensors; represent scalars as a
+  // rank-1 [1] view so elementwise kernels such as Minimum/Multiply/Add can
+  // still execute on MUSA instead of tripping the error path.
+  if (rank == 0) {
+    const int64_t one = 1;
+    rst.SetNdInfo(1, &one);
+  } else {
+    rst.SetNdInfo(rank, dims);
+  }
   return rst;
 }
 

test/ops/multiply_op_test.py

@@ -32,30 +32,30 @@ def _test_multiply(self, shape_x, shape_y, dtype, rtol=1e-5, atol=1e-8):
     else:
       x_np = np.random.uniform(-1, 1, size=shape_x).astype(dtype.as_numpy_dtype)
       y_np = np.random.uniform(-1, 1, size=shape_y).astype(dtype.as_numpy_dtype)
-    
+
     x = tf.constant(x_np, dtype=dtype)
     y = tf.constant(y_np, dtype=dtype)
-    
+
     # Test on CPU
     with tf.device('/CPU:0'):
       cpu_result = tf.multiply(x, y)
-    
+
     # Test on MUSA
     with tf.device('/device:MUSA:0'):
       musa_result = tf.multiply(x, y)
-    
+
     # Compare results
     if dtype in [tf.float16, tf.bfloat16]:
       cpu_result_f32 = tf.cast(cpu_result, tf.float32)
       musa_result_f32 = tf.cast(musa_result, tf.float32)
-      self.assertAllClose(cpu_result_f32.numpy(), 
+      self.assertAllClose(cpu_result_f32.numpy(),
                          musa_result_f32.numpy(),
-                         rtol=rtol, 
+                         rtol=rtol,
                          atol=atol)
     else:
-      self.assertAllClose(cpu_result.numpy(), 
+      self.assertAllClose(cpu_result.numpy(),
                          musa_result.numpy(),
-                         rtol=rtol, 
+                         rtol=rtol,
                          atol=atol)
 
   def testMultiplyBasic(self):
@@ -102,12 +102,12 @@ def testMultiplyZeroValues(self):
     y_data = [[1.0, 0.0], [0.0, 3.0]]
     x = tf.constant(x_data, dtype=tf.float32)
     y = tf.constant(y_data, dtype=tf.float32)
-    
+
     expected = [[0.0, 0.0], [0.0, 0.0]]
-    
+
     with tf.device('/device:MUSA:0'):
       result = tf.multiply(x, y)
-    
+
     self.assertAllClose(result.numpy(), expected)
 
 
@@ -125,6 +125,19 @@ def testMultiplyFastPathHotShapes(self):
         with self.subTest(shape_x=shape_x, shape_y=shape_y, dtype=dtype):
           self._test_multiply(shape_x, shape_y, dtype, rtol=rtol, atol=atol)
 
+  def testMultiplyFloat64StaysOnMusa(self):
+    x = tf.constant([2.0], dtype=tf.float64)
+    y = tf.constant([3.0], dtype=tf.float64)
+
+    with tf.device('/device:MUSA:0'):
+      result = tf.multiply(x, y)
+
+    self.assertIn(
+        'MUSA', result.device,
+        msg=f"Mul(float64) output landed on {result.device!r} instead of MUSA")
+    self.assertEqual(result.dtype, tf.float64)
+    self.assertAllClose(result.numpy(), [6.0], rtol=1e-12, atol=1e-12)
+
 
 if __name__ == "__main__":
-  tf.test.main()
+  tf.test.main()

test/ops/realdiv_op_test.py

@@ -6,6 +6,11 @@
 
 class RealDivOpTest(MUSATestCase):
 
+  def _assert_tensor_on_musa(self, tensor, op_name):
+    self.assertIn(
+        'MUSA', tensor.device,
+        msg=f"{op_name} output landed on {tensor.device!r} instead of MUSA")
+
   def _test_realdiv(self, shape_x, shape_y, dtype, rtol=1e-5, atol=1e-8):
     np_dtype = dtype.as_numpy_dtype
 
@@ -113,6 +118,32 @@ def testRealDivLargeTensors(self):
     for dtype in [tf.float32]:
       self._test_realdiv([1024, 1024, 2], [1024, 1024, 2], dtype)
 
+  def testRealDivFloat64ScalarStaysOnMusa(self):
+    with tf.device('/device:MUSA:0'):
+      dividend = tf.constant(100.0, dtype=tf.float64)
+      divisor = tf.constant(1000.0, dtype=tf.float64)
+      result = tf.math.truediv(dividend, divisor)
+
+    self._assert_tensor_on_musa(result, 'RealDiv')
+    self.assertEqual(result.dtype, tf.float64)
+    self.assertEqual(result.shape.rank, 0)
+    self.assertAllClose(result.numpy(), 0.1, rtol=1e-12, atol=1e-12)
+
+  def testInt64TrueDivInsideTfFunctionStaysOnMusa(self):
+    with tf.device('/device:MUSA:0'):
+      step = tf.Variable(100, dtype=tf.int64, name='iterations')
+
+    @tf.function
+    def warmup_scale(step_var):
+      return step_var / 1000
+
+    result = warmup_scale(step)
+
+    self._assert_tensor_on_musa(result, 'tf.function truediv')
+    self.assertEqual(result.dtype, tf.float64)
+    self.assertEqual(result.shape.rank, 0)
+    self.assertAllClose(result.numpy(), 0.1, rtol=1e-12, atol=1e-12)
+
 
 if __name__ == "__main__":
   tf.test.main()