def main(functions_or_sigs, results_file, cores, debug):
"""
Attempt to build all the signatures in functions_or_sigs, or all the signatures
associated with all the functions in functions_or_sigs, or if functions_or_sigs
is empty every signature the stanc3 compiler exposes.
Results are written to a results json file. Individual signatures are classified
as either compatible, incompatible, or irrelevant.
Compatible signatures can be compiled with varmat types in every argument that
could possibly be a varmat (the matrix-like ones).
Incompatible signatures cannot all be built, and for irrelevant signatures it does
not make sense to try to build them (there are no matrix arguments, or the function
does not support reverse mode autodiff, etc).
Compilation is done in parallel using the number of specified cores.
:param functions_or_sigs: List of function names and/or signatures to benchmark
:param results_file: File to use as a results cache
:param cores: Number of cores to use for compiling
:param debug: If true, don't delete temporary files
"""
all_signatures = get_signatures()
functions, signatures = handle_function_list(functions_or_sigs)
requested_functions = set(functions)
compatible_signatures = set()
incompatible_signatures = set()
irrelevant_signatures = set()
# Read the arguments and figure out the exact list of signatures to test
signatures_to_check = set()
for signature in all_signatures:
sp = SignatureParser(signature)
if len(requested_functions
) > 0 and sp.function_name not in requested_functions:
continue
signatures_to_check.add(signature)
work_queue = Queue.Queue()
# For each signature, generate cpp code to test
for signature in signatures_to_check:
sp = SignatureParser(signature)
if sp.is_high_order():
work_queue.put((n, signature, None))
continue
cpp_code = ""
any_overload_uses_varmat = False
for m, overloads in enumerate(
itertools.product(("Prim", "Rev", "RevVarmat"),
repeat=sp.number_arguments())):
cg = CodeGenerator()
arg_list_base = cg.build_arguments(sp, overloads, size=1)
arg_list = []
for overload, arg in zip(overloads, arg_list_base):
if arg.is_reverse_mode() and arg.is_varmat_compatible(
) and overload.endswith("Varmat"):
any_overload_uses_varmat = True
arg = cg.to_var_value(arg)
arg_list.append(arg)
cg.function_call_assign("stan::math::" + sp.function_name,
*arg_list)
cpp_code += TEST_TEMPLATE.format(
test_name=sp.function_name + repr(m),
code=cg.cpp(),
)
if any_overload_uses_varmat:
work_queue.put((work_queue.qsize(), signature, cpp_code))
else:
print("{0} ... Irrelevant".format(signature.strip()))
irrelevant_signatures.add(signature)
output_lock = threading.Lock()
if not os.path.exists(WORKING_FOLDER):
os.mkdir(WORKING_FOLDER)
work_queue_original_length = work_queue.qsize()
# Test if each cpp file builds and update the output file
# This part is done in parallel
def worker():
while True:
try:
n, signature, cpp_code = work_queue.get(False)
except Queue.Empty:
return # If queue is empty, worker quits
# Use signature as filename prefix to make it easier to find
prefix = re.sub('[^0-9a-zA-Z]+', '_', signature.strip())
# Test the signature
successful = build_signature(prefix, cpp_code, debug)
# Acquire a lock to do I/O
with output_lock:
if successful:
result_string = "Success"
compatible_signatures.add(signature)
else:
result_string = "Fail"
incompatible_signatures.add(signature)
print("Results of test {0} / {1}, {2} ... ".format(
n, work_queue_original_length, signature.strip()) +
result_string)
work_queue.task_done()
for i in range(cores):
threading.Thread(target=worker).start()
work_queue.join()
with open(results_file, "w") as f:
json.dump(
{
"compatible_signatures": list(compatible_signatures),
"incompatible_signatures": list(incompatible_signatures),
"irrelevant_signatures": list(irrelevant_signatures)
},
f,
indent=4,
sort_keys=True)
class CodeGeneratorTest(unittest.TestCase):
def setUp(self):
self.add = SignatureParser("add(real, vector) => vector")
self.int_var = IntVariable("myint", 5)
self.real_var1 = RealVariable("Rev", "myreal1", 0.5)
self.real_var2 = RealVariable("Rev", "myreal2", 0.5)
self.matrix_var = MatrixVariable("Rev", "mymatrix", "matrix", 2, 0.5)
self.cg = CodeGenerator()
def test_prim_prim(self):
self.cg.build_arguments(self.add, ["Prim", "Prim"], 1)
self.assertEqual(
self.cg.cpp(), """double real0 = 0.4;
auto matrix1 = stan::test::make_arg<Eigen::Matrix<double, Eigen::Dynamic, 1>>(0.4, 1);"""
)
def test_prim_rev(self):
self.cg.build_arguments(self.add, ["Prim", "Rev"], 1)
self.assertEqual(
self.cg.cpp(), """double real0 = 0.4;
auto matrix1 = stan::test::make_arg<Eigen::Matrix<stan::math::var, Eigen::Dynamic, 1>>(0.4, 1);"""
)
def test_rev_rev(self):
self.cg.build_arguments(self.add, ["Rev", "Rev"], 1)
self.assertEqual(
self.cg.cpp(), """stan::math::var real0 = 0.4;
auto matrix1 = stan::test::make_arg<Eigen::Matrix<stan::math::var, Eigen::Dynamic, 1>>(0.4, 1);"""
)
def test_size(self):
self.cg.build_arguments(self.add, ["Rev", "Rev"], 2)
self.assertEqual(
self.cg.cpp(), """stan::math::var real0 = 0.4;
auto matrix1 = stan::test::make_arg<Eigen::Matrix<stan::math::var, Eigen::Dynamic, 1>>(0.4, 2);"""
)
def test_add(self):
self.cg.add(self.real_var1, self.real_var2)
self.assertEqual(
self.cg.cpp(),
"auto sum_of_sums0 = stan::math::eval(stan::math::add(myreal1,myreal2));"
)
def test_convert_to_expression(self):
self.cg.convert_to_expression(self.matrix_var)
self.assertEqual(
self.cg.cpp(), """int mymatrix_expr0_counter = 0;
stan::test::counterOp<stan::math::var> mymatrix_expr0_counter_op(&mymatrix_expr0_counter);
auto mymatrix_expr0 = mymatrix.block(0,0,2,2).unaryExpr(mymatrix_expr0_counter_op);"""
)
def test_expect_adj_eq(self):
self.cg.expect_adj_eq(self.real_var1, self.real_var2)
self.assertEqual(self.cg.cpp(),
"stan::test::expect_adj_eq(myreal1,myreal2);")
def test_expect_eq(self):
self.cg.expect_eq(self.real_var1, self.real_var2)
self.assertEqual(self.cg.cpp(), "EXPECT_STAN_EQ(myreal1,myreal2);")
def test_expect_leq_one(self):
self.cg.expect_leq_one(self.int_var)
self.assertEqual(self.cg.cpp(), """int int0 = 1;
EXPECT_LE(myint,int0);""")
def test_function_call_assign(self):
self.cg.function_call_assign("stan::math::add", self.real_var1,
self.real_var2)
self.assertEqual(
self.cg.cpp(),
"auto result0 = stan::math::eval(stan::math::add(myreal1,myreal2));"
)
def test_grad(self):
self.cg.grad(self.real_var1)
self.assertEqual(self.cg.cpp(), "stan::test::grad(myreal1);")
def test_recover_memory(self):
self.cg.recover_memory()
self.assertEqual(self.cg.cpp(), "stan::math::recover_memory();")
def test_recursive_sum(self):
self.cg.recursive_sum(self.real_var1)
self.assertEqual(
self.cg.cpp(),
"auto summed_result0 = stan::math::eval(stan::test::recursive_sum(myreal1));"
)
def test_to_var_value(self):
self.cg.to_var_value(self.matrix_var)
self.assertEqual(
self.cg.cpp(),
"auto mymatrix_varmat0 = stan::math::eval(stan::math::to_var_value(mymatrix));"
)
