fix: fix rmsnorm dispatch to use one dispatch

Author: bitzyz

src/cambricon/rmsnorm/rms_norm.h

@@ -33,48 +33,18 @@ class Operator<RmsNorm, Device::Type::kCambricon> : public RmsNorm {
 
     DispatchFunc<
         List<DataType::kFloat16, DataType::kBFloat16, DataType::kFloat32>,
-        List<Device::Type::kCambricon>>(
-        {static_cast<int64_t>(input.dtype()),
-         static_cast<int64_t>(Device::Type::kCambricon)},
-        0,
-        [&](auto input_tag) {
-          constexpr DataType IDT = static_cast<DataType>(ListGet<0>(input_tag));
-          using InputT = TypeMapType<IDT>;
-          DispatchFunc<
-              List<DataType::kFloat16, DataType::kBFloat16, DataType::kFloat32>,
-              List<Device::Type::kCambricon>>(
-              {static_cast<int64_t>(weight.dtype()),
-               static_cast<int64_t>(Device::Type::kCambricon)},
-              0,
-              [&](auto weight_tag) {
-                constexpr DataType WDT =
-                    static_cast<DataType>(ListGet<0>(weight_tag));
-                using WeightT = TypeMapType<WDT>;
-
-                RmsnormUnion<InputT, WeightT>(
-                    workspace, core_per_cluster, cluster_count, queue,
-                    out.data(), input.data(), weight.data(), out_shape_.data(),
-                    out_strides_.data(), input_strides_.data(), eps, ndim_);
-              },
-              "CambriconRmsNorm::operator() - weight dispatch", List<>{});
+        List<DataType::kFloat16, DataType::kBFloat16, DataType::kFloat32>>(
+        {input.dtype(), weight.dtype()},
+        [&](auto input_tag, auto weight_tag) {
+          using InputT = typename decltype(input_tag)::type;
+          using WeightT = typename decltype(weight_tag)::type;
+
+          RmsnormUnion<InputT, WeightT>(
+              workspace, core_per_cluster, cluster_count, queue,
+              out.data(), input.data(), weight.data(), out_shape_.data(),
+              out_strides_.data(), input_strides_.data(), eps, ndim_);
         },
-        "CambriconRmsNorm::operator() - output dispatch", List<>{});
-    // DispatchFunc<List<DataType::kFloat16, DataType::kBFloat16,
-    // DataType::kFloat32>,
-    //              List<DataType::kFloat16, DataType::kBFloat16,
-    //              DataType::kFloat32>>(
-    //     {input.dtype(), weight.dtype()},
-    //     [&](auto input_tag, auto weight_tag) {
-    //       using InputT = typename decltype(input_tag)::type;
-    //       using WeightT = typename decltype(weight_tag)::type;
-
-    //       RmsnormUnion<InputT, WeightT>(
-    //           workspace, core_per_cluster, cluster_count, queue,
-    //           out.data(), input.data(), weight.data(),
-    //           out_shape_.data(), out_strides_.data(),
-    //           input_strides_.data(), eps, ndim_);
-    //     },
-    //     "CambriconRmsNorm::operator()");
+        "CambriconRmsNorm::operator() - output dispatch");
   }
 
   ~Operator() { cnrtFree(default_workspace_); }