From 743bc32e94d851bb71af0dba0c012c81fbc93cc3 Mon Sep 17 00:00:00 2001
From: Aaron Councilman <aaronjc4@illinois.edu>
Date: Sun, 2 Mar 2025 14:01:09 -0600
Subject: [PATCH] Add feature testing
---
juno_scheduler/src/compile.rs | 23 +++++++++++++++++++++++
juno_scheduler/src/ir.rs | 3 +++
juno_scheduler/src/pm.rs | 18 ++++++++++++++++++
3 files changed, 44 insertions(+)
diff --git a/juno_scheduler/src/compile.rs b/juno_scheduler/src/compile.rs
index 0db32bda..8b68ed71 100644
--- a/juno_scheduler/src/compile.rs
+++ b/juno_scheduler/src/compile.rs
@@ -22,6 +22,7 @@ pub enum ScheduleCompilerError {
actual: usize,
loc: Location,
},
+ SemanticError(String, Location),
}
impl fmt::Display for ScheduleCompilerError {
@@ -46,6 +47,11 @@ impl fmt::Display for ScheduleCompilerError {
"({}, {}) -- ({}, {}): Expected {} arguments, found {}",
loc.0 .0, loc.0 .1, loc.1 .0, loc.1 .1, expected, actual
),
+ ScheduleCompilerError::SemanticError(msg, loc) => write!(
+ f,
+ "({}, {}) -- ({}, {}): {}",
+ loc.0 .0, loc.0 .1, loc.1 .0, loc.1 .1, msg,
+ ),
}
}
}
@@ -76,6 +82,8 @@ enum Appliable {
// DeleteUncalled requires special handling because it changes FunctionIDs, so it is not
// treated like a pass
DeleteUncalled,
+ // Test whether a feature is enabled
+ Feature,
Schedule(Schedule),
Device(Device),
}
@@ -85,6 +93,8 @@ impl Appliable {
fn is_valid_num_args(&self, num: usize) -> bool {
match self {
Appliable::Pass(pass) => pass.is_valid_num_args(num),
+ // Testing whether a feature is enabled takes the feature instead of a selection, so it
+ // has 0 arguments
// Delete uncalled, Schedules, and devices do not take arguments
_ => num == 0,
}
@@ -158,6 +168,8 @@ impl FromStr for Appliable {
"serialize" => Ok(Appliable::Pass(ir::Pass::Serialize)),
"write-predication" => Ok(Appliable::Pass(ir::Pass::WritePredication)),
+ "feature" => Ok(Appliable::Feature),
+
"print" => Ok(Appliable::Pass(ir::Pass::Print)),
"cpu" | "llvm" => Ok(Appliable::Device(Device::LLVM)),
@@ -409,6 +421,17 @@ fn compile_expr(
on: selection,
}))
}
+ Appliable::Feature => match selection {
+ ir::Selector::Selection(mut args) if args.len() == 1 => {
+ Ok(ExprResult::Expr(ir::ScheduleExp::Feature {
+ feature: Box::new(args.pop().unwrap()),
+ }))
+ }
+ _ => Err(ScheduleCompilerError::SemanticError(
+ "feature requires exactly one argument as its selection".to_string(),
+ lexer.line_col(span),
+ )),
+ },
Appliable::Schedule(sched) => Ok(ExprResult::Stmt(ir::ScheduleStmt::AddSchedule {
sched,
on: selection,
diff --git a/juno_scheduler/src/ir.rs b/juno_scheduler/src/ir.rs
index 98f8050f..bacb4142 100644
--- a/juno_scheduler/src/ir.rs
+++ b/juno_scheduler/src/ir.rs
@@ -121,6 +121,9 @@ pub enum ScheduleExp {
DeleteUncalled {
on: Selector,
},
+ Feature {
+ feature: Box<ScheduleExp>,
+ },
Record {
fields: Vec<(String, ScheduleExp)>,
},
diff --git a/juno_scheduler/src/pm.rs b/juno_scheduler/src/pm.rs
index db1dfc9a..bbcdb95f 100644
--- a/juno_scheduler/src/pm.rs
+++ b/juno_scheduler/src/pm.rs
@@ -16,6 +16,7 @@ use juno_utils::stringtab::StringTable;
use std::cell::RefCell;
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
+use std::env;
use std::fmt;
use std::fs::File;
use std::io::Write;
@@ -1460,6 +1461,23 @@ fn interp_expr(
changed,
))
}
+ ScheduleExp::Feature { feature } => {
+ let (feature, modified) = interp_expr(pm, &*feature, stringtab, env, functions)?;
+ let Value::String { val } = feature else {
+ return Err(SchedulerError::SemanticError(
+ "Feature expects a single string argument (instead of a selection)".to_string(),
+ ));
+ };
+ // To test for features, the scheduler needs to be invoked from a build script so that
+ // Cargo provides the enabled features via environment variables
+ let key = val.to_uppercase().replace("-", "_");
+ Ok((
+ Value::Boolean {
+ val: env::var(key).is_ok(),
+ },
+ modified,
+ ))
+ }
ScheduleExp::Record { fields } => {
let mut result = HashMap::new();
let mut changed = false;
--
GitLab