refine eager_math_op_patch.cc convert disttensor (#59062)

wanghuancoder · web-flow · commit 2811d1f7cb9a · 2023-11-17T10:21:24.000+08:00
diff --git a/paddle/fluid/pybind/eager_math_op_patch.cc b/paddle/fluid/pybind/eager_math_op_patch.cc
@@ -247,6 +247,11 @@ static PyObject* tensor__add__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -348,6 +353,11 @@ static PyObject* tensor__sub__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -445,6 +455,11 @@ static PyObject* tensor__rsub__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -548,6 +563,11 @@ static PyObject* tensor__mul__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -579,11 +599,6 @@ static PyObject* tensor__mul__method(TensorObject* self,
     }
   }
 
-  const phi::distributed::ProcessMesh* mesh = nullptr;
-  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
-    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
-  }
-
   // 4. calculation
   VLOG(6) << "Calling multiply_ad_func in tensor__mul__method";
   {
@@ -659,6 +674,11 @@ static PyObject* tensor__div__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -779,6 +799,11 @@ static PyObject* tensor__rdiv__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -897,6 +922,11 @@ static PyObject* tensor__gt__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -987,6 +1017,11 @@ static PyObject* tensor__ge__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -1078,6 +1113,11 @@ static PyObject* tensor__mod__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -1168,6 +1208,11 @@ static PyObject* tensor__matmul__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -1276,6 +1321,11 @@ static PyObject* tensor__lt__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -1366,6 +1416,11 @@ static PyObject* tensor__le__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -1457,6 +1512,11 @@ static PyObject* tensor__floordiv__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -1546,6 +1606,11 @@ static PyObject* tensor__pow__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -1638,6 +1703,11 @@ static PyObject* tensor__rpow__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -1728,6 +1798,11 @@ static PyObject* tensor__ne__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();
@@ -1818,6 +1893,11 @@ static PyObject* tensor__eq__method(TensorObject* self,
     }
   }
 
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+    ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+  }
+
   // 3. promote types or unify right var type to left var
   phi::DataType lhs_dtype = self_tensor.dtype();
   phi::DataType rhs_dtype = other_tensor.dtype();

Original file line number	Diff line number	Diff line change
`@@ -247,6 +247,11 @@ static PyObject* tensor__add__method(TensorObject* self,`
`247`	`247`	`}`
`248`	`248`	`}`
`249`	`249`
	`250`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`251`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`252`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`253`	`+ }`
	`254`	`+`
`250`	`255`	`// 3. promote types or unify right var type to left var`
`251`	`256`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`252`	`257`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -348,6 +353,11 @@ static PyObject* tensor__sub__method(TensorObject* self,`
`348`	`353`	`}`
`349`	`354`	`}`
`350`	`355`
	`356`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`357`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`358`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`359`	`+ }`
	`360`	`+`
`351`	`361`	`// 3. promote types or unify right var type to left var`
`352`	`362`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`353`	`363`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -445,6 +455,11 @@ static PyObject* tensor__rsub__method(TensorObject* self,`
`445`	`455`	`}`
`446`	`456`	`}`
`447`	`457`
	`458`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`459`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`460`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`461`	`+ }`
	`462`	`+`
`448`	`463`	`// 3. promote types or unify right var type to left var`
`449`	`464`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`450`	`465`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -548,6 +563,11 @@ static PyObject* tensor__mul__method(TensorObject* self,`
`548`	`563`	`}`
`549`	`564`	`}`
`550`	`565`
	`566`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`567`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`568`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`569`	`+ }`
	`570`	`+`
`551`	`571`	`// 3. promote types or unify right var type to left var`
`552`	`572`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`553`	`573`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -579,11 +599,6 @@ static PyObject* tensor__mul__method(TensorObject* self,`
`579`	`599`	`}`
`580`	`600`	`}`
`581`	`601`
`582`		`- const phi::distributed::ProcessMesh* mesh = nullptr;`
`583`		`- if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
`584`		`- ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
`585`		`- }`
`586`		`-`
`587`	`602`	`// 4. calculation`
`588`	`603`	`VLOG(6) << "Calling multiply_ad_func in tensor__mul__method";`
`589`	`604`	`{`
`@@ -659,6 +674,11 @@ static PyObject* tensor__div__method(TensorObject* self,`
`659`	`674`	`}`
`660`	`675`	`}`
`661`	`676`
	`677`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`678`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`679`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`680`	`+ }`
	`681`	`+`
`662`	`682`	`// 3. promote types or unify right var type to left var`
`663`	`683`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`664`	`684`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -779,6 +799,11 @@ static PyObject* tensor__rdiv__method(TensorObject* self,`
`779`	`799`	`}`
`780`	`800`	`}`
`781`	`801`
	`802`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`803`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`804`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`805`	`+ }`
	`806`	`+`
`782`	`807`	`// 3. promote types or unify right var type to left var`
`783`	`808`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`784`	`809`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -897,6 +922,11 @@ static PyObject* tensor__gt__method(TensorObject* self,`
`897`	`922`	`}`
`898`	`923`	`}`
`899`	`924`
	`925`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`926`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`927`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`928`	`+ }`
	`929`	`+`
`900`	`930`	`// 3. promote types or unify right var type to left var`
`901`	`931`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`902`	`932`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -987,6 +1017,11 @@ static PyObject* tensor__ge__method(TensorObject* self,`
`987`	`1017`	`}`
`988`	`1018`	`}`
`989`	`1019`
	`1020`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1021`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1022`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1023`	`+ }`
	`1024`	`+`
`990`	`1025`	`// 3. promote types or unify right var type to left var`
`991`	`1026`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`992`	`1027`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -1078,6 +1113,11 @@ static PyObject* tensor__mod__method(TensorObject* self,`
`1078`	`1113`	`}`
`1079`	`1114`	`}`
`1080`	`1115`
	`1116`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1117`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1118`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1119`	`+ }`
	`1120`	`+`
`1081`	`1121`	`// 3. promote types or unify right var type to left var`
`1082`	`1122`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`1083`	`1123`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -1168,6 +1208,11 @@ static PyObject* tensor__matmul__method(TensorObject* self,`
`1168`	`1208`	`}`
`1169`	`1209`	`}`
`1170`	`1210`
	`1211`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1212`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1213`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1214`	`+ }`
	`1215`	`+`
`1171`	`1216`	`// 3. promote types or unify right var type to left var`
`1172`	`1217`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`1173`	`1218`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -1276,6 +1321,11 @@ static PyObject* tensor__lt__method(TensorObject* self,`
`1276`	`1321`	`}`
`1277`	`1322`	`}`
`1278`	`1323`
	`1324`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1325`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1326`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1327`	`+ }`
	`1328`	`+`
`1279`	`1329`	`// 3. promote types or unify right var type to left var`
`1280`	`1330`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`1281`	`1331`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -1366,6 +1416,11 @@ static PyObject* tensor__le__method(TensorObject* self,`
`1366`	`1416`	`}`
`1367`	`1417`	`}`
`1368`	`1418`
	`1419`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1420`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1421`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1422`	`+ }`
	`1423`	`+`
`1369`	`1424`	`// 3. promote types or unify right var type to left var`
`1370`	`1425`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`1371`	`1426`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -1457,6 +1512,11 @@ static PyObject* tensor__floordiv__method(TensorObject* self,`
`1457`	`1512`	`}`
`1458`	`1513`	`}`
`1459`	`1514`
	`1515`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1516`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1517`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1518`	`+ }`
	`1519`	`+`
`1460`	`1520`	`// 3. promote types or unify right var type to left var`
`1461`	`1521`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`1462`	`1522`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -1546,6 +1606,11 @@ static PyObject* tensor__pow__method(TensorObject* self,`
`1546`	`1606`	`}`
`1547`	`1607`	`}`
`1548`	`1608`
	`1609`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1610`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1611`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1612`	`+ }`
	`1613`	`+`
`1549`	`1614`	`// 3. promote types or unify right var type to left var`
`1550`	`1615`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`1551`	`1616`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -1638,6 +1703,11 @@ static PyObject* tensor__rpow__method(TensorObject* self,`
`1638`	`1703`	`}`
`1639`	`1704`	`}`
`1640`	`1705`
	`1706`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1707`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1708`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1709`	`+ }`
	`1710`	`+`
`1641`	`1711`	`// 3. promote types or unify right var type to left var`
`1642`	`1712`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`1643`	`1713`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -1728,6 +1798,11 @@ static PyObject* tensor__ne__method(TensorObject* self,`
`1728`	`1798`	`}`
`1729`	`1799`	`}`
`1730`	`1800`
	`1801`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1802`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1803`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1804`	`+ }`
	`1805`	`+`
`1731`	`1806`	`// 3. promote types or unify right var type to left var`
`1732`	`1807`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`1733`	`1808`	`phi::DataType rhs_dtype = other_tensor.dtype();`
`@@ -1818,6 +1893,11 @@ static PyObject* tensor__eq__method(TensorObject* self,`
`1818`	`1893`	`}`
`1819`	`1894`	`}`
`1820`	`1895`
	`1896`	`+ const phi::distributed::ProcessMesh* mesh = nullptr;`
	`1897`	`+ if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {`
	`1898`	`+ ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);`
	`1899`	`+ }`
	`1900`	`+`
`1821`	`1901`	`// 3. promote types or unify right var type to left var`
`1822`	`1902`	`phi::DataType lhs_dtype = self_tensor.dtype();`
`1823`	`1903`	`phi::DataType rhs_dtype = other_tensor.dtype();`