Skip to content

Commit

Permalink
Generalize actions to functions. Implement function lowering
Browse files Browse the repository at this point in the history 
Signed-off-by: Anton Korobeynikov <[email protected]>
  • Loading branch information
@asl
asl committed Feb 10, 2025
1 parent 041ed80 commit 1f7501c
Show file tree
Hide file tree
Showing 17 changed files with 521 additions and 145 deletions.
75 changes: 60 additions & 15 deletions include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ def AllocaOp : P4HIR_Op<"alloca", [

let skipDefaultBuilders = 1;
let builders = [
OpBuilder<(ins "mlir::Type":$ref, "mlir::Type":$objectType, "llvm::StringRef":$name)>,
OpBuilder<(ins "mlir::Type":$ref, "mlir::Type":$objectType, "const llvm::Twine &":$name)>,
];

let assemblyFormat = [{
Expand Down Expand Up @@ -534,11 +534,16 @@ def IfOp : P4HIR_Op<"if",
}

def ReturnOp : P4HIR_Op<"return", [ParentOneOf<["ScopeOp", "IfOp",
"ActionOp",
"FuncOp",
// "SwitchOp", "CaseOp",
// "ForInOp", "ForOp",
]>,
Terminator, ReturnLike]> {
// Note that ReturnOp is not ReturnLike: currently there is no way to
// represent early exits in MLIR "properly"
// We might not be able to have it a Terminator at this level in order
// to represent dead code. We might lower it to proper terminator later (!)
// See https://discourse.llvm.org/t/rfc-region-based-control-flow-with-early-exits-in-mlir/76998
Terminator]> {
let summary = "Return from function or action";
let description = [{
The "return" operation represents a return operation within a function or action.
Expand Down Expand Up @@ -579,13 +584,12 @@ def ReturnOp : P4HIR_Op<"return", [ParentOneOf<["ScopeOp", "IfOp",
let hasVerifier = 1;
}

def ActionOp : P4HIR_Op<"action", [
def FuncOp : P4HIR_Op<"func", [
AutomaticAllocationScope, CallableOpInterface, FunctionOpInterface,
IsolatedFromAbove
]> {
let summary = "Define an action";
let summary = "Define a function-like object (action, function)";
let description = [{

Similar to `mlir::FuncOp` built-in:
> Operations within the function cannot implicitly capture values defined
> outside of the function, i.e. Functions are `IsolatedFromAbove`. All
Expand All @@ -600,10 +604,13 @@ def ActionOp : P4HIR_Op<"action", [
> Only dialect attribute names may be specified in the attribute dictionaries
> for function arguments, results, or the function itself.

Action parameters might have direction that is specified via `p4hir.dir`
Parameters might have direction that is specified via `p4hir.dir`
attribute. Out and inout parameters must have a reference type. All
refence-typed parameters must have a direction and it should be `out` or `input`.

An action must be marked as `action`, should always have a body and cannot return
anything.

Example:

```mlir
Expand All @@ -624,22 +631,55 @@ def ActionOp : P4HIR_Op<"action", [
}];

let arguments = (ins SymbolNameAttr:$sym_name,
TypeAttrOf<P4HIR_ActionType>:$function_type,
TypeAttrOf<FuncType>:$function_type,
UnitAttr:$action,
OptionalAttr<DictArrayAttr>:$arg_attrs,
OptionalAttr<DictArrayAttr>:$res_attrs,
OptionalAttr<ArrayAttr>:$annotations);
let regions = (region AnyRegion:$body);
let skipDefaultBuilders = 1;

let builders = [
OpBuilder<(ins "llvm::StringRef":$name, "ActionType":$type,
OpBuilder<(ins "llvm::StringRef":$name, "FuncType":$type,
CArg<"llvm::ArrayRef<mlir::NamedAttribute>", "{}">:$attrs,
CArg<"llvm::ArrayRef<mlir::DictionaryAttr>", "{}">:$argAttrs)>
];

let extraClassDeclaration = [{
// TODO: move to custom builder
static FuncOp buildAction(mlir::OpBuilder &builder,
mlir::Location loc,
llvm::StringRef name,
P4HIR::FuncType type,
llvm::ArrayRef<mlir::NamedAttribute> attrs = {},
llvm::ArrayRef<mlir::DictionaryAttr> argAttrs = {}) {
auto op = builder.create<FuncOp>(loc, name, type, attrs, argAttrs);
op.createEntryBlock();
op.setAction(true);
return op;
}

/// Returns the region on the current operation that is callable. Always
/// non-null for actions.
mlir::Region *getCallableRegion();
mlir::Region *getCallableRegion() { return isExternal() ? nullptr : &getBody(); }

/// Returns the results types that the callable region produces when
/// executed.
llvm::ArrayRef<mlir::Type> getCallableResults() {
return getFunctionType().getReturnTypes();
}

/// Returns the argument attributes for all callable region arguments or
/// null if there are none.
::mlir::ArrayAttr getCallableArgAttrs() {
return getArgAttrs().value_or(nullptr);
}

/// Returns the result attributes for all callable region results or null if
/// there are none.
::mlir::ArrayAttr getCallableResAttrs() {
return getResAttrs().value_or(nullptr);
}

/// Returns the argument types of this action.
llvm::ArrayRef<mlir::Type> getArgumentTypes() {
Expand Down Expand Up @@ -667,17 +707,21 @@ def ActionOp : P4HIR_Op<"action", [
return ParamDirection::None;
}

/// Returns the result types of this action. Required for FunctionOp
/// interface, so return empty.
/// Returns 0 or 1 result type of this function (0 in the case of a function
/// returing void or action)
llvm::ArrayRef<mlir::Type> getResultTypes() {
return {};
return getFunctionType().getReturnTypes();
}

/// Hook for OpTrait::FunctionOpInterfaceTrait, called after verifying that
/// the 'type' attribute is present and checks if it holds a function type.
/// Ensures getType, getNumFuncArguments, and getNumFuncResults can be
/// called safely.
llvm::LogicalResult verifyType();

bool isDeclaration() { return isExternal(); }

void createEntryBlock();
}];

let hasCustomAssemblyFormat = 1;
Expand Down Expand Up @@ -706,7 +750,7 @@ def CallOp : P4HIR_Op<"call",
}];

// TODO: Refine result types, refine parameter type
let results = (outs Optional<AnyP4Type>:$result);
let results = (outs Optional<CallResultP4Type>:$result);
let arguments = (ins OptionalAttr<FlatSymbolRefAttr>:$callee, Variadic<AnyType>:$operands);

let skipDefaultBuilders = 1;
Expand All @@ -718,7 +762,8 @@ def CallOp : P4HIR_Op<"call",
CArg<"mlir::ValueRange", "{}">:$operands), [{
$_state.addOperands(operands);
$_state.addAttribute("callee", callee);
$_state.addTypes(resType);
if (resType && !isa<VoidType>(resType))
$_state.addTypes(resType);
}]>,
// Everything else that does not produce result
OpBuilder<(ins "mlir::SymbolRefAttr":$callee,
Expand Down
61 changes: 47 additions & 14 deletions include/p4mlir/Dialect/P4HIR/P4HIR_Types.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,15 @@ def DontcareType : P4HIR_Type<"Dontcare", "dontcare"> {}
def ErrorType : P4HIR_Type<"Error", "error"> {}
def UnknownType : P4HIR_Type<"Unknown", "unknown"> {}

def VoidType : P4HIR_Type<"Void", "void"> {
let summary = "void type";
let description = [{
Represents absense of result of actions and methods, or `void` type for functions.
}];
let extraClassDeclaration = [{
std::string getAlias() const { return "void"; };
}];
}
//===----------------------------------------------------------------------===//
// ReferenceType
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -130,29 +139,40 @@ def ReferenceType : P4HIR_Type<"Reference", "ref"> {
let skipDefaultBuilders = 1;
}

//===----------------------------------------------------------------------===//
// ParameterType
//===----------------------------------------------------------------------===//

def P4HIR_ActionType : P4HIR_Type<"Action", "action"> {
let summary = "P4 action type";
def FuncType : P4HIR_Type<"Func", "func"> {
let summary = "P4 function-like type (actions, methods, functions)";
let description = [{
The `!p4hir.action` is an action type. It is essentially a list of parameter
types. There is no return type for actions.

The `!p4hir.func` is a function type.
Example:

```mlir
!p4hir.action<()>
!p4hir.action<(!p4hir.bit<32>, !p4hir.int<42>)>
!p4hir.func<()>
!p4hir.func<!p4hir.bit<32>(!p4hir.bit<32>, !p4hir.int<42>)>
```
}];

let parameters = (ins ArrayRefParameter<"mlir::Type">:$inputs);
let parameters = (ins ArrayRefParameter<"mlir::Type">:$inputs,
"mlir::Type":$optionalReturnType);

let builders = [
// Construct with an actual return type or explicit !p4hir.void
TypeBuilderWithInferredContext<(ins
"llvm::ArrayRef<mlir::Type>":$inputs, "mlir::Type":$returnType), [{
return $_get(returnType.getContext(), inputs,
mlir::isa<P4HIR::VoidType>(returnType) ? nullptr
: returnType);
}]>,

// Construct without return type
TypeBuilder<(ins "llvm::ArrayRef<mlir::Type>":$inputs), [{
return $_get($_ctxt, inputs, nullptr);
}]>

];

// Use a custom parser to handle the argument types in better way.
let assemblyFormat = [{
`<` custom<ActionType>($inputs) `>`
`<` custom<FuncType>($optionalReturnType, $inputs) `>`
}];

let extraClassDeclaration = [{
Expand All @@ -162,9 +182,20 @@ def P4HIR_ActionType : P4HIR_Type<"Action", "action"> {
/// Returns the number of arguments to the function.
unsigned getNumInputs() const { return getInputs().size(); }

/// Returns the result type of the function as an actual return type or
/// explicit !p4hir.void
mlir::Type getReturnType() const;

/// Returns the result type of the function as an ArrayRef, enabling better
/// integration with generic MLIR utilities.
llvm::ArrayRef<mlir::Type> getReturnTypes() const;

/// Returns a clone of this action type with the given argument
/// and result types. Required for FunctionOp interface
ActionType clone(mlir::TypeRange inputs, mlir::TypeRange outputs) const;
FuncType clone(mlir::TypeRange inputs, mlir::TypeRange outputs) const;

/// Returns whether the function returns void.
bool isVoid() const;
}];
}

Expand All @@ -174,6 +205,8 @@ def P4HIR_ActionType : P4HIR_Type<"Action", "action"> {

def AnyP4Type : AnyTypeOf<[BitsType, BooleanType, InfIntType,
DontcareType, ErrorType, UnknownType]> {}
def CallResultP4Type : AnyTypeOf<[BitsType, BooleanType, InfIntType, VoidType]> {}
def LoadableP4Type : AnyTypeOf<[BitsType, BooleanType, InfIntType]> {}


#endif // P4MLIR_DIALECT_P4HIR_P4HIR_TYPES_TD
Loading

0 comments on commit 1f7501c

Please sign in to comment.