From 40b52fb0cf86fa8805aeec4e1b932cf2d5340c6f Mon Sep 17 00:00:00 2001
From: Taoqiuyu <1042656162@qq.com>
Date: Mon, 11 Aug 2025 10:40:36 +0800
Subject: [PATCH 1/2] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E7=89=88=E6=9D=83?=
 =?UTF-8?q?=E4=BF=A1=E6=81=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 torchrec/hybrid_torchrec/src/ids_process/ids_mapper.cpp | 8 ++++++++
 torchrec/hybrid_torchrec/src/ids_process/ids_mapper.h   | 8 ++++++++
 torchrec/hybrid_torchrec/src/ids_process/unique.cpp     | 9 +++++++++
 torchrec/hybrid_torchrec/src/ids_process/unique.h       | 8 ++++++++
 4 files changed, 33 insertions(+)

diff --git a/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.cpp b/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.cpp
index ceae77d5..9505152c 100644
--- a/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.cpp
+++ b/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.cpp
@@ -1,3 +1,11 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright (c) Huawei Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
 #include "ids_mapper.h"
 
 #include <algorithm>
diff --git a/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.h b/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.h
index 16de28c6..42d3a4e9 100644
--- a/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.h
+++ b/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.h
@@ -1,3 +1,11 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright (c) Huawei Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
 #ifndef IDS_MAPPER_H
 #define IDS_MAPPER_H
 
diff --git a/torchrec/hybrid_torchrec/src/ids_process/unique.cpp b/torchrec/hybrid_torchrec/src/ids_process/unique.cpp
index b5b22194..50790542 100644
--- a/torchrec/hybrid_torchrec/src/ids_process/unique.cpp
+++ b/torchrec/hybrid_torchrec/src/ids_process/unique.cpp
@@ -1,3 +1,12 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright (c) Huawei Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 #include "unique.h"
 
 #include <ATen/TensorIndexing.h>
diff --git a/torchrec/hybrid_torchrec/src/ids_process/unique.h b/torchrec/hybrid_torchrec/src/ids_process/unique.h
index 58cf7c92..74d53cb4 100644
--- a/torchrec/hybrid_torchrec/src/ids_process/unique.h
+++ b/torchrec/hybrid_torchrec/src/ids_process/unique.h
@@ -1,3 +1,11 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright (c) Huawei Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
 #ifndef HYBRID_UNIQUE_H
 #define HYBRID_UNIQUE_H
 #include <omp.h>
-- 
Gitee


From 463c352dd6cfd96484d3f56622e24344107c4716 Mon Sep 17 00:00:00 2001
From: Taoqiuyu <1042656162@qq.com>
Date: Mon, 11 Aug 2025 11:48:23 +0800
Subject: [PATCH 2/2] =?UTF-8?q?=E7=9B=B8=E4=BC=BC=E7=89=87=E6=AE=B5?=
 =?UTF-8?q?=E5=BC=95=E7=94=A8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../hybrid_torchrec/src/ids_process/unique.cpp | 18 +++---------------
 1 file changed, 3 insertions(+), 15 deletions(-)

diff --git a/torchrec/hybrid_torchrec/src/ids_process/unique.cpp b/torchrec/hybrid_torchrec/src/ids_process/unique.cpp
index 50790542..ce786f21 100644
--- a/torchrec/hybrid_torchrec/src/ids_process/unique.cpp
+++ b/torchrec/hybrid_torchrec/src/ids_process/unique.cpp
@@ -30,24 +30,12 @@ struct IsUnique {};
 
 template <typename scalar_t>
 struct IsUnique<scalar_t, false> {
-    inline bool operator()(scalar_t* dataPtr, int64_t i)
+    inline bool operator()(scalar_t* sortedDataPtr, int64_t idx)
     {
-        if (i == 0) {
+        if (idx == 0) {
             return true;
         }
-        return c10::load(&dataPtr[i]) != c10::load(&dataPtr[i - 1]);
-    }
-};
-
-template <typename scalar_t>
-struct IsUnique<scalar_t, true> {
-    inline bool operator()(scalar_t* dataPtr, int64_t i)
-    {
-        if (i == 0) {
-            return true;
-        }
-        return (c10::load(&dataPtr[i]) != c10::load(&dataPtr[i - 1])) &&
-               !(_isnan(dataPtr[i]) && _isnan(dataPtr[i - 1]));
+        return c10::load(&sortedDataPtr[idx]) != c10::load(&sortedDataPtr[idx - 1]);
     }
 };
 
-- 
Gitee