[CIR][Dialect] Emit OpenCL kernel argument metadata (#767)

Similar to #705, this PR implements the remaining
`genKernelArgMetadata()` logic.

The attribute `cir.cl.kernel_arg_metadata` is also intentionally placed
in the `cir.func`'s `extra_attrs` rather than `cir.func`'s standard
`arg_attrs` list. Also, the metadata is stored by `Array` with proper
verification on it. See the tablegen doc string for details.

This is in order to
* keep it side-by-side with `cl.kernel_metadata`.
* still emit metadata when kernel has an *empty* arg list (see the test
`kernel-arg-meatadata.cl`).
* avoid horrors of repeating the long name `cir.cl.kernel_arg_metadata`
for `numArgs` times.

Because clangir doesn't support OpenCL built-in types and the `half`
floating point type yet, their changes and test cases are not included.
Corresponding missing feature flag is added.

Author: seven-mile

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

@@ -92,4 +92,59 @@ def OpenCLKernelMetadataAttr
 
 }
 
+//===----------------------------------------------------------------------===//
+// OpenCLKernelArgMetadataAttr
+//===----------------------------------------------------------------------===//
+
+def OpenCLKernelArgMetadataAttr
+    : CIR_Attr<"OpenCLKernelArgMetadata", "cl.kernel_arg_metadata"> {
+  
+  let summary = "OpenCL kernel argument metadata";
+  let description = [{
+    Provide the required information of an OpenCL kernel argument for the SPIR-V
+    backend.
+
+    All parameters are arrays, containing the information of the argument in
+    the same order as they appear in the source code.
+
+    The `addr_space` parameter is an array of I32 that provides the address
+    space of the argument. It's useful for special types like `image`, which
+    have implicit global address space.
+
+    Other parameters are arrays of strings that pass through the information
+    from the source code correspondingly.
+
+    All the fields are mandatory except for `name`, which is optional.
+
+    Example:
+    ```
+    #fn_attr = #cir<extra({cl.kernel_arg_metadata = #cir.cl.kernel_arg_metadata<
+      addr_space = [1 : i32],
+      access_qual = ["none"],
+      type = ["char*"],
+      base_type = ["char*"],
+      type_qual = [""],
+      name = ["in"]
+    >})>
+
+    cir.func @kernel(%arg0: !s32i) extra(#fn_attr) {
+      cir.return
+    }
+    ```
+  }];
+
+  let parameters = (ins
+    "ArrayAttr":$addr_space,
+    "ArrayAttr":$access_qual,
+    "ArrayAttr":$type,
+    "ArrayAttr":$base_type,
+    "ArrayAttr":$type_qual,
+    OptionalParameter<"ArrayAttr">:$name
+  );
+
+  let assemblyFormat = "`<` struct(params) `>`";
+
+  let genVerifyDecl = 1;
+}
+
 #endif // MLIR_CIR_DIALECT_CIR_OPENCL_ATTRS

clang/include/clang/CIR/MissingFeatures.h

@@ -142,7 +142,7 @@ struct MissingFeatures {
   static bool getFPFeaturesInEffect() { return false; }
   static bool cxxABI() { return false; }
   static bool openCL() { return false; }
-  static bool openCLGenKernelMetadata() { return false; }
+  static bool openCLBuiltinTypes() { return false; }
   static bool CUDA() { return false; }
   static bool openMP() { return false; }
   static bool openMPRuntime() { return false; }

clang/lib/CIR/CodeGen/CIRGenFunction.cpp

@@ -1725,8 +1725,7 @@ void CIRGenFunction::buildKernelMetadata(const FunctionDecl *FD,
   if (!FD->hasAttr<OpenCLKernelAttr>() && !FD->hasAttr<CUDAGlobalAttr>())
     return;
 
-  // TODO(cir): CGM.genKernelArgMetadata(Fn, FD, this);
-  assert(!MissingFeatures::openCLGenKernelMetadata());
+  CGM.genKernelArgMetadata(Fn, FD, this);
 
   if (!getLangOpts().OpenCL)
     return;

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -3061,3 +3061,174 @@ mlir::cir::SourceLanguage CIRGenModule::getCIRSourceLanguage() {
   // TODO(cir): support remaining source languages.
   llvm_unreachable("CIR does not yet support the given source language");
 }
+
+// Returns the address space id that should be produced to the
+// kernel_arg_addr_space metadata. This is always fixed to the ids
+// as specified in the SPIR 2.0 specification in order to differentiate
+// for example in clGetKernelArgInfo() implementation between the address
+// spaces with targets without unique mapping to the OpenCL address spaces
+// (basically all single AS CPUs).
+static unsigned ArgInfoAddressSpace(LangAS AS) {
+  switch (AS) {
+  case LangAS::opencl_global:
+    return 1;
+  case LangAS::opencl_constant:
+    return 2;
+  case LangAS::opencl_local:
+    return 3;
+  case LangAS::opencl_generic:
+    return 4; // Not in SPIR 2.0 specs.
+  case LangAS::opencl_global_device:
+    return 5;
+  case LangAS::opencl_global_host:
+    return 6;
+  default:
+    return 0; // Assume private.
+  }
+}
+
+void CIRGenModule::genKernelArgMetadata(mlir::cir::FuncOp Fn,
+                                        const FunctionDecl *FD,
+                                        CIRGenFunction *CGF) {
+  assert(((FD && CGF) || (!FD && !CGF)) &&
+         "Incorrect use - FD and CGF should either be both null or not!");
+  // Create MDNodes that represent the kernel arg metadata.
+  // Each MDNode is a list in the form of "key", N number of values which is
+  // the same number of values as their are kernel arguments.
+
+  const PrintingPolicy &Policy = getASTContext().getPrintingPolicy();
+
+  // Integer values for the kernel argument address space qualifiers.
+  SmallVector<int32_t, 8> addressQuals;
+
+  // Attrs for the kernel argument access qualifiers (images only).
+  SmallVector<mlir::Attribute, 8> accessQuals;
+
+  // Attrs for the kernel argument type names.
+  SmallVector<mlir::Attribute, 8> argTypeNames;
+
+  // Attrs for the kernel argument base type names.
+  SmallVector<mlir::Attribute, 8> argBaseTypeNames;
+
+  // Attrs for the kernel argument type qualifiers.
+  SmallVector<mlir::Attribute, 8> argTypeQuals;
+
+  // Attrs for the kernel argument names.
+  SmallVector<mlir::Attribute, 8> argNames;
+
+  // OpenCL image and pipe types require special treatments for some metadata
+  assert(!MissingFeatures::openCLBuiltinTypes());
+
+  if (FD && CGF)
+    for (unsigned i = 0, e = FD->getNumParams(); i != e; ++i) {
+      const ParmVarDecl *parm = FD->getParamDecl(i);
+      // Get argument name.
+      argNames.push_back(builder.getStringAttr(parm->getName()));
+
+      if (!getLangOpts().OpenCL)
+        continue;
+      QualType ty = parm->getType();
+      std::string typeQuals;
+
+      // Get image and pipe access qualifier:
+      if (ty->isImageType() || ty->isPipeType()) {
+        llvm_unreachable("NYI");
+      } else
+        accessQuals.push_back(builder.getStringAttr("none"));
+
+      auto getTypeSpelling = [&](QualType Ty) {
+        auto typeName = Ty.getUnqualifiedType().getAsString(Policy);
+
+        if (Ty.isCanonical()) {
+          StringRef typeNameRef = typeName;
+          // Turn "unsigned type" to "utype"
+          if (typeNameRef.consume_front("unsigned "))
+            return std::string("u") + typeNameRef.str();
+          if (typeNameRef.consume_front("signed "))
+            return typeNameRef.str();
+        }
+
+        return typeName;
+      };
+
+      if (ty->isPointerType()) {
+        QualType pointeeTy = ty->getPointeeType();
+
+        // Get address qualifier.
+        addressQuals.push_back(
+            ArgInfoAddressSpace(pointeeTy.getAddressSpace()));
+
+        // Get argument type name.
+        std::string typeName = getTypeSpelling(pointeeTy) + "*";
+        std::string baseTypeName =
+            getTypeSpelling(pointeeTy.getCanonicalType()) + "*";
+        argTypeNames.push_back(builder.getStringAttr(typeName));
+        argBaseTypeNames.push_back(builder.getStringAttr(baseTypeName));
+
+        // Get argument type qualifiers:
+        if (ty.isRestrictQualified())
+          typeQuals = "restrict";
+        if (pointeeTy.isConstQualified() ||
+            (pointeeTy.getAddressSpace() == LangAS::opencl_constant))
+          typeQuals += typeQuals.empty() ? "const" : " const";
+        if (pointeeTy.isVolatileQualified())
+          typeQuals += typeQuals.empty() ? "volatile" : " volatile";
+      } else {
+        uint32_t AddrSpc = 0;
+        bool isPipe = ty->isPipeType();
+        if (ty->isImageType() || isPipe)
+          llvm_unreachable("NYI");
+
+        addressQuals.push_back(AddrSpc);
+
+        // Get argument type name.
+        ty = isPipe ? ty->castAs<PipeType>()->getElementType() : ty;
+        std::string typeName = getTypeSpelling(ty);
+        std::string baseTypeName = getTypeSpelling(ty.getCanonicalType());
+
+        // Remove access qualifiers on images
+        // (as they are inseparable from type in clang implementation,
+        // but OpenCL spec provides a special query to get access qualifier
+        // via clGetKernelArgInfo with CL_KERNEL_ARG_ACCESS_QUALIFIER):
+        if (ty->isImageType()) {
+          llvm_unreachable("NYI");
+        }
+
+        argTypeNames.push_back(builder.getStringAttr(typeName));
+        argBaseTypeNames.push_back(builder.getStringAttr(baseTypeName));
+
+        if (isPipe)
+          llvm_unreachable("NYI");
+      }
+      argTypeQuals.push_back(builder.getStringAttr(typeQuals));
+    }
+
+  bool shouldEmitArgName = getCodeGenOpts().EmitOpenCLArgMetadata ||
+                           getCodeGenOpts().HIPSaveKernelArgName;
+
+  if (getLangOpts().OpenCL) {
+    // The kernel arg name is emitted only when `-cl-kernel-arg-info` is on,
+    // since it is only used to support `clGetKernelArgInfo` which requires
+    // `-cl-kernel-arg-info` to work. The other metadata are mandatory because
+    // they are necessary for OpenCL runtime to set kernel argument.
+    mlir::ArrayAttr resArgNames = {};
+    if (shouldEmitArgName)
+      resArgNames = builder.getArrayAttr(argNames);
+
+    // Update the function's extra attributes with the kernel argument metadata.
+    auto value = mlir::cir::OpenCLKernelArgMetadataAttr::get(
+        Fn.getContext(), builder.getI32ArrayAttr(addressQuals),
+        builder.getArrayAttr(accessQuals), builder.getArrayAttr(argTypeNames),
+        builder.getArrayAttr(argBaseTypeNames),
+        builder.getArrayAttr(argTypeQuals), resArgNames);
+    mlir::NamedAttrList items{Fn.getExtraAttrs().getElements().getValue()};
+    auto oldValue = items.set(value.getMnemonic(), value);
+    if (oldValue != value) {
+      Fn.setExtraAttrsAttr(mlir::cir::ExtraFuncAttributesAttr::get(
+          builder.getContext(), builder.getDictionaryAttr(items)));
+    }
+  } else {
+    if (shouldEmitArgName)
+      llvm_unreachable("NYI HIPSaveKernelArgName");
+  }
+}

clang/lib/CIR/CodeGen/CIRGenModule.h

@@ -688,6 +688,20 @@ class CIRGenModule : public CIRGenTypeCache {
     return *openMPRuntime;
   }
 
+  /// OpenCL v1.2 s5.6.4.6 allows the compiler to store kernel argument
+  /// information in the program executable. The argument information stored
+  /// includes the argument name, its type, the address and access qualifiers
+  /// used. This helper can be used to generate metadata for source code kernel
+  /// function as well as generated implicitly kernels. If a kernel is generated
+  /// implicitly null value has to be passed to the last two parameters,
+  /// otherwise all parameters must have valid non-null values.
+  /// \param FN is a pointer to IR function being generated.
+  /// \param FD is a pointer to function declaration if any.
+  /// \param CGF is a pointer to CIRGenFunction that generates this function.
+  void genKernelArgMetadata(mlir::cir::FuncOp FN,
+                            const FunctionDecl *FD = nullptr,
+                            CIRGenFunction *CGF = nullptr);
+
 private:
   // An ordered map of canonical GlobalDecls to their mangled names.
   llvm::MapVector<clang::GlobalDecl, llvm::StringRef> MangledDeclNames;

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

@@ -554,6 +554,43 @@ LogicalResult OpenCLKernelMetadataAttr::verify(
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// OpenCLKernelArgMetadataAttr definitions
+//===----------------------------------------------------------------------===//
+
+LogicalResult OpenCLKernelArgMetadataAttr::verify(
+    ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError,
+    ArrayAttr addrSpaces, ArrayAttr accessQuals, ArrayAttr types,
+    ArrayAttr baseTypes, ArrayAttr typeQuals, ArrayAttr argNames) {
+  auto isIntArray = [](ArrayAttr elt) {
+    return llvm::all_of(
+        elt, [](Attribute elt) { return mlir::isa<IntegerAttr>(elt); });
+  };
+  auto isStrArray = [](ArrayAttr elt) {
+    return llvm::all_of(
+        elt, [](Attribute elt) { return mlir::isa<StringAttr>(elt); });
+  };
+
+  if (!isIntArray(addrSpaces))
+    return emitError() << "addr_space must be integer arrays";
+  if (!llvm::all_of<ArrayRef<ArrayAttr>>(
+          {accessQuals, types, baseTypes, typeQuals}, isStrArray))
+    return emitError()
+           << "access_qual, type, base_type, type_qual must be string arrays";
+  if (argNames && !isStrArray(argNames)) {
+    return emitError() << "name must be a string array";
+  }
+
+  if (!llvm::all_of<ArrayRef<ArrayAttr>>(
+          {addrSpaces, accessQuals, types, baseTypes, typeQuals, argNames},
+          [&](ArrayAttr arr) {
+            return !arr || arr.size() == addrSpaces.size();
+          })) {
+    return emitError() << "all arrays must have the same number of elements";
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // AddressSpaceAttr definitions
 //===----------------------------------------------------------------------===//