Merge pull request #30 from daneill/paddlebox

support embex 2 and 4
jack603047588 · Feb 11, 2022 · 78f584c · 78f584c
2 parents a283817 + 8b0ec8e
commit 78f584c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/paddle/fluid/framework/fleet/box_wrapper.cc b/paddle/fluid/framework/fleet/box_wrapper.cc
@@ -443,6 +443,8 @@ void BoxWrapper::PullSparse(const paddle::platform::Place& place,
   CheckEmbedSizeIsValid(hidden_size - cvm_offset_, expand_embed_dim);
   switch (embedx_dim_) {
     EMBEDX_CASE(0, PULLSPARSE_CASE(0););
+    EMBEDX_CASE(2, PULLSPARSE_CASE(0););
+    EMBEDX_CASE(4, PULLSPARSE_CASE(0););
     EMBEDX_CASE(8, PULLSPARSE_CASE(0); PULLSPARSE_CASE(1); PULLSPARSE_CASE(2);
                 PULLSPARSE_CASE(3); PULLSPARSE_CASE(4); PULLSPARSE_CASE(5);
                 PULLSPARSE_CASE(6); PULLSPARSE_CASE(7); PULLSPARSE_CASE(8);
@@ -514,6 +516,8 @@ void BoxWrapper::PushSparseGrad(const paddle::platform::Place& place,
   CheckEmbedSizeIsValid(hidden_size - cvm_offset_, expand_embed_dim);
   switch (embedx_dim_) {
     EMBEDX_CASE(0, PUSHSPARSE_CASE(0););
+    EMBEDX_CASE(2, PUSHSPARSE_CASE(0););
+    EMBEDX_CASE(4, PUSHSPARSE_CASE(0););
     EMBEDX_CASE(8, PUSHSPARSE_CASE(0); PUSHSPARSE_CASE(1); PUSHSPARSE_CASE(2);
                 PUSHSPARSE_CASE(3); PUSHSPARSE_CASE(4); PUSHSPARSE_CASE(5);
                 PUSHSPARSE_CASE(6); PUSHSPARSE_CASE(7); PUSHSPARSE_CASE(8);

diff --git a/paddle/fluid/framework/fleet/box_wrapper.cu b/paddle/fluid/framework/fleet/box_wrapper.cu
@@ -1254,6 +1254,8 @@ void BoxWrapper::CopyForPull(const paddle::platform::Place& place,
 
   switch (embedx_dim_) {
     EMBEDX_CASE(0, EXPAND_EMBED_PULL_CASE(0););
+    EMBEDX_CASE(2, EXPAND_EMBED_PULL_CASE(0););
+    EMBEDX_CASE(4, EXPAND_EMBED_PULL_CASE(0););
     EMBEDX_CASE(8, EXPAND_EMBED_PULL_CASE(0); EXPAND_EMBED_PULL_SHARE(1);
                 EXPAND_EMBED_PULL_SHARE(2); EXPAND_EMBED_PULL_SHARE(3);
                 EXPAND_EMBED_PULL_SHARE(4); EXPAND_EMBED_PULL_SHARE(5);
@@ -1553,6 +1555,8 @@ void BoxWrapper::CopyForPush(const paddle::platform::Place& place,
   } break
   switch (embedx_dim_) {
     EMBEDX_CASE(0, EXPAND_EMBED_PUSH_CASE(0););
+    EMBEDX_CASE(2, EXPAND_EMBED_PUSH_CASE(0););
+    EMBEDX_CASE(4, EXPAND_EMBED_PUSH_CASE(0););
     EMBEDX_CASE(8, EXPAND_EMBED_PUSH_CASE(0); EXPAND_EMBED_PUSH_SHARE(1);
                 EXPAND_EMBED_PUSH_SHARE(2); EXPAND_EMBED_PUSH_SHARE(3);
                 EXPAND_EMBED_PUSH_SHARE(4); EXPAND_EMBED_PUSH_SHARE(5);