[mlir][LLVM] Add result attribute import support

[Dinistro]

Summary

This commit introduces support for importing result attributes.

Depends on D142372

[gysit]

mlir/lib/Target/LLVMIR/ModuleImport.cpp

1488

nit: I would probably inline this list onto the convertParameterAttribute method.

mlir/test/Target/LLVMIR/Import/function-attributes.ll

64

nit: maybe try if two of the result attributes can appear on one function by combining two tests?

[Dinistro]

address additional comments and extend tests

mlir/include/mlir/Target/LLVMIR/ModuleImport.h

LogicalResult convertBranchArgs(llvm::Instruction *branch,
                                   llvm::BasicBlock *target,
                                   SmallVectorImpl<Value> &blockArguments);
   /// Appends the converted result type and operands of `callInst` to the
   /// `types` and `operands` arrays. For indirect calls, the method additionally
   /// inserts the called function at the beginning of the `operands` array.
   LogicalResult convertCallTypeAndOperands(llvm::CallBase *callInst,
                                            SmallVectorImpl<Type> &types,
                                            SmallVectorImpl<Value> &operands);
+  /// Converts the parameter attributes attached to `func` and adds them to the
+  /// `funcOp`.
+  void convertParameterAttributes(llvm::Function *func, LLVMFuncOp funcOp,
+                                  OpBuilder &builder);
+  /// Converts the AttributeSet of one parameter in LLVM IR to a corresponding
+  /// DictionaryAttr for the LLVM dialect.
+  DictionaryAttr convertParameterAttribute(llvm::AttributeSet llvmParamAttrs,
+                                           OpBuilder &builder);
   /// Returns the builtin type equivalent to be used in attributes for the given
   /// LLVM IR dialect type.
   Type getStdTypeForAttr(Type type);
   /// Returns `value` as an attribute to attach to a GlobalOp.
   Attribute getConstantAsAttr(llvm::Constant *value);
   /// Returns the topologically sorted set of transitive dependencies needed to
   /// convert the given constant.
   SetVector<llvm::Constant *> getConstantsToConvert(llvm::Constant *constant);

mlir/lib/Target/LLVMIR/ModuleImport.cpp

 }
 
 void ModuleImport::processFunctionAttributes(llvm::Function *func,
                                              LLVMFuncOp funcOp) {
   processMemoryEffects(func, funcOp);
   processPassthroughAttrs(func, funcOp);
 }
 
+DictionaryAttr
+ModuleImport::convertParameterAttribute(llvm::AttributeSet llvmParamAttrs,
+                                        OpBuilder &builder) {
+  using ElemTy = std::pair<llvm::Attribute::AttrKind, StringRef>;
+  // Mapping from llvm attribute kinds to their corresponding MLIR name.
+  static const SmallVector<ElemTy> kindNamePairs = {
+      {llvm::Attribute::AttrKind::NoAlias, LLVMDialect::getNoAliasAttrName()},
+      {llvm::Attribute::AttrKind::ReadOnly, LLVMDialect::getReadonlyAttrName()},
+      {llvm::Attribute::AttrKind::Nest, LLVMDialect::getNestAttrName()},
+      {llvm::Attribute::AttrKind::SExt, LLVMDialect::getSExtAttrName()},
+      {llvm::Attribute::AttrKind::ZExt, LLVMDialect::getZExtAttrName()},
+      {llvm::Attribute::AttrKind::NoUndef, LLVMDialect::getNoUndefAttrName()},
+      {llvm::Attribute::AttrKind::StructRet,
+       LLVMDialect::getStructRetAttrName()},
+      {llvm::Attribute::AttrKind::ByVal, LLVMDialect::getByValAttrName()},
+      {llvm::Attribute::AttrKind::ByRef, LLVMDialect::getByRefAttrName()},
+      {llvm::Attribute::AttrKind::InAlloca, LLVMDialect::getInAllocaAttrName()},
+      {llvm::Attribute::AttrKind::Alignment, LLVMDialect::getAlignAttrName()}};
+
+  SmallVector<NamedAttribute> paramAttrs;
+  for (auto [llvmKind, mlirName] : kindNamePairs) {
+    auto llvmAttr = llvmParamAttrs.getAttribute(llvmKind);
+    // Skip attributes that are not attached.
+    if (!llvmAttr.isValid())
+      continue;
+    Attribute mlirAttr;
+    if (llvmAttr.isTypeAttribute())
+      mlirAttr = TypeAttr::get(convertType(llvmAttr.getValueAsType()));
+    else if (llvmAttr.isIntAttribute())
+      mlirAttr = builder.getI64IntegerAttr(llvmAttr.getValueAsInt());
+    else if (llvmAttr.isEnumAttribute())
+      mlirAttr = builder.getUnitAttr();
+    else
+      llvm_unreachable("unexpected parameter attribute kind");
+    paramAttrs.push_back(builder.getNamedAttr(mlirName, mlirAttr));
+  }
+
+  return builder.getDictionaryAttr(paramAttrs);
+}
+
+void ModuleImport::convertParameterAttributes(llvm::Function *func,
+                                              LLVMFuncOp funcOp,
+                                              OpBuilder &builder) {
+  auto llvmAttrs = func->getAttributes();
+  for (size_t i = 0, e = funcOp.getNumArguments(); i < e; ++i) {
+    llvm::AttributeSet llvmArgAttrs = llvmAttrs.getParamAttrs(i);
+    funcOp.setArgAttrs(i, convertParameterAttribute(llvmArgAttrs, builder));
+  }
+  // Convert the result attributes and attach them wrapped in an ArrayAttribute
+  // to the funcOp.
+  llvm::AttributeSet llvmResAttr = llvmAttrs.getRetAttrs();
+  funcOp.setResAttrsAttr(
+      builder.getArrayAttr(convertParameterAttribute(llvmResAttr, builder)));
+}
+
 LogicalResult ModuleImport::processFunction(llvm::Function *func) {
   clearBlockAndValueMapping();
 
   auto functionType =
       convertType(func->getFunctionType()).dyn_cast<LLVMFunctionType>();
   if (func->isIntrinsic() &&
       iface.isConvertibleIntrinsic(func->getIntrinsicID()))
     return success();
 
   bool dsoLocal = func->hasLocalLinkage();
   CConv cconv = convertCConvFromLLVM(func->getCallingConv());
 
   // Insert the function at the end of the module.
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPoint(mlirModule.getBody(), mlirModule.getBody()->end());
 
   LLVMFuncOp funcOp = builder.create<LLVMFuncOp>(
       mlirModule.getLoc(), func->getName(), functionType,
       convertLinkageFromLLVM(func->getLinkage()), dsoLocal, cconv);
 
   // Set the function debug information if available.
   debugImporter->translate(func, funcOp);
 
-  for (const auto &it : llvm::enumerate(functionType.getParams())) {
-    llvm::SmallVector<NamedAttribute, 1> argAttrs;
-    if (auto *type = func->getParamByValType(it.index())) {
-      Type mlirType = convertType(type);
-      argAttrs.push_back(
-          NamedAttribute(builder.getStringAttr(LLVMDialect::getByValAttrName()),
-                         TypeAttr::get(mlirType)));
-    }
-    if (auto *type = func->getParamByRefType(it.index())) {
-      Type mlirType = convertType(type);
-      argAttrs.push_back(
-          NamedAttribute(builder.getStringAttr(LLVMDialect::getByRefAttrName()),
-                         TypeAttr::get(mlirType)));
-    }
-    if (auto *type = func->getParamStructRetType(it.index())) {
-      Type mlirType = convertType(type);
-      argAttrs.push_back(NamedAttribute(
-          builder.getStringAttr(LLVMDialect::getStructRetAttrName()),
-          TypeAttr::get(mlirType)));
-    }
-    if (auto *type = func->getParamInAllocaType(it.index())) {
-      Type mlirType = convertType(type);
-      argAttrs.push_back(NamedAttribute(
-          builder.getStringAttr(LLVMDialect::getInAllocaAttrName()),
-          TypeAttr::get(mlirType)));
-    }
-
-    funcOp.setArgAttrs(it.index(), argAttrs);
-  }
+  convertParameterAttributes(func, funcOp, builder);
 
   if (FlatSymbolRefAttr personality = getPersonalityAsAttr(func))
     funcOp.setPersonalityAttr(personality);
   else if (func->hasPersonalityFn())
     emitWarning(funcOp.getLoc(), "could not deduce personality, skipping it");
 
   if (func->hasGC())
     funcOp.setGarbageCollector(StringRef(func->getGC()));

mlir/test/Target/LLVMIR/Import/function-attributes.ll

 
 ; // -----
 
 ; CHECK-LABEL: @func_arg_attrs
 ; CHECK-SAME:  !llvm.ptr {llvm.byval = i64}
 ; CHECK-SAME:  !llvm.ptr {llvm.byref = i64}
 ; CHECK-SAME:  !llvm.ptr {llvm.sret = i64}
 ; CHECK-SAME:  !llvm.ptr {llvm.inalloca = i64}
+; CHECK-SAME:  !llvm.ptr {llvm.noalias}
+; CHECK-SAME:  !llvm.ptr {llvm.readonly}
+; CHECK-SAME:  !llvm.ptr {llvm.nest}
+; CHECK-SAME:  i32 {llvm.signext}
+; CHECK-SAME:  i64 {llvm.zeroext}
+; CHECK-SAME:  !llvm.ptr {llvm.align = 64 : i64, llvm.noundef}
 define void @func_arg_attrs(
     ptr byval(i64) %arg0,
     ptr byref(i64) %arg1,
     ptr sret(i64) %arg2,
-    ptr inalloca(i64) %arg3) {
+    ptr inalloca(i64) %arg3,
+    ptr noalias %arg4,
+    ptr readonly %arg5,
+    ptr nest %arg6,
+    i32 signext %arg7,
+    i64 zeroext %arg8,
+    ptr align(64) noundef %arg9) {
   ret void
 }
 
 ; // -----
 
+; CHECK-LABEL: @func_res_attr_align
+; CHECK-SAME:  !llvm.ptr {llvm.align = 16 : i64}
+declare align(16) ptr @func_res_attr_align()
+
+; // -----
+
+; CHECK-LABEL: @func_res_attr_noalias
+; CHECK-SAME:  !llvm.ptr {llvm.noalias}
+declare noalias ptr @func_res_attr_noalias()
+
+; // -----
+
+; CHECK-LABEL: @func_res_attr_signext
+; CHECK-DAG: llvm.noundef
+; CHECK-DAG: llvm.signext
+declare noundef signext i32 @func_res_attr_signext()
+
+; // -----
+
+; CHECK-LABEL: @func_res_attr_zeroext
+; CHECK-SAME:  i32 {llvm.zeroext}
+declare zeroext i32 @func_res_attr_zeroext()
+
+
+; // -----
+
 ; CHECK-LABEL: @entry_count
 ; CHECK-SAME:  attributes {function_entry_count = 4242 : i64}
 define void @entry_count() !prof !1 {
   ret void
 }
 
 !1 = !{!"function_entry_count", i64 4242}