[CIR][Dialect] Add minimal definitions of unified address space offload_* cases (#738)

This PR adds definitions of unified address space cases `offload_*`
discussed in [this RFC
thread](https://discourse.llvm.org/t/rfc-clangir-unified-address-space-design-in-clangir/79728).

It also refactors the `getValueFromLangAS` method in tablegen to a
hand-written method, because it should be a non-trivial map for unified
AS.

---------

Co-authored-by: Julian Oppermann <[email protected]>

Author: seven-mile

clang/include/clang/CIR/Dialect/IR/CIRAttrs.td

@@ -652,19 +652,27 @@ def DynamicCastInfoAttr
 // AddressSpaceAttr
 //===----------------------------------------------------------------------===//
 
-// TODO: other CIR AS cases
-def AS_Target : I32EnumAttrCase<"target", 21>;
+def AS_OffloadPrivate : I32EnumAttrCase<"offload_private", 1>;
+def AS_OffloadLocal : I32EnumAttrCase<"offload_local", 2>;
+def AS_OffloadGlobal : I32EnumAttrCase<"offload_global", 3>;
+def AS_OffloadConstant : I32EnumAttrCase<"offload_constant", 4>;
+def AS_OffloadGeneric : I32EnumAttrCase<"offload_generic", 5>;
+def AS_Target : I32EnumAttrCase<"target", 6>;
 
 def AddressSpaceAttr : CIR_Attr<"AddressSpace", "addrspace"> {
 
   let summary = "Address space attribute for pointer types";
   let description = [{
-    The address space attribute models `clang::LangAS` rather than the LLVM
-    address space, which means it's not yet converted by the address space map
-    to carry target-specific semantics.
+    The address space attribute is used in pointer types. It essentially
+    provides a unified model on top of `clang::LangAS`, rather than LLVM address
+    spaces.
 
-    The representation is one-to-one except for `LangAS::Default`, which
-    corresponds to a null attribute instead.
+    The representation is further simplified: `LangAS::Default` is encoded as
+    a null attribute; many address spaces from different offloading languages
+    are unified as `offload_*`; etc.
+
+    The meaning of `value` parameter is defined as an extensible enum `Kind`,
+    which encodes target AS as offset to the last language AS.
   }];
 
   let parameters = (ins "int32_t":$value);
@@ -690,7 +698,8 @@ def AddressSpaceAttr : CIR_Attr<"AddressSpace", "addrspace"> {
   // simplified assembly format `custom<PointerAddrSpace>`.
 
   list<I32EnumAttrCase> langASCases = [
-    // TODO: includes all non-target CIR AS cases here
+    AS_OffloadPrivate, AS_OffloadLocal, AS_OffloadGlobal, AS_OffloadConstant,
+    AS_OffloadGeneric
   ];
 
   I32EnumAttrCase targetASCase = AS_Target;
@@ -703,9 +712,23 @@ def AddressSpaceAttr : CIR_Attr<"AddressSpace", "addrspace"> {
     bool isTarget() const;
     unsigned getTargetValue() const;
 
-    static std::optional<int32_t> parseValueFromString(llvm::StringRef s);
+    /// Convert a clang LangAS to its corresponding CIR AS storage value. This
+    /// helper does not perform any language-specific mappings (e.g. determining
+    /// the default AS for offloading languages), so these must be handled in
+    /// the caller.
     static std::optional<int32_t> getValueFromLangAS(clang::LangAS v);
+
+    /// Helper methods for the assembly format `custom<PointerAddrSpace>`.
+    static std::optional<int32_t> parseValueFromString(llvm::StringRef s);
     static std::optional<llvm::StringRef> stringifyValue(int32_t v);
+
+    struct Kind {
+    }]#!interleave(
+      !foreach(case, langASCases,
+        "static constexpr int32_t "#case.symbol#" = "#case.value#";"
+      ), "\n"
+    )#[{
+    };
   }];
 
   let extraClassDefinition = [{
@@ -757,31 +780,6 @@ def AddressSpaceAttr : CIR_Attr<"AddressSpace", "addrspace"> {
         return std::nullopt;
       }
     }
-
-    std::optional<int32_t>
-    $cppClass::getValueFromLangAS(clang::LangAS langAS) {
-      assert((langAS == clang::LangAS::Default ||
-             clang::isTargetAddressSpace(langAS)) &&
-             "Language-specific address spaces are not supported");
-      switch (langAS) {
-      }]
-      #
-      !interleave(
-        !foreach(case, langASCases,
-          "case clang::LangAS::"#case.symbol
-          # [{: llvm_unreachable("Not Yet Supported");}] ),
-        "\n"
-      )
-      #
-      [{
-      case clang::LangAS::Default:
-        // Default address space should be encoded as a null attribute.
-        return std::nullopt;
-      default:
-        // Target address space offset arithmetics
-        return clang::toTargetAddressSpace(langAS) + kFirstTargetASValue;
-      }
-    }
   }];
 }
 

clang/lib/CIR/Dialect/IR/CIRAttrs.cpp

@@ -553,6 +553,44 @@ LogicalResult OpenCLKernelMetadataAttr::verify(
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// AddressSpaceAttr definitions
+//===----------------------------------------------------------------------===//
+
+std::optional<int32_t>
+AddressSpaceAttr::getValueFromLangAS(clang::LangAS langAS) {
+  using clang::LangAS;
+  switch (langAS) {
+  case LangAS::Default:
+    // Default address space should be encoded as a null attribute.
+    return std::nullopt;
+  case LangAS::opencl_global:
+  case LangAS::opencl_local:
+  case LangAS::opencl_constant:
+  case LangAS::opencl_private:
+  case LangAS::opencl_generic:
+  case LangAS::opencl_global_device:
+  case LangAS::opencl_global_host:
+  case LangAS::cuda_device:
+  case LangAS::cuda_constant:
+  case LangAS::cuda_shared:
+  case LangAS::sycl_global:
+  case LangAS::sycl_global_device:
+  case LangAS::sycl_global_host:
+  case LangAS::sycl_local:
+  case LangAS::sycl_private:
+  case LangAS::ptr32_sptr:
+  case LangAS::ptr32_uptr:
+  case LangAS::ptr64:
+  case LangAS::hlsl_groupshared:
+  case LangAS::wasm_funcref:
+    llvm_unreachable("NYI");
+  default:
+    // Target address space offset arithmetics
+    return clang::toTargetAddressSpace(langAS) + kFirstTargetASValue;
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // CIR Dialect
 //===----------------------------------------------------------------------===//

clang/test/CIR/IR/address-space.cir

@@ -13,4 +13,29 @@ module {
   cir.func @test_format2(%arg0: !cir.ptr<!s32i>) {
     cir.return
   }
+
+  // CHECK: @test_format3(%arg0: !cir.ptr<!s32i, addrspace(offload_private)>)
+  cir.func @test_format3(%arg0: !cir.ptr<!s32i, addrspace(offload_private)>) {
+    cir.return
+  }
+
+  // CHECK: @test_format4(%arg0: !cir.ptr<!s32i, addrspace(offload_local)>)
+  cir.func @test_format4(%arg0: !cir.ptr<!s32i, addrspace(offload_local)>) {
+    cir.return
+  }
+
+  // CHECK: @test_format5(%arg0: !cir.ptr<!s32i, addrspace(offload_global)>)
+  cir.func @test_format5(%arg0: !cir.ptr<!s32i, addrspace(offload_global)>) {
+    cir.return
+  }
+
+  // CHECK: @test_format6(%arg0: !cir.ptr<!s32i, addrspace(offload_constant)>)
+  cir.func @test_format6(%arg0: !cir.ptr<!s32i, addrspace(offload_constant)>) {
+    cir.return
+  }
+
+  // CHECK: @test_format7(%arg0: !cir.ptr<!s32i, addrspace(offload_generic)>)
+  cir.func @test_format7(%arg0: !cir.ptr<!s32i, addrspace(offload_generic)>) {
+    cir.return
+  }
 }