[MLIR][LLVM] add dwarfAddressSpace to DIDerivedType (#92043)

This field is present in LLVM, but was missing from the MLIR wrapper
type. This addition allows MLIR languages to add proper DWARF info for
GPU programs.

Author: willghatch

mlir/include/mlir-c/Dialect/LLVM.h

@@ -241,11 +241,16 @@ MLIR_CAPI_EXPORTED MlirAttribute mlirLLVMDICompositeTypeAttrGet(
     MlirAttribute baseType, int64_t flags, uint64_t sizeInBits,
     uint64_t alignInBits, intptr_t nElements, MlirAttribute const *elements);
 
-/// Creates a LLVM DIDerivedType attribute.
+/// Creates a LLVM DIDerivedType attribute.  Note that `dwarfAddressSpace` is an
+/// optional field, where `MLIR_CAPI_DWARF_ADDRESS_SPACE_NULL` indicates null
+/// and non-negative values indicate a value present.
 MLIR_CAPI_EXPORTED MlirAttribute mlirLLVMDIDerivedTypeAttrGet(
     MlirContext ctx, unsigned int tag, MlirAttribute name,
     MlirAttribute baseType, uint64_t sizeInBits, uint32_t alignInBits,
-    uint64_t offsetInBits, MlirAttribute extraData);
+    uint64_t offsetInBits, int64_t dwarfAddressSpace, MlirAttribute extraData);
+
+/// Constant to represent std::nullopt for dwarfAddressSpace to omit the field.
+#define MLIR_CAPI_DWARF_ADDRESS_SPACE_NULL -1
 
 /// Gets the base type from a LLVM DIDerivedType attribute.
 MLIR_CAPI_EXPORTED MlirAttribute

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

@@ -418,6 +418,7 @@ def LLVM_DIDerivedTypeAttr : LLVM_Attr<"DIDerivedType", "di_derived_type",
     OptionalParameter<"uint64_t">:$sizeInBits,
     OptionalParameter<"uint32_t">:$alignInBits,
     OptionalParameter<"uint64_t">:$offsetInBits,
+    OptionalParameter<"std::optional<unsigned>">:$dwarfAddressSpace,
     OptionalParameter<"DINodeAttr">:$extraData
   );
   let assemblyFormat = "`<` struct(params) `>`";

mlir/lib/CAPI/Dialect/LLVM.cpp

@@ -176,15 +176,17 @@ MlirAttribute mlirLLVMDICompositeTypeAttrGet(
                           [](Attribute a) { return cast<DINodeAttr>(a); })));
 }
 
-MlirAttribute
-mlirLLVMDIDerivedTypeAttrGet(MlirContext ctx, unsigned int tag,
-                             MlirAttribute name, MlirAttribute baseType,
-                             uint64_t sizeInBits, uint32_t alignInBits,
-                             uint64_t offsetInBits, MlirAttribute extraData) {
+MlirAttribute mlirLLVMDIDerivedTypeAttrGet(
+    MlirContext ctx, unsigned int tag, MlirAttribute name,
+    MlirAttribute baseType, uint64_t sizeInBits, uint32_t alignInBits,
+    uint64_t offsetInBits, int64_t dwarfAddressSpace, MlirAttribute extraData) {
+  std::optional<unsigned> addressSpace = std::nullopt;
+  if (dwarfAddressSpace >= 0)
+    addressSpace = (unsigned)dwarfAddressSpace;
   return wrap(DIDerivedTypeAttr::get(
       unwrap(ctx), tag, cast<StringAttr>(unwrap(name)),
       cast<DITypeAttr>(unwrap(baseType)), sizeInBits, alignInBits, offsetInBits,
-      cast<DINodeAttr>(unwrap(extraData))));
+      addressSpace, cast<DINodeAttr>(unwrap(extraData))));
 }
 
 MlirAttribute

mlir/lib/Target/LLVMIR/DebugImporter.cpp

@@ -106,7 +106,7 @@ DIDerivedTypeAttr DebugImporter::translateImpl(llvm::DIDerivedType *node) {
   return DIDerivedTypeAttr::get(
       context, node->getTag(), getStringAttrOrNull(node->getRawName()),
       baseType, node->getSizeInBits(), node->getAlignInBits(),
-      node->getOffsetInBits(), extraData);
+      node->getOffsetInBits(), node->getDWARFAddressSpace(), extraData);
 }
 
 DIFileAttr DebugImporter::translateImpl(llvm::DIFile *node) {

mlir/lib/Target/LLVMIR/DebugTranslation.cpp

@@ -165,7 +165,7 @@ llvm::DIDerivedType *DebugTranslation::translateImpl(DIDerivedTypeAttr attr) {
       /*File=*/nullptr, /*Line=*/0,
       /*Scope=*/nullptr, translate(attr.getBaseType()), attr.getSizeInBits(),
       attr.getAlignInBits(), attr.getOffsetInBits(),
-      /*DWARFAddressSpace=*/std::nullopt, /*PtrAuthData=*/std::nullopt,
+      attr.getDwarfAddressSpace(), /*PtrAuthData=*/std::nullopt,
       /*Flags=*/llvm::DINode::FlagZero, translate(attr.getExtraData()));
 }
 

mlir/test/CAPI/llvm.c

@@ -296,8 +296,14 @@ static void testDebugInfoAttributes(MlirContext ctx) {
   mlirAttributeDump(mlirLLVMDILocalVariableAttrGet(ctx, compile_unit, foo, file,
                                                    1, 0, 8, di_type));
   // CHECK: #llvm.di_derived_type<{{.*}}>
+  // CHECK-NOT: dwarfAddressSpace
+  mlirAttributeDump(mlirLLVMDIDerivedTypeAttrGet(
+      ctx, 0, bar, di_type, 64, 8, 0, MLIR_CAPI_DWARF_ADDRESS_SPACE_NULL,
+      di_type));
+
+  // CHECK: #llvm.di_derived_type<{{.*}} dwarfAddressSpace = 3{{.*}}>
   mlirAttributeDump(
-      mlirLLVMDIDerivedTypeAttrGet(ctx, 0, bar, di_type, 64, 8, 0, di_type));
+      mlirLLVMDIDerivedTypeAttrGet(ctx, 0, bar, di_type, 64, 8, 0, 3, di_type));
 
   // CHECK: #llvm.di_composite_type<{{.*}}>
   mlirAttributeDump(mlirLLVMDICompositeTypeAttrGet(

mlir/test/Dialect/LLVMIR/debuginfo.mlir

@@ -37,6 +37,13 @@
   tag = DW_TAG_pointer_type, name = "ptr1"
 >
 
+// CHECK-DAG: #[[PTR2:.*]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[INT0]], sizeInBits = 64, alignInBits = 32, offsetInBits = 4, dwarfAddressSpace = 3, extraData = #[[INT1]]>
+#ptr2 = #llvm.di_derived_type<
+  tag = DW_TAG_pointer_type, baseType = #int0,
+  sizeInBits = 64, alignInBits = 32, offsetInBits = 4,
+  dwarfAddressSpace = 3, extraData = #int1
+>
+
 // CHECK-DAG: #[[COMP0:.*]] = #llvm.di_composite_type<tag = DW_TAG_array_type, name = "array0", line = 10, sizeInBits = 128, alignInBits = 32>
 #comp0 = #llvm.di_composite_type<
   tag = DW_TAG_array_type, name = "array0",
@@ -73,9 +80,9 @@
   flags = "TypePassByReference|NonTrivial"
 >
 
-// CHECK-DAG: #[[SPTYPE0:.*]] = #llvm.di_subroutine_type<callingConvention = DW_CC_normal, types = #[[NULL]], #[[INT0]], #[[PTR0]], #[[PTR1]], #[[COMP0:.*]], #[[COMP1:.*]], #[[COMP2:.*]]>
+// CHECK-DAG: #[[SPTYPE0:.*]] = #llvm.di_subroutine_type<callingConvention = DW_CC_normal, types = #[[NULL]], #[[INT0]], #[[PTR0]], #[[PTR1]], #[[PTR2]], #[[COMP0:.*]], #[[COMP1:.*]], #[[COMP2:.*]]>
 #spType0 = #llvm.di_subroutine_type<
-  callingConvention = DW_CC_normal, types = #null, #int0, #ptr0, #ptr1, #comp0, #comp1, #comp2
+  callingConvention = DW_CC_normal, types = #null, #int0, #ptr0, #ptr1, #ptr2, #comp0, #comp1, #comp2
 >
 
 // CHECK-DAG: #[[SPTYPE1:.*]] = #llvm.di_subroutine_type<types = #[[INT1]], #[[INT1]]>

mlir/test/Target/LLVMIR/Import/debug-info.ll

@@ -137,7 +137,8 @@ define void @basic_type() !dbg !3 {
 ; CHECK: #[[INT:.+]] = #llvm.di_basic_type<tag = DW_TAG_base_type, name = "int">
 ; CHECK: #[[PTR1:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[INT]]>
 ; CHECK: #[[PTR2:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, name = "mypointer", baseType = #[[INT]], sizeInBits = 64, alignInBits = 32, offsetInBits = 4, extraData = #[[INT]]>
-; CHECK: #llvm.di_subroutine_type<types = #[[PTR1]], #[[PTR2]]>
+; CHECK: #[[PTR3:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[INT]], dwarfAddressSpace = 3>
+; CHECK: #llvm.di_subroutine_type<types = #[[PTR1]], #[[PTR2]], #[[PTR3]]>
 
 define void @derived_type() !dbg !3 {
   ret void
@@ -150,10 +151,11 @@ define void @derived_type() !dbg !3 {
 !2 = !DIFile(filename: "debug-info.ll", directory: "/")
 !3 = distinct !DISubprogram(name: "derived_type", scope: !2, file: !2, spFlags: DISPFlagDefinition, unit: !1, type: !4)
 !4 = !DISubroutineType(types: !5)
-!5 = !{!7, !8}
+!5 = !{!7, !8, !9}
 !6 = !DIBasicType(name: "int")
 !7 = !DIDerivedType(tag: DW_TAG_pointer_type, baseType: !6)
 !8 = !DIDerivedType(name: "mypointer", tag: DW_TAG_pointer_type, baseType: !6, size: 64, align: 32, offset: 4, extraData: !6)
+!9 = !DIDerivedType(tag: DW_TAG_pointer_type, baseType: !6, dwarfAddressSpace: 3)
 
 ; // -----
 

mlir/test/Target/LLVMIR/llvmir-debug.mlir

@@ -35,6 +35,11 @@ llvm.func @func_no_debug() {
   // Specify the name parameter.
   tag = DW_TAG_pointer_type, name = "named", baseType = #si32
 >
+#ptrWithAddressSpace = #llvm.di_derived_type<
+  tag = DW_TAG_pointer_type, baseType = #si32,
+  sizeInBits = 64, alignInBits = 32, offsetInBits = 8,
+  dwarfAddressSpace = 3
+>
 #cu = #llvm.di_compile_unit<
   id = distinct[0]<>, sourceLanguage = DW_LANG_C, file = #file,
   producer = "MLIR", isOptimized = true, emissionKind = Full,
@@ -51,7 +56,7 @@ llvm.func @func_no_debug() {
   elements = #llvm.di_subrange<lowerBound = 0, upperBound = 4, stride = 1>
 >
 #null = #llvm.di_null_type
-#spType0 = #llvm.di_subroutine_type<callingConvention = DW_CC_normal, types = #null, #si64, #ptr, #named, #composite, #vector>
+#spType0 = #llvm.di_subroutine_type<callingConvention = DW_CC_normal, types = #null, #si64, #ptr, #named, #ptrWithAddressSpace, #composite, #vector>
 #toplevel_namespace = #llvm.di_namespace<
   name = "toplevel", exportSymbols = true
 >
@@ -138,11 +143,12 @@ llvm.func @empty_types() {
 // CHECK: ![[NESTED_NAMESPACE]] = !DINamespace(name: "nested", scope: ![[TOPLEVEL_NAMESPACE:.*]])
 // CHECK: ![[TOPLEVEL_NAMESPACE]] = !DINamespace(name: "toplevel", scope: null, exportSymbols: true)
 // CHECK: ![[FUNC_TYPE]] = !DISubroutineType(cc: DW_CC_normal, types: ![[FUNC_ARGS:.*]])
-// CHECK: ![[FUNC_ARGS]] = !{null, ![[ARG_TYPE:.*]], ![[PTR_TYPE:.*]], ![[NAMED_TYPE:.*]], ![[COMPOSITE_TYPE:.*]], ![[VECTOR_TYPE:.*]]}
+// CHECK: ![[FUNC_ARGS]] = !{null, ![[ARG_TYPE:.*]], ![[PTR_TYPE:.*]], ![[NAMED_TYPE:.*]], ![[PTR_WITH_ADDR_SPACE:.*]], ![[COMPOSITE_TYPE:.*]], ![[VECTOR_TYPE:.*]]}
 // CHECK: ![[ARG_TYPE]] = !DIBasicType(name: "si64")
 // CHECK: ![[PTR_TYPE]] = !DIDerivedType(tag: DW_TAG_pointer_type, baseType: ![[BASE_TYPE:.*]], size: 64, align: 32, offset: 8, extraData: ![[BASE_TYPE]])
 // CHECK: ![[BASE_TYPE]] = !DIBasicType(name: "si32", size: 32, encoding: DW_ATE_signed)
 // CHECK: ![[NAMED_TYPE]] = !DIDerivedType(tag: DW_TAG_pointer_type, name: "named", baseType: ![[BASE_TYPE:.*]])
+// CHECK: ![[PTR_WITH_ADDR_SPACE]] = !DIDerivedType(tag: DW_TAG_pointer_type, baseType: ![[BASE_TYPE:.*]], size: 64, align: 32, offset: 8, dwarfAddressSpace: 3)
 // CHECK: ![[COMPOSITE_TYPE]] = distinct !DICompositeType(tag: DW_TAG_structure_type, name: "composite", file: ![[CU_FILE_LOC]], line: 42, size: 64, align: 32, elements: ![[COMPOSITE_ELEMENTS:.*]])
 // CHECK: ![[COMPOSITE_ELEMENTS]] = !{![[COMPOSITE_ELEMENT:.*]]}
 // CHECK: ![[COMPOSITE_ELEMENT]] = !DISubrange(count: 4)