Skip to content

Commit 1377056

Browse files
committed
[mlir][LLVM] add argument and result attributes to llvm.call
1 parent b652f02 commit 1377056

File tree

12 files changed

+230
-41
lines changed

12 files changed

+230
-41
lines changed

llvm/include/llvm/IR/InstrTypes.h

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1490,6 +1490,11 @@ class CallBase : public Instruction {
14901490
Attrs = Attrs.addRetAttribute(getContext(), Attr);
14911491
}
14921492

1493+
/// Adds attributes to the return value.
1494+
void addRetAttrs(const AttrBuilder &B) {
1495+
Attrs = Attrs.addRetAttributes(getContext(), B);
1496+
}
1497+
14931498
/// Adds the attribute to the indicated argument
14941499
void addParamAttr(unsigned ArgNo, Attribute::AttrKind Kind) {
14951500
assert(ArgNo < arg_size() && "Out of bounds");
@@ -1502,6 +1507,12 @@ class CallBase : public Instruction {
15021507
Attrs = Attrs.addParamAttribute(getContext(), ArgNo, Attr);
15031508
}
15041509

1510+
/// Adds attributes to the indicated argument
1511+
void addParamAttrs(unsigned ArgNo, const AttrBuilder &B) {
1512+
assert(ArgNo < arg_size() && "Out of bounds");
1513+
Attrs = Attrs.addParamAttributes(getContext(), ArgNo, B);
1514+
}
1515+
15051516
/// removes the attribute from the list of attributes.
15061517
void removeAttributeAtIndex(unsigned i, Attribute::AttrKind Kind) {
15071518
Attrs = Attrs.removeAttributeAtIndex(getContext(), i, Kind);

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -755,7 +755,9 @@ def LLVM_CallOp : LLVM_MemAccessOpBase<"call",
755755
VariadicOfVariadic<LLVM_Type,
756756
"op_bundle_sizes">:$op_bundle_operands,
757757
DenseI32ArrayAttr:$op_bundle_sizes,
758-
OptionalAttr<ArrayAttr>:$op_bundle_tags);
758+
OptionalAttr<ArrayAttr>:$op_bundle_tags,
759+
OptionalAttr<DictArrayAttr>:$arg_attrs,
760+
OptionalAttr<DictArrayAttr>:$res_attrs);
759761
// Append the aliasing related attributes defined in LLVM_MemAccessOpBase.
760762
let arguments = !con(args, aliasAttrs);
761763
let results = (outs Optional<LLVM_Type>:$result);

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -326,14 +326,18 @@ class ModuleImport {
326326
SmallVectorImpl<Type> &types,
327327
SmallVectorImpl<Value> &operands,
328328
bool allowInlineAsm = false);
329-
/// Converts the parameter attributes attached to `func` and adds them to the
330-
/// `funcOp`.
329+
/// Converts the parameter and result attributes attached to `func` and adds
330+
/// them to the `funcOp`.
331331
void convertParameterAttributes(llvm::Function *func, LLVMFuncOp funcOp,
332332
OpBuilder &builder);
333333
/// Converts the AttributeSet of one parameter in LLVM IR to a corresponding
334334
/// DictionaryAttr for the LLVM dialect.
335335
DictionaryAttr convertParameterAttribute(llvm::AttributeSet llvmParamAttrs,
336336
OpBuilder &builder);
337+
/// Converts the parameter and result attributes attached to `call` and adds
338+
/// them to the `callOp`.
339+
void convertParameterAttributes(llvm::CallBase *call, CallOpInterface callOp,
340+
OpBuilder &builder);
337341
/// Returns the builtin type equivalent to the given LLVM dialect type or
338342
/// nullptr if there is no equivalent. The returned type can be used to create
339343
/// an attribute for a GlobalOp or a ConstantOp.

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,11 @@ class ModuleTranslation {
228228
/*recordInsertions=*/false);
229229
}
230230

231+
/// Translates parameter attributes of a call and adds them to the returned
232+
/// AttrBuilder. Returns failure if any of the translations failed.
233+
FailureOr<llvm::AttrBuilder> convertParameterAttrs(CallOp callOp, int argIdx,
234+
DictionaryAttr paramAttrs);
235+
231236
/// Gets the named metadata in the LLVM IR module being constructed, creating
232237
/// it if it does not exist.
233238
llvm::NamedMDNode *getOrInsertNamedModuleMetadata(StringRef name);
@@ -346,8 +351,8 @@ class ModuleTranslation {
346351
convertDialectAttributes(Operation *op,
347352
ArrayRef<llvm::Instruction *> instructions);
348353

349-
/// Translates parameter attributes and adds them to the returned AttrBuilder.
350-
/// Returns failure if any of the translations failed.
354+
/// Translates parameter attributes of a function and adds them to the
355+
/// returned AttrBuilder. Returns failure if any of the translations failed.
351356
FailureOr<llvm::AttrBuilder>
352357
convertParameterAttrs(LLVMFuncOp func, int argIdx, DictionaryAttr paramAttrs);
353358

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

Lines changed: 45 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1033,6 +1033,7 @@ void CallOp::build(OpBuilder &builder, OperationState &state, TypeRange results,
10331033
/*memory_effects=*/nullptr,
10341034
/*convergent=*/nullptr, /*no_unwind=*/nullptr, /*will_return=*/nullptr,
10351035
/*op_bundle_operands=*/{}, /*op_bundle_tags=*/{},
1036+
/*arg_attrs=*/nullptr, /*res_attrs=*/nullptr,
10361037
/*access_groups=*/nullptr, /*alias_scopes=*/nullptr,
10371038
/*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);
10381039
}
@@ -1060,6 +1061,7 @@ void CallOp::build(OpBuilder &builder, OperationState &state,
10601061
/*convergent=*/nullptr,
10611062
/*no_unwind=*/nullptr, /*will_return=*/nullptr,
10621063
/*op_bundle_operands=*/{}, /*op_bundle_tags=*/{},
1064+
/*arg_attrs=*/nullptr, /*res_attrs=*/nullptr,
10631065
/*access_groups=*/nullptr,
10641066
/*alias_scopes=*/nullptr, /*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);
10651067
}
@@ -1073,6 +1075,7 @@ void CallOp::build(OpBuilder &builder, OperationState &state,
10731075
/*CConv=*/nullptr, /*TailCallKind=*/nullptr, /*memory_effects=*/nullptr,
10741076
/*convergent=*/nullptr, /*no_unwind=*/nullptr, /*will_return=*/nullptr,
10751077
/*op_bundle_operands=*/{}, /*op_bundle_tags=*/{},
1078+
/*arg_attrs=*/nullptr, /*res_attrs=*/nullptr,
10761079
/*access_groups=*/nullptr, /*alias_scopes=*/nullptr,
10771080
/*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);
10781081
}
@@ -1087,6 +1090,7 @@ void CallOp::build(OpBuilder &builder, OperationState &state, LLVMFuncOp func,
10871090
/*convergent=*/nullptr, /*no_unwind=*/nullptr, /*will_return=*/nullptr,
10881091
/*op_bundle_operands=*/{}, /*op_bundle_tags=*/{},
10891092
/*access_groups=*/nullptr, /*alias_scopes=*/nullptr,
1093+
/*arg_attrs=*/nullptr, /*res_attrs=*/nullptr,
10901094
/*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);
10911095
}
10921096

@@ -1331,55 +1335,65 @@ void CallOp::print(OpAsmPrinter &p) {
13311335
getVarCalleeTypeAttrName(), getCConvAttrName(),
13321336
getOperandSegmentSizesAttrName(),
13331337
getOpBundleSizesAttrName(),
1334-
getOpBundleTagsAttrName()});
1338+
getOpBundleTagsAttrName(), getArgAttrsAttrName(),
1339+
getResAttrsAttrName()});
13351340

13361341
p << " : ";
13371342
if (!isDirect)
13381343
p << getOperand(0).getType() << ", ";
13391344

13401345
// Reconstruct the function MLIR function type from operand and result types.
1341-
p.printFunctionalType(args.getTypes(), getResultTypes());
1346+
call_interface_impl::printFunctionSignature(
1347+
p, *this, args.getTypes(), /*isVariadic=*/false, getResultTypes());
13421348
}
13431349

13441350
/// Parses the type of a call operation and resolves the operands if the parsing
13451351
/// succeeds. Returns failure otherwise.
13461352
static ParseResult parseCallTypeAndResolveOperands(
13471353
OpAsmParser &parser, OperationState &result, bool isDirect,
1348-
ArrayRef<OpAsmParser::UnresolvedOperand> operands) {
1354+
ArrayRef<OpAsmParser::UnresolvedOperand> operands,
1355+
SmallVectorImpl<DictionaryAttr> &argAttrs,
1356+
SmallVectorImpl<DictionaryAttr> &resultAttrs) {
13491357
SMLoc trailingTypesLoc = parser.getCurrentLocation();
13501358
SmallVector<Type> types;
1351-
if (parser.parseColonTypeList(types))
1359+
if (parser.parseColon())
13521360
return failure();
1353-
1354-
if (isDirect && types.size() != 1)
1355-
return parser.emitError(trailingTypesLoc,
1356-
"expected direct call to have 1 trailing type");
1357-
if (!isDirect && types.size() != 2)
1358-
return parser.emitError(trailingTypesLoc,
1359-
"expected indirect call to have 2 trailing types");
1360-
1361-
auto funcType = llvm::dyn_cast<FunctionType>(types.pop_back_val());
1362-
if (!funcType)
1361+
if (!isDirect) {
1362+
types.emplace_back();
1363+
if (parser.parseType(types.back()))
1364+
return failure();
1365+
if (parser.parseOptionalComma())
1366+
return parser.emitError(
1367+
trailingTypesLoc, "expected indirect call to have 2 trailing types");
1368+
}
1369+
SmallVector<Type> argTypes;
1370+
SmallVector<Type> resTypes;
1371+
if (call_interface_impl::parseFunctionSignature(parser, argTypes, argAttrs,
1372+
resTypes, resultAttrs)) {
1373+
if (isDirect)
1374+
return parser.emitError(trailingTypesLoc,
1375+
"expected direct call to have 1 trailing types");
13631376
return parser.emitError(trailingTypesLoc,
13641377
"expected trailing function type");
1365-
if (funcType.getNumResults() > 1)
1378+
}
1379+
1380+
if (resTypes.size() > 1)
13661381
return parser.emitError(trailingTypesLoc,
13671382
"expected function with 0 or 1 result");
1368-
if (funcType.getNumResults() == 1 &&
1369-
llvm::isa<LLVM::LLVMVoidType>(funcType.getResult(0)))
1383+
if (resTypes.size() == 1 && llvm::isa<LLVM::LLVMVoidType>(resTypes[0]))
13701384
return parser.emitError(trailingTypesLoc,
13711385
"expected a non-void result type");
13721386

13731387
// The head element of the types list matches the callee type for
13741388
// indirect calls, while the types list is emtpy for direct calls.
13751389
// Append the function input types to resolve the call operation
13761390
// operands.
1377-
llvm::append_range(types, funcType.getInputs());
1391+
llvm::append_range(types, argTypes);
13781392
if (parser.resolveOperands(operands, types, parser.getNameLoc(),
13791393
result.operands))
13801394
return failure();
1381-
if (funcType.getNumResults() != 0)
1382-
result.addTypes(funcType.getResults());
1395+
if (resTypes.size() != 0)
1396+
result.addTypes(resTypes);
13831397

13841398
return success();
13851399
}
@@ -1493,8 +1507,14 @@ ParseResult CallOp::parse(OpAsmParser &parser, OperationState &result) {
14931507
return failure();
14941508

14951509
// Parse the trailing type list and resolve the operands.
1496-
if (parseCallTypeAndResolveOperands(parser, result, isDirect, operands))
1510+
SmallVector<DictionaryAttr> argAttrs;
1511+
SmallVector<DictionaryAttr> resultAttrs;
1512+
if (parseCallTypeAndResolveOperands(parser, result, isDirect, operands,
1513+
argAttrs, resultAttrs))
14971514
return failure();
1515+
call_interface_impl::addArgAndResultAttrs(
1516+
parser.getBuilder(), result, argAttrs, resultAttrs,
1517+
getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));
14981518
if (resolveOpBundleOperands(parser, opBundlesLoc, result, opBundleOperands,
14991519
opBundleOperandTypes,
15001520
getOpBundleSizesAttrName(result.name)))
@@ -1714,7 +1734,10 @@ ParseResult InvokeOp::parse(OpAsmParser &parser, OperationState &result) {
17141734
return failure();
17151735

17161736
// Parse the trailing type list and resolve the function operands.
1717-
if (parseCallTypeAndResolveOperands(parser, result, isDirect, operands))
1737+
SmallVector<DictionaryAttr> argAttrs;
1738+
SmallVector<DictionaryAttr> resultAttrs;
1739+
if (parseCallTypeAndResolveOperands(parser, result, isDirect, operands,
1740+
argAttrs, resultAttrs))
17181741
return failure();
17191742
if (resolveOpBundleOperands(parser, opBundlesLoc, result, opBundleOperands,
17201743
opBundleOperandTypes,

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,27 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
265265
if (callOp.getWillReturnAttr())
266266
call->addFnAttr(llvm::Attribute::WillReturn);
267267

268+
if (ArrayAttr argAttrsArray = callOp.getArgAttrsAttr())
269+
for (auto [argIdx, argAttrsAttr] : llvm::enumerate(argAttrsArray)) {
270+
if (auto argAttrs = llvm::cast<DictionaryAttr>(argAttrsAttr)) {
271+
FailureOr<llvm::AttrBuilder> attrBuilder =
272+
moduleTranslation.convertParameterAttrs(callOp, argIdx, argAttrs);
273+
if (failed(attrBuilder))
274+
return failure();
275+
call->addParamAttrs(argIdx, *attrBuilder);
276+
}
277+
}
278+
279+
ArrayAttr resAttrsArray = callOp.getResAttrsAttr();
280+
if (resAttrsArray && resAttrsArray.size() == 1)
281+
if (auto resAttrs = llvm::cast<DictionaryAttr>(resAttrsArray[0])) {
282+
FailureOr<llvm::AttrBuilder> attrBuilder =
283+
moduleTranslation.convertParameterAttrs(callOp, -1, resAttrs);
284+
if (failed(attrBuilder))
285+
return failure();
286+
call->addRetAttrs(*attrBuilder);
287+
}
288+
268289
if (MemoryEffectsAttr memAttr = callOp.getMemoryEffectsAttr()) {
269290
llvm::MemoryEffects memEffects =
270291
llvm::MemoryEffects(llvm::MemoryEffects::Location::ArgMem,

mlir/lib/Target/LLVMIR/ModuleImport.cpp

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1641,6 +1641,9 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
16411641
if (callInst->hasFnAttr(llvm::Attribute::WillReturn))
16421642
callOp.setWillReturn(true);
16431643

1644+
// Handle parameter and result attributes.
1645+
convertParameterAttributes(callInst, callOp, builder);
1646+
16441647
llvm::MemoryEffects memEffects = callInst->getMemoryEffects();
16451648
ModRefInfo othermem = convertModRefInfoFromLLVM(
16461649
memEffects.getModRef(llvm::MemoryEffects::Location::Other));
@@ -2084,6 +2087,38 @@ void ModuleImport::convertParameterAttributes(llvm::Function *func,
20842087
builder.getArrayAttr(convertParameterAttribute(llvmResAttr, builder)));
20852088
}
20862089

2090+
void ModuleImport::convertParameterAttributes(llvm::CallBase *call,
2091+
CallOpInterface callOp,
2092+
OpBuilder &builder) {
2093+
auto llvmAttrs = call->getAttributes();
2094+
SmallVector<llvm::AttributeSet> llvmArgAttrsSet;
2095+
bool anyArgAttrs = false;
2096+
for (size_t i = 0, e = call->arg_size(); i < e; ++i) {
2097+
llvmArgAttrsSet.emplace_back(llvmAttrs.getParamAttrs(i));
2098+
if (llvmArgAttrsSet.back().hasAttributes())
2099+
anyArgAttrs = true;
2100+
}
2101+
auto getArrayAttr = [&](ArrayRef<DictionaryAttr> dictAttrs) {
2102+
SmallVector<Attribute> attrs;
2103+
for (auto &dict : dictAttrs)
2104+
attrs.push_back(dict ? dict : builder.getDictionaryAttr({}));
2105+
return builder.getArrayAttr(attrs);
2106+
};
2107+
if (anyArgAttrs) {
2108+
SmallVector<DictionaryAttr> argAttrs;
2109+
for (auto &llvmArgAttrs : llvmArgAttrsSet)
2110+
argAttrs.emplace_back(convertParameterAttribute(llvmArgAttrs, builder));
2111+
callOp.setArgAttrsAttr(getArrayAttr(argAttrs));
2112+
}
2113+
2114+
llvm::AttributeSet llvmResAttr = llvmAttrs.getRetAttrs();
2115+
if (!llvmResAttr.hasAttributes())
2116+
return;
2117+
SmallVector<DictionaryAttr, 1> resAttrs;
2118+
resAttrs.emplace_back(convertParameterAttribute(llvmResAttr, builder));
2119+
callOp.setResAttrsAttr(getArrayAttr(resAttrs));
2120+
}
2121+
20872122
LogicalResult ModuleImport::processFunction(llvm::Function *func) {
20882123
clearRegionState();
20892124

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

Lines changed: 38 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1563,6 +1563,26 @@ static void convertFunctionKernelAttributes(LLVMFuncOp func,
15631563
}
15641564
}
15651565

1566+
static void convertParameterAttr(llvm::AttrBuilder &attrBuilder,
1567+
llvm::Attribute::AttrKind llvmKind,
1568+
NamedAttribute namedAttr,
1569+
ModuleTranslation &moduleTranslation) {
1570+
llvm::TypeSwitch<Attribute>(namedAttr.getValue())
1571+
.Case<TypeAttr>([&](auto typeAttr) {
1572+
attrBuilder.addTypeAttr(
1573+
llvmKind, moduleTranslation.convertType(typeAttr.getValue()));
1574+
})
1575+
.Case<IntegerAttr>([&](auto intAttr) {
1576+
attrBuilder.addRawIntAttr(llvmKind, intAttr.getInt());
1577+
})
1578+
.Case<UnitAttr>([&](auto) { attrBuilder.addAttribute(llvmKind); })
1579+
.Case<LLVM::ConstantRangeAttr>([&](auto rangeAttr) {
1580+
attrBuilder.addConstantRangeAttr(
1581+
llvmKind,
1582+
llvm::ConstantRange(rangeAttr.getLower(), rangeAttr.getUpper()));
1583+
});
1584+
}
1585+
15661586
FailureOr<llvm::AttrBuilder>
15671587
ModuleTranslation::convertParameterAttrs(LLVMFuncOp func, int argIdx,
15681588
DictionaryAttr paramAttrs) {
@@ -1573,20 +1593,7 @@ ModuleTranslation::convertParameterAttrs(LLVMFuncOp func, int argIdx,
15731593
auto it = attrNameToKindMapping.find(namedAttr.getName());
15741594
if (it != attrNameToKindMapping.end()) {
15751595
llvm::Attribute::AttrKind llvmKind = it->second;
1576-
1577-
llvm::TypeSwitch<Attribute>(namedAttr.getValue())
1578-
.Case<TypeAttr>([&](auto typeAttr) {
1579-
attrBuilder.addTypeAttr(llvmKind, convertType(typeAttr.getValue()));
1580-
})
1581-
.Case<IntegerAttr>([&](auto intAttr) {
1582-
attrBuilder.addRawIntAttr(llvmKind, intAttr.getInt());
1583-
})
1584-
.Case<UnitAttr>([&](auto) { attrBuilder.addAttribute(llvmKind); })
1585-
.Case<LLVM::ConstantRangeAttr>([&](auto rangeAttr) {
1586-
attrBuilder.addConstantRangeAttr(
1587-
llvmKind, llvm::ConstantRange(rangeAttr.getLower(),
1588-
rangeAttr.getUpper()));
1589-
});
1596+
convertParameterAttr(attrBuilder, llvmKind, namedAttr, *this);
15901597
} else if (namedAttr.getNameDialect()) {
15911598
if (failed(iface.convertParameterAttr(func, argIdx, namedAttr, *this)))
15921599
return failure();
@@ -1596,6 +1603,23 @@ ModuleTranslation::convertParameterAttrs(LLVMFuncOp func, int argIdx,
15961603
return attrBuilder;
15971604
}
15981605

1606+
FailureOr<llvm::AttrBuilder>
1607+
ModuleTranslation::convertParameterAttrs(CallOp, int argIdx,
1608+
DictionaryAttr paramAttrs) {
1609+
llvm::AttrBuilder attrBuilder(llvmModule->getContext());
1610+
auto attrNameToKindMapping = getAttrNameToKindMapping();
1611+
1612+
for (auto namedAttr : paramAttrs) {
1613+
auto it = attrNameToKindMapping.find(namedAttr.getName());
1614+
if (it != attrNameToKindMapping.end()) {
1615+
llvm::Attribute::AttrKind llvmKind = it->second;
1616+
convertParameterAttr(attrBuilder, llvmKind, namedAttr, *this);
1617+
}
1618+
}
1619+
1620+
return attrBuilder;
1621+
}
1622+
15991623
LogicalResult ModuleTranslation::convertFunctionSignatures() {
16001624
// Declare all functions first because there may be function calls that form a
16011625
// call graph with cycles, or global initializers that reference functions.

mlir/test/Dialect/LLVMIR/invalid.mlir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,7 @@ func.func @call_missing_ptr_type(%callee : !llvm.func<i8 (i8)>, %arg : i8) {
235235
func.func private @standard_func_callee()
236236

237237
func.func @call_missing_ptr_type(%arg : i8) {
238+
// expected-error@+2 {{expected '('}}
238239
// expected-error@+1 {{expected direct call to have 1 trailing type}}
239240
llvm.call @standard_func_callee(%arg) : !llvm.ptr, (i8) -> (i8)
240241
llvm.return
@@ -251,6 +252,7 @@ func.func @call_non_pointer_type(%callee : !llvm.func<i8 (i8)>, %arg : i8) {
251252
// -----
252253

253254
func.func @call_non_function_type(%callee : !llvm.ptr, %arg : i8) {
255+
// expected-error@+2 {{expected '('}}
254256
// expected-error@+1 {{expected trailing function type}}
255257
llvm.call %callee(%arg) : !llvm.ptr, !llvm.func<i8 (i8)>
256258
llvm.return

0 commit comments

Comments
 (0)