[NPU] Fix aclnnInplaceMuls parameter bug in dropout kernel (#1413)

Author: warrentdrew

backends/npu/kernels/dropout_kernel.cc

@@ -402,13 +402,21 @@ void DropoutRawKernel(const Context& dev_ctx,
     }
 
     if (!is_upscale) {
-      phi::Scalar revert_scale = static_cast<T>(1.0 - dropout_prob);
-      EXEC_NPU_CMD(aclnnInplaceMuls, dev_ctx, *out, revert_scale);
+      auto revert_scale = static_cast<T>(1.0 - dropout_prob);
+      aclDataType acl_data_type = ConvertToNpuDtype(x.dtype());
+      static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar);
+      aclScalar* acl_scalar_revert_scale =
+          aclCreateScalar(&revert_scale, acl_data_type);
+      EXEC_NPU_CMD(aclnnInplaceMuls, dev_ctx, *out, acl_scalar_revert_scale);
     }
   } else {
     if (!is_upscale) {
-      phi::Scalar down_scale = static_cast<T>(1.0 - dropout_prob);
-      EXEC_NPU_CMD(aclnnMuls, dev_ctx, x, down_scale, *out);
+      auto down_scale = static_cast<T>(1.0 - dropout_prob);
+      aclDataType acl_data_type = ConvertToNpuDtype(x.dtype());
+      static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar);
+      aclScalar* acl_scalar_down_scale =
+          aclCreateScalar(&down_scale, acl_data_type);
+      EXEC_NPU_CMD(aclnnMuls, dev_ctx, x, acl_scalar_down_scale, *out);
       return;
     }
     TensorCopy(dev_ctx, x, false, out);
@@ -565,8 +573,12 @@ void DropoutGradRawKernel(const Context& dev_ctx,
   }
 
   if (!is_upscale) {
-    phi::Scalar revert_scale = static_cast<T>(1.0 - dropout_prob);
-    EXEC_NPU_CMD(aclnnInplaceMuls, dev_ctx, *dx, revert_scale);
+    auto revert_scale = static_cast<T>(1.0 - dropout_prob);
+    aclDataType acl_data_type = ConvertToNpuDtype(dx->dtype());
+    static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar);
+    aclScalar* acl_scalar_revert_scale =
+        aclCreateScalar(&revert_scale, acl_data_type);
+    EXEC_NPU_CMD(aclnnInplaceMuls, dev_ctx, *dx, acl_scalar_revert_scale);
   }
   return;
 }

backends/npu/tests/unittests/test_dropout_op_npu.py

@@ -148,6 +148,28 @@ def setUp(self):
         }
 
 
+class TestDropoutModeDownOp1Fp16(TestDropoutOp):
+    # the dropout_prob is 0.2
+    def init_dtype(self):
+        self.dtype = np.float16
+
+    def setUp(self):
+        self.op_type = "dropout"
+        self.set_npu()
+        self.init_dtype()
+        self.inputs = {"X": np.random.random((32, 64)).astype(self.dtype)}
+        self.attrs = {
+            "dropout_prob": 0.2,
+            "fix_seed": True,
+            "is_test": False,
+            "dropout_implementation": "downgrade_in_infer",
+        }
+        self.outputs = {
+            "Out": np.zeros((32, 64)).astype("float16"),
+            "Mask": convert_to_npu_mask(np.zeros((32, 64)).astype("uint8")),
+        }
+
+
 class TestDropoutOp2(TestDropoutOp):
     # the dropout_prob is 1.0
     def setUp(self):