From 4d2548aaa86e19958856d7af61f33d3042b3dd2a Mon Sep 17 00:00:00 2001
From: Wenju He <wenju.he@intel.com>
Date: Mon, 11 Apr 2022 12:06:30 +0800
Subject: [PATCH] [SYCL] Add fsycl-id-queries-fit-in-int support for id class

This patch also removes two __SYCL_ASSUME_INT from item class. They
becomes redundant since MIndex is implemented using id.
---
 sycl/include/CL/sycl/id.hpp                   | 21 +++++++++++
 sycl/include/CL/sycl/item.hpp                 |  8 ++---
 .../check_device_code/id_queries_fit_int.cpp  | 36 +++++++++++++++++++
 3 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/sycl/include/CL/sycl/id.hpp b/sycl/include/CL/sycl/id.hpp
index 4fe1934a39147..8726cdb1b6d50 100644
--- a/sycl/include/CL/sycl/id.hpp
+++ b/sycl/include/CL/sycl/id.hpp
@@ -103,6 +103,27 @@ template <int dimensions = 1> class id : public detail::array<dimensions> {
   }
 
 #ifndef __SYCL_DISABLE_ID_TO_INT_CONV__
+  size_t __SYCL_ALWAYS_INLINE get(int dimension) const {
+    base::check_dimension(dimension);
+    size_t Result = this->common_array[dimension];
+    __SYCL_ASSUME_INT(Result);
+    return Result;
+  }
+
+  size_t __SYCL_ALWAYS_INLINE &operator[](int dimension) {
+    base::check_dimension(dimension);
+    size_t &Result = this->common_array[dimension];
+    __SYCL_ASSUME_INT(Result);
+    return Result;
+  }
+
+  size_t __SYCL_ALWAYS_INLINE operator[](int dimension) const {
+    base::check_dimension(dimension);
+    size_t Result = this->common_array[dimension];
+    __SYCL_ASSUME_INT(Result);
+    return Result;
+  }
+
   /* Template operator is not allowed because it disables further type
    * conversion. For example, the next code will not work in case of template
    * conversion:
diff --git a/sycl/include/CL/sycl/item.hpp b/sycl/include/CL/sycl/item.hpp
index 17deeeda98455..3250e74207ae0 100644
--- a/sycl/include/CL/sycl/item.hpp
+++ b/sycl/include/CL/sycl/item.hpp
@@ -50,15 +50,11 @@ template <int dimensions = 1, bool with_offset = true> class item {
   id<dimensions> get_id() const { return MImpl.MIndex; }
 
   size_t __SYCL_ALWAYS_INLINE get_id(int dimension) const {
-    size_t Id = MImpl.MIndex[dimension];
-    __SYCL_ASSUME_INT(Id);
-    return Id;
+    return MImpl.MIndex[dimension];
   }
 
   size_t __SYCL_ALWAYS_INLINE operator[](int dimension) const {
-    size_t Id = MImpl.MIndex[dimension];
-    __SYCL_ASSUME_INT(Id);
-    return Id;
+    return MImpl.MIndex[dimension];
   }
 
   range<dimensions> get_range() const { return MImpl.MExtent; }
diff --git a/sycl/test/check_device_code/id_queries_fit_int.cpp b/sycl/test/check_device_code/id_queries_fit_int.cpp
index 1d8359be58eb6..9ce2132c02cba 100644
--- a/sycl/test/check_device_code/id_queries_fit_int.cpp
+++ b/sycl/test/check_device_code/id_queries_fit_int.cpp
@@ -43,3 +43,39 @@ SYCL_EXTERNAL void testNDItem(nd_item<1> TestNDItem) {
   // CHECK: call void @llvm.assume(i1 {{.*}})
   int OffsetConferted = TestNDItem.get_offset();
 }
+
+// CHECK-LABEL: _Z10TestIdDim1N2cl4sycl2idILi1EEE
+SYCL_EXTERNAL void TestIdDim1(id<1> TestId) {
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id0Get = TestId.get(0);
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id0 = TestId[0];
+}
+
+// CHECK-LABEL: _Z10TestIdDim2N2cl4sycl2idILi2EEE
+SYCL_EXTERNAL void TestIdDim2(id<2> TestId) {
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id0Get = TestId.get(0);
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id1Get = TestId.get(1);
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id0 = TestId[0];
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id1 = TestId[1];
+}
+
+// CHECK-LABEL: _Z10TestIdDim3N2cl4sycl2idILi3EEE
+SYCL_EXTERNAL void TestIdDim3(id<3> TestId) {
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id0Get = TestId.get(0);
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id1Get = TestId.get(1);
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id2Get = TestId.get(2);
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id0 = TestId[0];
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id1 = TestId[1];
+  // CHECK: call void @llvm.assume(i1 {{.*}})
+  int Id2 = TestId[2];
+}