fix(ascend): adopt PR #63/#60 master API — GetWorkspacePool/Ensure rename + drop registry.h (SFINAE autodetect)

Author: zhangyue

src/ascend/add/kernel.h

@@ -65,7 +65,7 @@ class Operator<Add, Device::Type::kAscend> : public Add {
       aclSetOutputTensorAddr(executor_, 0, t_out, out.data());
     }
 
-    auto& arena = ascend::workspacePool().ensure(stream, ws_size_);
+    auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_);
     aclnnAdd(arena.buf, ws_size_, executor_, stream);
   }
 

src/ascend/add_rms_norm/kernel.h

@@ -7,9 +7,9 @@
 #include "aclnn/aclnn_base.h"
 #include "aclnn_add.h"
 #include "aclnn_rms_norm.h"
-#include "ascend/add_rms_norm/registry.h"
 #include "ascend/common.h"
 #include "ascend/workspace_pool_.h"
+#include "base/add_rms_norm.h"
 #include "operator.h"
 
 namespace infini::ops {
@@ -74,12 +74,12 @@ class Operator<AddRmsNorm, Device::Type::kAscend, 0> : public AddRmsNorm {
       aclSetInputTensorAddr(add_exec_, 1, t_x2, const_cast<void*>(x2.data()));
       aclSetOutputTensorAddr(add_exec_, 0, t_x_out, x_out.data());
     }
-    auto& add_arena = ascend::workspacePool().ensure(stream, add_ws_);
+    auto& add_arena = ascend::GetWorkspacePool().Ensure(stream, add_ws_);
     aclnnAdd(add_arena.buf, add_ws_, add_exec_, stream);
 
     // Obtain shared rstd buffer from pool.
     auto& rstd_arena =
-        ascend::workspacePool().ensure(stream, rstd_size_, "temp");
+        ascend::GetWorkspacePool().Ensure(stream, rstd_size_, "temp");
 
     // Lazily create rstd tensor descriptor on first call.
     if (!rstd_tensor_) {
@@ -102,7 +102,7 @@ class Operator<AddRmsNorm, Device::Type::kAscend, 0> : public AddRmsNorm {
       aclSetOutputTensorAddr(norm_exec_, 0, t_y_out, y_out.data());
       aclSetOutputTensorAddr(norm_exec_, 1, rstd_tensor_, rstd_arena.buf);
     }
-    auto& norm_arena = ascend::workspacePool().ensure(stream, norm_ws_);
+    auto& norm_arena = ascend::GetWorkspacePool().Ensure(stream, norm_ws_);
     aclnnRmsNorm(norm_arena.buf, norm_ws_, norm_exec_, stream);
   }
 

src/ascend/add_rms_norm/kernel_custom.h

@@ -10,7 +10,6 @@
 #include "acl/acl.h"
 #include "aclnn/aclnn_base.h"
 #include "aclnnop/aclnn_cast.h"
-#include "ascend/add_rms_norm/registry.h"
 #include "ascend/common.h"
 #include "ascend/workspace_pool_.h"
 #include "base/add_rms_norm.h"
@@ -121,7 +120,7 @@ class Operator<AddRmsNorm, Device::Type::kAscend, 2> : public AddRmsNorm {
           aclSetOutputTensorAddr(cast_exec_, 0, t_dst, weight_fp32_data_);
         }
 
-        auto& arena = ascend::workspacePool().ensure(stream, cast_ws_);
+        auto& arena = ascend::GetWorkspacePool().Ensure(stream, cast_ws_);
         aclnnCast(arena.buf, cast_ws_, cast_exec_, stream);
         last_weight_ptr_ = cur_weight;
       }

src/ascend/add_rms_norm/kernel_fused.h

@@ -6,9 +6,9 @@
 #include "acl/acl.h"
 #include "aclnn/aclnn_base.h"
 #include "aclnnop/aclnn_add_rms_norm.h"
-#include "ascend/add_rms_norm/registry.h"
 #include "ascend/common.h"
 #include "ascend/workspace_pool_.h"
+#include "base/add_rms_norm.h"
 #include "operator.h"
 
 namespace infini::ops {
@@ -98,7 +98,7 @@ class Operator<AddRmsNorm, Device::Type::kAscend, 1> : public AddRmsNorm {
       aclSetOutputTensorAddr(executor_, 2, t_x_out, x_out.data());
     }
 
-    auto& arena = ascend::workspacePool().ensure(stream, ws_size_);
+    auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_);
     aclnnAddRmsNorm(arena.buf, ws_size_, executor_, stream);
   }
 

src/ascend/add_rms_norm/registry.h

@@ -117,7 +117,7 @@ class Operator<ApplyRotaryPosEmb, Device::Type::kAscend>
       aclSetInputTensorAddr(v2_exec_, 3, t_sin, const_cast<void*>(sin.data()));
     }
 
-    auto& arena = ascend::workspacePool().ensure(stream, v2_ws_);
+    auto& arena = ascend::GetWorkspacePool().Ensure(stream, v2_ws_);
     auto exec_ret =
         aclnnApplyRotaryPosEmbV2(arena.buf, v2_ws_, v2_exec_, stream);
     assert(exec_ret == 0 && "aclnnApplyRotaryPosEmbV2 failed");

src/ascend/apply_rotary_pos_emb/kernel.h

@@ -9,7 +9,6 @@
 #include <vector>
 
 #include "acl/acl.h"
-#include "ascend/apply_rotary_pos_emb/registry.h"
 #include "ascend/atb_common_.h"
 #include "ascend/common.h"
 #include "ascend/workspace_pool_.h"
@@ -141,7 +140,7 @@ class Operator<ApplyRotaryPosEmb, Device::Type::kAscend, 1>
     uint8_t* ws_ptr = nullptr;
 
     if (ws_size > 0) {
-      auto& arena = ascend::workspacePool().ensure(stream, ws_size);
+      auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size);
       ws_ptr = static_cast<uint8_t*>(arena.buf);
     }
 

src/ascend/apply_rotary_pos_emb/kernel_atb.h

@@ -43,7 +43,7 @@ class Operator<Cast, Device::Type::kAscend> : public Cast {
       aclSetOutputTensorAddr(executor_, 0, t_out, out.data());
     }
 
-    auto& arena = ascend::workspacePool().ensure(stream, ws_size_);
+    auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_);
     aclnnCast(arena.buf, ws_size_, executor_, stream);
   }
 

src/ascend/apply_rotary_pos_emb/registry.h

@@ -77,7 +77,7 @@ class Operator<Cat, Device::Type::kAscend> : public Cat {
       aclSetOutputTensorAddr(executor_, 0, t_out, out.data());
     }
 
-    auto& arena = ascend::workspacePool().ensure(stream, ws_size_);
+    auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_);
     aclnnCat(arena.buf, ws_size_, executor_, stream);
   }