Skip to content

Commit

Permalink
feat: Handle no_predicates attribute (#4942)
Browse files Browse the repository at this point in the history
# Description

## Problem\*

Resolves #4911 and
#4688

## Summary\*

~~We recently included the `#[inline(never)]` attribute to enable
developers to optimize codegen.~~ This has now been switched to the name
`no_predicates`. The main use-case in mind is for circuits in issue
#4688 where inlining a function with heavy array operations when
dependent upon witnesses is not always ideal. Specifically when the
function being inlined does not need to rely on the predicate for
correctness.

Originally I had in mind to delay inlining all the way to after ACIR gen
and inline the ACIR artifacts. However, this feels overly complex now as
we have all the infrastructure to inline functions as we wish during
SSA, we could just need to delay the inlining of certain functions to
happen after flattening. This PR does exactly what was just mentioned.

For example, the new test `no_predicates_numeric_generic_poseidon` gave
these results when `poseidon_hash` was not marked with
`#[no_predicates_numeric_generic_poseidon]`:
<img width="785" alt="ExistingInlineStrategy"
src="https://github.com/noir-lang/noir/assets/43554004/f2fc1358-c86c-4f02-999e-414056b87a01">

While when `poseidon_hash` was marked with
`#[no_predicates_numeric_generic_poseidon]`:
<img width="788" alt="InlineNeverBench"
src="https://github.com/noir-lang/noir/assets/43554004/21d729f9-32db-4a32-b592-56f76bf5663d">


## Additional Context



## Documentation\*

Check one:
- [ ] No documentation needed.
- [ ] Documentation included in this PR.
- [X] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [X] I have tested the changes locally.
- [X] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
  • Loading branch information
vezenovm committed Apr 30, 2024
1 parent aaac0f6 commit 0ce04d3
Show file tree
Hide file tree
Showing 19 changed files with 164 additions and 62 deletions.
3 changes: 3 additions & 0 deletions compiler/noirc_evaluator/src/ssa.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,9 @@ pub(crate) fn optimize_into_acir(
.run_pass(Ssa::remove_bit_shifts, "After Removing Bit Shifts:")
// Run mem2reg once more with the flattened CFG to catch any remaining loads/stores
.run_pass(Ssa::mem2reg, "After Mem2Reg:")
// Run the inlining pass again to handle functions with `InlineType::NoPredicates`.
// Before flattening is run, we treat functions marked with the `InlineType::NoPredicates` as an entry point.
.run_pass(Ssa::inline_functions_with_no_predicates, "After Inlining:")
.run_pass(Ssa::fold_constants, "After Constant Folding:")
.run_pass(Ssa::remove_enable_side_effects, "After EnableSideEffects removal:")
.run_pass(Ssa::fold_constants_using_constraints, "After Constraint Folding:")
Expand Down
17 changes: 12 additions & 5 deletions compiler/noirc_evaluator/src/ssa/acir_gen/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -390,7 +390,10 @@ impl<'a> Context<'a> {
panic!("ACIR function should have been inlined earlier if not marked otherwise");
}
}
InlineType::Fold | InlineType::Never => {}
InlineType::NoPredicates => {
panic!("All ACIR functions marked with #[no_predicates] should be inlined before ACIR gen. This is an SSA exclusive codegen attribute");
}
InlineType::Fold => {}
}
// We only want to convert entry point functions. This being `main` and those marked with `InlineType::Fold`
Ok(Some(self.convert_acir_main(function, ssa, brillig)?))
Expand Down Expand Up @@ -2653,10 +2656,14 @@ mod test {
basic_call_with_outputs_assert(InlineType::Fold);
call_output_as_next_call_input(InlineType::Fold);
basic_nested_call(InlineType::Fold);
}

call_output_as_next_call_input(InlineType::Never);
basic_nested_call(InlineType::Never);
basic_call_with_outputs_assert(InlineType::Never);
#[test]
#[should_panic]
fn basic_calls_no_predicates() {
call_output_as_next_call_input(InlineType::NoPredicates);
basic_nested_call(InlineType::NoPredicates);
basic_call_with_outputs_assert(InlineType::NoPredicates);
}

#[test]
Expand Down Expand Up @@ -2795,7 +2802,7 @@ mod test {
let (acir_functions, _) = ssa
.into_acir(&Brillig::default(), noirc_frontend::ast::Distinctness::Distinct)
.expect("Should compile manually written SSA into ACIR");
// The expected result should look very similar to the abvoe test expect that the input witnesses of the `Call`
// The expected result should look very similar to the above test expect that the input witnesses of the `Call`
// opcodes will be different. The changes can discerned from the checks below.

let main_acir = &acir_functions[0];
Expand Down
7 changes: 7 additions & 0 deletions compiler/noirc_evaluator/src/ssa/ir/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,13 @@ impl Function {
self.runtime = runtime;
}

pub(crate) fn is_no_predicates(&self) -> bool {
match self.runtime() {
RuntimeType::Acir(inline_type) => matches!(inline_type, InlineType::NoPredicates),
RuntimeType::Brillig => false,
}
}

/// Retrieves the entry block of a function.
///
/// A function's entry block contains the instructions
Expand Down
68 changes: 56 additions & 12 deletions compiler/noirc_evaluator/src/ssa/opt/inlining.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,30 @@ impl Ssa {
/// changes. This is because if the function's id later becomes known by a later
/// pass, we would need to re-run all of inlining anyway to inline it, so we might
/// as well save the work for later instead of performing it twice.
///
/// There are some attributes that allow inlining a function at a different step of codegen.
/// Currently this is just `InlineType::NoPredicates` for which we have a flag indicating
/// whether treating that inline functions. The default is to treat these functions as entry points.
#[tracing::instrument(level = "trace", skip(self))]
pub(crate) fn inline_functions(mut self) -> Ssa {
self.functions = btree_map(get_entry_point_functions(&self), |entry_point| {
let new_function = InlineContext::new(&self, entry_point).inline_all(&self);
(entry_point, new_function)
});
pub(crate) fn inline_functions(self) -> Ssa {
Self::inline_functions_inner(self, true)
}

// Run the inlining pass where functions marked with `InlineType::NoPredicates` as not entry points
pub(crate) fn inline_functions_with_no_predicates(self) -> Ssa {
Self::inline_functions_inner(self, false)
}

fn inline_functions_inner(mut self, no_predicates_is_entry_point: bool) -> Ssa {
self.functions = btree_map(
get_entry_point_functions(&self, no_predicates_is_entry_point),
|entry_point| {
let new_function =
InlineContext::new(&self, entry_point, no_predicates_is_entry_point)
.inline_all(&self);
(entry_point, new_function)
},
);
self
}
}
Expand All @@ -60,6 +77,8 @@ struct InlineContext {

// The FunctionId of the entry point function we're inlining into in the old, unmodified Ssa.
entry_point: FunctionId,

no_predicates_is_entry_point: bool,
}

/// The per-function inlining context contains information that is only valid for one function.
Expand Down Expand Up @@ -97,10 +116,19 @@ struct PerFunctionContext<'function> {
/// should be left in the final program.
/// This is the `main` function, any Acir functions with a [fold inline type][InlineType::Fold],
/// and any brillig functions used.
fn get_entry_point_functions(ssa: &Ssa) -> BTreeSet<FunctionId> {
fn get_entry_point_functions(
ssa: &Ssa,
no_predicates_is_entry_point: bool,
) -> BTreeSet<FunctionId> {
let functions = ssa.functions.iter();
let mut entry_points = functions
.filter(|(_, function)| function.runtime().is_entry_point())
.filter(|(_, function)| {
// If we have not already finished the flattening pass, functions marked
// to not have predicates should be marked as entry points.
let no_predicates_is_entry_point =
no_predicates_is_entry_point && function.is_no_predicates();
function.runtime().is_entry_point() || no_predicates_is_entry_point
})
.map(|(id, _)| *id)
.collect::<BTreeSet<_>>();

Expand All @@ -114,11 +142,21 @@ impl InlineContext {
/// The function being inlined into will always be the main function, although it is
/// actually a copy that is created in case the original main is still needed from a function
/// that could not be inlined calling it.
fn new(ssa: &Ssa, entry_point: FunctionId) -> InlineContext {
fn new(
ssa: &Ssa,
entry_point: FunctionId,
no_predicates_is_entry_point: bool,
) -> InlineContext {
let source = &ssa.functions[&entry_point];
let mut builder = FunctionBuilder::new(source.name().to_owned(), entry_point);
builder.set_runtime(source.runtime());
Self { builder, recursion_level: 0, entry_point, call_stack: CallStack::new() }
Self {
builder,
recursion_level: 0,
entry_point,
call_stack: CallStack::new(),
no_predicates_is_entry_point,
}
}

/// Start inlining the entry point function and all functions reachable from it.
Expand Down Expand Up @@ -351,11 +389,17 @@ impl<'function> PerFunctionContext<'function> {
for id in block.instructions() {
match &self.source_function.dfg[*id] {
Instruction::Call { func, arguments } => match self.get_function(*func) {
Some(function) => {
if ssa.functions[&function].runtime().is_entry_point() {
Some(func_id) => {
let function = &ssa.functions[&func_id];
// If we have not already finished the flattening pass, functions marked
// to not have predicates should be marked as entry points.
let no_predicates_is_entry_point =
self.context.no_predicates_is_entry_point
&& function.is_no_predicates();
if function.runtime().is_entry_point() || no_predicates_is_entry_point {
self.push_instruction(*id);
} else {
self.inline_function(ssa, *id, function, arguments);
self.inline_function(ssa, *id, func_id, arguments);
}
}
None => self.push_instruction(*id),
Expand Down
2 changes: 1 addition & 1 deletion compiler/noirc_frontend/src/ast/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ impl From<FunctionDefinition> for NoirFunction {
Some(FunctionAttribute::Oracle(_)) => FunctionKind::Oracle,
Some(FunctionAttribute::Recursive) => FunctionKind::Recursive,
Some(FunctionAttribute::Fold) => FunctionKind::Normal,
Some(FunctionAttribute::Inline(_)) => FunctionKind::Normal,
Some(FunctionAttribute::NoPredicates) => FunctionKind::Normal,
None => FunctionKind::Normal,
};

Expand Down
12 changes: 6 additions & 6 deletions compiler/noirc_frontend/src/hir/resolution/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -90,8 +90,8 @@ pub enum ResolverError {
MutableGlobal { span: Span },
#[error("Self-referential structs are not supported")]
SelfReferentialStruct { span: Span },
#[error("#[inline(tag)] attribute is only allowed on constrained functions")]
InlineAttributeOnUnconstrained { ident: Ident },
#[error("#[no_predicates] attribute is only allowed on constrained functions")]
NoPredicatesAttributeOnUnconstrained { ident: Ident },
#[error("#[fold] attribute is only allowed on constrained functions")]
FoldAttributeOnUnconstrained { ident: Ident },
}
Expand Down Expand Up @@ -362,16 +362,16 @@ impl<'a> From<&'a ResolverError> for Diagnostic {
*span,
)
},
ResolverError::InlineAttributeOnUnconstrained { ident } => {
ResolverError::NoPredicatesAttributeOnUnconstrained { ident } => {
let name = &ident.0.contents;

let mut diag = Diagnostic::simple_error(
format!("misplaced #[inline(tag)] attribute on unconstrained function {name}. Only allowed on constrained functions"),
"misplaced #[inline(tag)] attribute".to_string(),
format!("misplaced #[no_predicates] attribute on unconstrained function {name}. Only allowed on constrained functions"),
"misplaced #[no_predicates] attribute".to_string(),
ident.0.span(),
);

diag.add_note("The `#[inline(tag)]` attribute specifies to the compiler whether it should diverge from auto-inlining constrained functions".to_owned());
diag.add_note("The `#[no_predicates]` attribute specifies to the compiler whether it should diverge from auto-inlining constrained functions".to_owned());
diag
}
ResolverError::FoldAttributeOnUnconstrained { ident } => {
Expand Down
12 changes: 6 additions & 6 deletions compiler/noirc_frontend/src/hir/resolution/resolver.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1000,11 +1000,11 @@ impl<'a> Resolver<'a> {
let name_ident = HirIdent::non_trait_method(id, location);

let attributes = func.attributes().clone();
let has_inline_attribute = attributes.is_inline();
let has_no_predicates_attribute = attributes.is_no_predicates();
let should_fold = attributes.is_foldable();
if !self.inline_attribute_allowed(func) {
if has_inline_attribute {
self.push_err(ResolverError::InlineAttributeOnUnconstrained {
if has_no_predicates_attribute {
self.push_err(ResolverError::NoPredicatesAttributeOnUnconstrained {
ident: func.name_ident().clone(),
});
} else if should_fold {
Expand All @@ -1013,10 +1013,10 @@ impl<'a> Resolver<'a> {
});
}
}
// Both the #[fold] and #[inline(tag)] alter a function's inline type and code generation in similar ways.
// Both the #[fold] and #[no_predicates] alter a function's inline type and code generation in similar ways.
// In certain cases such as type checking (for which the following flag will be used) both attributes
// indicate we should code generate in the same way. Thus, we unify the attributes into one flag here.
let has_inline_or_fold_attribute = has_inline_attribute || should_fold;
let has_inline_attribute = has_no_predicates_attribute || should_fold;

let mut generics = vecmap(&self.generics, |(_, typevar, _)| typevar.clone());
let mut parameters = vec![];
Expand Down Expand Up @@ -1111,7 +1111,7 @@ impl<'a> Resolver<'a> {
has_body: !func.def.body.is_empty(),
trait_constraints: self.resolve_trait_constraints(&func.def.where_clause),
is_entry_point: self.is_entry_point_function(func),
has_inline_or_fold_attribute,
has_inline_attribute,
}
}

Expand Down
4 changes: 2 additions & 2 deletions compiler/noirc_frontend/src/hir/type_check/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ fn check_if_type_is_valid_for_program_input(
) {
let meta = type_checker.interner.function_meta(&func_id);
if (meta.is_entry_point && !param.1.is_valid_for_program_input())
|| (meta.has_inline_or_fold_attribute && !param.1.is_valid_non_inlined_function_input())
|| (meta.has_inline_attribute && !param.1.is_valid_non_inlined_function_input())
{
let span = param.0.span();
errors.push(TypeCheckError::InvalidTypeForEntryPoint { span });
Expand Down Expand Up @@ -546,7 +546,7 @@ pub mod test {
trait_constraints: Vec::new(),
direct_generics: Vec::new(),
is_entry_point: true,
has_inline_or_fold_attribute: false,
has_inline_attribute: false,
};
interner.push_fn_meta(func_meta, func_id);

Expand Down
5 changes: 3 additions & 2 deletions compiler/noirc_frontend/src/hir_def/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,8 +126,9 @@ pub struct FuncMeta {
pub is_entry_point: bool,

/// True if this function is marked with an attribute
/// that indicates it should not be inlined, such as `fold` or `inline(never)`
pub has_inline_or_fold_attribute: bool,
/// that indicates it should be inlined differently than the default (inline everything).
/// For example, such as `fold` (never inlined) or `no_predicates` (inlined after flattening)
pub has_inline_attribute: bool,
}

impl FuncMeta {
Expand Down
19 changes: 8 additions & 11 deletions compiler/noirc_frontend/src/lexer/token.rs
Original file line number Diff line number Diff line change
Expand Up @@ -582,8 +582,8 @@ impl Attributes {
self.function.as_ref().map_or(false, |func_attribute| func_attribute.is_foldable())
}

pub fn is_inline(&self) -> bool {
self.function.as_ref().map_or(false, |func_attribute| func_attribute.is_inline())
pub fn is_no_predicates(&self) -> bool {
self.function.as_ref().map_or(false, |func_attribute| func_attribute.is_no_predicates())
}
}

Expand Down Expand Up @@ -645,10 +645,7 @@ impl Attribute {
["test"] => Attribute::Function(FunctionAttribute::Test(TestScope::None)),
["recursive"] => Attribute::Function(FunctionAttribute::Recursive),
["fold"] => Attribute::Function(FunctionAttribute::Fold),
["inline", tag] => {
validate(tag)?;
Attribute::Function(FunctionAttribute::Inline(tag.to_string()))
}
["no_predicates"] => Attribute::Function(FunctionAttribute::NoPredicates),
["test", name] => {
validate(name)?;
let malformed_scope =
Expand Down Expand Up @@ -701,7 +698,7 @@ pub enum FunctionAttribute {
Test(TestScope),
Recursive,
Fold,
Inline(String),
NoPredicates,
}

impl FunctionAttribute {
Expand Down Expand Up @@ -738,8 +735,8 @@ impl FunctionAttribute {
/// Check whether we have an `inline` attribute
/// Although we also do not want to inline foldable functions,
/// we keep the two attributes distinct for clarity.
pub fn is_inline(&self) -> bool {
matches!(self, FunctionAttribute::Inline(_))
pub fn is_no_predicates(&self) -> bool {
matches!(self, FunctionAttribute::NoPredicates)
}
}

Expand All @@ -752,7 +749,7 @@ impl fmt::Display for FunctionAttribute {
FunctionAttribute::Oracle(ref k) => write!(f, "#[oracle({k})]"),
FunctionAttribute::Recursive => write!(f, "#[recursive]"),
FunctionAttribute::Fold => write!(f, "#[fold]"),
FunctionAttribute::Inline(ref k) => write!(f, "#[inline({k})]"),
FunctionAttribute::NoPredicates => write!(f, "#[no_predicates]"),
}
}
}
Expand Down Expand Up @@ -798,7 +795,7 @@ impl AsRef<str> for FunctionAttribute {
FunctionAttribute::Test { .. } => "",
FunctionAttribute::Recursive => "",
FunctionAttribute::Fold => "",
FunctionAttribute::Inline(string) => string,
FunctionAttribute::NoPredicates => "",
}
}
}
Expand Down
21 changes: 9 additions & 12 deletions compiler/noirc_frontend/src/monomorphization/ast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -213,24 +213,21 @@ pub enum InlineType {
Inline,
/// Functions marked as foldable will not be inlined and compiled separately into ACIR
Fold,
/// Similar to `Fold`, these functions will not be inlined and compile separately into ACIR.
/// They are different from `Fold` though as they are expected to be inlined into the program
/// Functions marked to have no predicates will not be inlined in the default inlining pass
/// and will be separately inlined after the flattening pass.
/// They are different from `Fold` as they are expected to be inlined into the program
/// entry point before being used in the backend.
Never,
/// This attribute is unsafe and can cause a function whose logic relies on predicates from
/// the flattening pass to fail.
NoPredicates,
}

impl From<&Attributes> for InlineType {
fn from(attributes: &Attributes) -> Self {
attributes.function.as_ref().map_or(InlineType::default(), |func_attribute| {
match func_attribute {
FunctionAttribute::Fold => InlineType::Fold,
FunctionAttribute::Inline(tag) => {
if tag == "never" {
InlineType::Never
} else {
InlineType::default()
}
}
FunctionAttribute::NoPredicates => InlineType::NoPredicates,
_ => InlineType::default(),
}
})
Expand All @@ -242,7 +239,7 @@ impl InlineType {
match self {
InlineType::Inline => false,
InlineType::Fold => true,
InlineType::Never => true,
InlineType::NoPredicates => false,
}
}
}
Expand All @@ -252,7 +249,7 @@ impl std::fmt::Display for InlineType {
match self {
InlineType::Inline => write!(f, "inline"),
InlineType::Fold => write!(f, "fold"),
InlineType::Never => write!(f, "inline(never)"),
InlineType::NoPredicates => write!(f, "no_predicates"),
}
}
}
Expand Down

0 comments on commit 0ce04d3

Please sign in to comment.