solve value-src dtype fault (#56613)

zoooo0820 · web-flow · commit 594177b8cdf0 · 2023-09-14T18:27:19.000+08:00
* solve value-src dtype fault

* support basic-indexing with different dtype, add UT

* skip amp convert when level is O0

* fix ut of amp
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -1665,14 +1665,17 @@ static PyObject* tensor_method__setitem_eager_tensor(TensorObject* self,
       // use inplace set_value_ operator
       if (value_tensor.initialized() &&
           (self->tensor.dtype() != value_tensor.dtype())) {
-        paddle::small_vector<std::vector<paddle::Tensor>,
-                             egr::kSlotSmallVectorSize>
-            tmps = {{self->tensor}, {value_tensor}};
-        auto amp_dtype = egr::GetAmpDestDtype("set_value", tmps);
-        self->tensor = egr::EagerAmpAutoCast(
-            self->tensor.name(), self->tensor, amp_dtype, "set_value");
-        value_tensor = egr::EagerAmpAutoCast(
-            value_tensor.name(), value_tensor, amp_dtype, "set_value");
+        if (egr::Controller::Instance().GetAMPLevel() !=
+            paddle::imperative::AmpLevel::O0) {
+          paddle::small_vector<std::vector<paddle::Tensor>,
+                               egr::kSlotSmallVectorSize>
+              tmps = {{self->tensor}, {value_tensor}};
+          auto amp_dtype = egr::GetAmpDestDtype("set_value", tmps);
+          self->tensor = egr::EagerAmpAutoCast(
+              self->tensor.name(), self->tensor, amp_dtype, "set_value");
+          value_tensor = egr::EagerAmpAutoCast(
+              value_tensor.name(), value_tensor, amp_dtype, "set_value");
+        }
         if (self->tensor.dtype() != value_tensor.dtype()) {
           value_tensor = cast_ad_func(value_tensor, self->tensor.dtype());
         }
diff --git a/python/paddle/base/variable_index.py b/python/paddle/base/variable_index.py
@@ -892,6 +892,7 @@ def _setitem_static(x, indices, values):
             attrs["shape"] = shape
 
         elif isinstance(values, Variable):
+            values = values.astype(dtype)
             inputs["ValueTensor"] = values
             value_tensor = values
 
@@ -967,6 +968,9 @@ def _setitem_static(x, indices, values):
         if not isinstance(values, Variable):
             values = paddle.assign(values).astype(transed_sub_tensor.dtype)
 
+        if values.dtype != transed_sub_tensor.dtype:
+            values = values.astype(transed_sub_tensor.dtype)
+
         if paddle.in_dynamic_mode():
             return transed_sub_tensor.index_put_(
                 adjusted_advanced_index, values
diff --git a/test/amp/test_amp_api.py b/test/amp/test_amp_api.py
@@ -281,7 +281,7 @@ def forward(self, x):
 class TestDy2STWithSetValue(AmpTestBase):
     def test_op_called_as_expected(self):
         expected_fp16_calls = {
-            "cast": 0,
+            "cast": 1,
             "layer_norm": 1,
             "scale": 3,
             "set_value": 1,
diff --git a/test/indexing/test_setitem.py b/test/indexing/test_setitem.py
@@ -60,6 +60,31 @@ def test_index_has_range(self):
 
         np.testing.assert_allclose(x.numpy(), np_data)
 
+    def test_src_value_with_different_dtype_1(self):
+        # basic-indexing, with set_value op
+        np_data = np.ones((3, 4, 5, 6), dtype='int32')
+        np_value = np.zeros((6,), dtype='float32')
+        x = paddle.to_tensor(np_data)
+        v = paddle.to_tensor(np_value)
+
+        np_data[0, 2, 3] = np_value
+        x[0, 2, 3] = v
+
+        np.testing.assert_allclose(x.numpy(), np_data)
+
+    def test_src_value_with_different_dtype_2(self):
+        # combined-indexing, with index_put op
+        np_data = np.ones((3, 4, 5, 6), dtype='float32')
+        np_value = np.zeros((6,), dtype='int64')
+
+        x = paddle.to_tensor(np_data)
+        v = paddle.to_tensor(np_value)
+
+        np_data[:, [1, 0], 3] = np_value
+        x[:, [1, 0], 3] = v
+
+        np.testing.assert_allclose(x.numpy(), np_data)
+
     def test_indexing_with_bool_list1(self):
         # test bool-list indexing when axes num less than x.rank
         np_data = np.arange(3 * 4 * 5 * 6).reshape((3, 4, 5, 6))
@@ -201,6 +226,46 @@ def test_index_has_range(self):
 
         np.testing.assert_allclose(res[0], np_data)
 
+    def test_src_value_with_different_dtype_1(self):
+        # basic-indexing, with set_value op
+        np_data = np.ones((3, 4, 5, 6), dtype='int32')
+        np_value = np.zeros((6,), dtype='float32')
+        np_data[0, 2, 3] = np_value
+
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            x = paddle.ones((3, 4, 5, 6), dtype='int32')
+            v = paddle.zeros((6,), dtype='float32')
+            y = _setitem_static(
+                x,
+                (0, 2, 3),
+                v,
+            )
+            res = self.exe.run(fetch_list=[y.name])
+
+        np.testing.assert_allclose(res[0], np_data)
+
+    def test_src_value_with_different_dtype_2(self):
+        # combined-indexing, with index_put op
+        np_data = np.ones((3, 4, 5, 6), dtype='float32')
+        np_value = np.zeros((6,), dtype='int64')
+        np_data[:, [1, 0], 3] = np_value
+
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            x = paddle.ones((3, 4, 5, 6), dtype='float32')
+            v = paddle.zeros((6,), dtype='int64')
+            y = _setitem_static(
+                x,
+                (slice(None, None), [1, 0], 3),
+                v,
+            )
+            res = self.exe.run(fetch_list=[y.name])
+
+        np.testing.assert_allclose(res[0], np_data)
+
     def test_indexing_with_bool_list1(self):
         # test bool-list indexing when axes num less than x.rank
         np_data = np.arange(3 * 4 * 5 * 6).reshape((3, 4, 5, 6))
