[SYCL] Refine annotated_arg addrspace with USM kind property

sycl/include/sycl/ext/oneapi/experimental/annotated_arg/annotated_arg.hpp

@@ -70,11 +70,22 @@ __SYCL_TYPE(annotated_arg) annotated_arg<T *, detail::properties_t<Props...>> {
 
   template <typename T2, typename PropertyListT> friend class annotated_arg;
 
+#ifdef __ENABLE_USM_ADDR_SPACE__
+  using global_pointer_t = std::conditional_t<
+      detail::IsUsmKindDevice<property_list_t>::value,
+      typename sycl::ext::intel::decorated_device_ptr<T>::pointer,
+      std::conditional_t<
+          detail::IsUsmKindHost<property_list_t>::value,
+          typename sycl::ext::intel::decorated_host_ptr<T>::pointer,
+          typename decorated_global_ptr<T>::pointer>>;
+#else  // __ENABLE_USM_ADDR_SPACE__
+  using global_pointer_t = typename decorated_global_ptr<T>::pointer;
+#endif // __ENABLE_USM_ADDR_SPACE__
+
 #ifdef __SYCL_DEVICE_ONLY__
   void __init([[__sycl_detail__::add_ir_attributes_kernel_parameter(
       detail::PropertyMetaInfo<Props>::name...,
-      detail::PropertyMetaInfo<Props>::value...)]]
-              typename decorated_global_ptr<T>::pointer _obj) {
+      detail::PropertyMetaInfo<Props>::value...)]] global_pointer_t _obj) {
     obj = _obj;
   }
 #endif

sycl/test/extensions/annotated_arg/usm_kind_intel_fpga.cpp

@@ -0,0 +1,55 @@
+// RUN: %clangxx -fsycl-device-only -fsycl-targets=spir64_fpga -S -emit-llvm %s -o - | FileCheck %s  --check-prefix CHECK-IR
+
+// Tests that the address space of `annotated_arg` kernel argument is refined
+// when:
+// 1. `usm_kind` property is specified in annotated_arg type, and
+// 2. both `__SYCL_DEVICE_ONLY__` and `__ENABLE_USM_ADDR_SPACE__` are turned on
+// (i.e. equiv to flags "-fsycl-device-only -fsycl-targets=spir64_fpga")
+
+#include "sycl/sycl.hpp"
+
+#include <sycl/ext/intel/fpga_extensions.hpp>
+
+#include <iostream>
+
+using namespace sycl;
+using namespace ext::oneapi::experimental;
+using namespace ext::intel::experimental;
+
+using annotated_arg_t1 =
+    annotated_arg<int *,
+                  decltype(properties(buffer_location<0>, usm_kind_device))>;
+using annotated_arg_t2 =
+    annotated_arg<int *,
+                  decltype(properties(buffer_location<1>, usm_kind_host))>;
+using annotated_arg_t3 =
+    annotated_arg<int *,
+                  decltype(properties(buffer_location<2>, usm_kind_shared))>;
+
+struct MyIP {
+
+  // CHECK-IR: spir_kernel void @_ZTS4MyIP(ptr addrspace(5) {{.*}} %_arg_a, ptr addrspace(6) {{.*}} %_arg_b, ptr addrspace(1) {{.*}} %_arg_c)
+  annotated_arg_t1 a;
+  annotated_arg_t2 b;
+  annotated_arg_t3 c;
+
+  void operator()() const { *a = *b + *c; }
+};
+
+void TestVectorAddWithAnnotatedMMHosts() {
+  queue q;
+  auto p1 = malloc_device<int>(5, q);
+  auto p2 = malloc_host<int>(5, q);
+  auto p3 = malloc_shared<int>(5, q);
+
+  q.submit([&](handler &h) { h.single_task(MyIP{p1, p2, p3}); }).wait();
+
+  free(p1, q);
+  free(p2, q);
+  free(p3, q);
+}
+
+int main() {
+  TestVectorAddWithAnnotatedMMHosts();
+  return 0;
+}