add GetDataTypeOfVar

paddle/fluid/framework/operator.cc

@@ -35,6 +35,17 @@ std::vector<std::tuple<platform::Place, LibraryType>> kKernelPriority = {
     std::make_tuple(platform::CPUPlace(), LibraryType::kPlain),
 };
 
+proto::VarType::Type GetDataTypeOfVar(const Variable* var) {
+  if (var->IsType<framework::LoDTensor>()) {
+    return framework::ToDataType(var->Get<framework::LoDTensor>().type());
+  } else if (var->IsType<framework::SelectedRows>()) {
+    return framework::ToDataType(
+        var->Get<framework::SelectedRows>().value().type());
+  } else {
+    PADDLE_THROW("Var should be LoDTensor or SelectedRows");
+  }
+}
+
 static DDim GetDims(const Scope& scope, const std::string& name) {
   Variable* var = scope.FindVar(name);
   if (var == nullptr) {

paddle/fluid/framework/operator.h

@@ -61,6 +61,8 @@ inline std::string GradVarName(const std::string& var_name) {
   return var_name + kGradVarSuffix;
 }
 
+proto::VarType::Type GetDataTypeOfVar(const Variable* var);
+
 class OperatorBase;
 class ExecutionContext;
 

paddle/fluid/operators/lookup_table_op.cc

@@ -18,22 +18,6 @@ limitations under the License. */
 namespace paddle {
 namespace operators {
 
-static inline framework::OpKernelType ExpectedKernelType(
-    const framework::ExecutionContext& ctx) {
-  auto* table_var = ctx.InputVar("W");
-  if (table_var->IsType<LoDTensor>()) {
-    return framework::OpKernelType(
-        framework::ToDataType(table_var->Get<LoDTensor>().type()),
-        ctx.device_context());
-  } else if (table_var->IsType<SelectedRows>()) {
-    return framework::OpKernelType(
-        framework::ToDataType(table_var->Get<SelectedRows>().value().type()),
-        ctx.device_context());
-  } else {
-    PADDLE_THROW("W should be LoDTensor or SelectedRows");
-  }
-}
-
 class LookupTableOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
@@ -67,7 +51,8 @@ class LookupTableOp : public framework::OperatorWithKernel {
  protected:
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext& ctx) const override {
-    return ExpectedKernelType(ctx);
+    auto data_type = framework::GetDataTypeOfVar(ctx.InputVar("W"));
+    return framework::OpKernelType(data_type, ctx.device_context());
   }
 };
 
@@ -138,7 +123,8 @@ class LookupTableOpGrad : public framework::OperatorWithKernel {
  protected:
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext& ctx) const override {
-    return ExpectedKernelType(ctx);
+    auto data_type = framework::GetDataTypeOfVar(ctx.InputVar("W"));
+    return framework::OpKernelType(data_type, ctx.device_context());
   }
 };
 

paddle/fluid/operators/sgd_op.cc

@@ -43,19 +43,8 @@ class SGDOp : public framework::OperatorWithKernel {
  protected:
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext& ctx) const override {
-    auto* table_var = ctx.InputVar("Param");
-    if (table_var->IsType<framework::LoDTensor>()) {
-      return framework::OpKernelType(
-          framework::ToDataType(table_var->Get<framework::LoDTensor>().type()),
-          ctx.device_context());
-    } else if (table_var->IsType<framework::SelectedRows>()) {
-      return framework::OpKernelType(
-          framework::ToDataType(
-              table_var->Get<framework::SelectedRows>().value().type()),
-          ctx.device_context());
-    } else {
-      PADDLE_THROW("Param should be LoDTensor or SelectedRows");
-    }
+    auto data_type = framework::GetDataTypeOfVar(ctx.InputVar("Param"));
+    return framework::OpKernelType(data_type, ctx.device_context());
   }
 };