Revision a0518f338c6c26b2464fa6d19d8707b567e4dded authored by Steven Johnson on 31 October 2022, 20:08:04 UTC, committed by Steven Johnson on 31 October 2022, 20:08:04 UTC
Associativity.cpp
#include "Associativity.h"
#include "CSE.h"
#include "ExprUsesVar.h"
#include "IREquality.h"
#include "IRMatch.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "IRPrinter.h"
#include "Simplify.h"
#include "Solve.h"
#include "Substitute.h"
#include "Util.h"
#include <algorithm>
#include <iterator>
namespace Halide {
namespace Internal {
using std::map;
using std::pair;
using std::set;
using std::string;
using std::vector;
namespace {
template<typename T>
vector<T> get_subvector(const vector<T> &v, const set<int> &indices) {
vector<T> sub;
for (const auto &index : indices) {
internal_assert(index < (int)v.size());
sub.push_back(v[index]);
}
return sub;
}
// Replace self-references to 'func' with arguments 'args' at
// 'value_index' in the Expr/Stmt with some Var
class ConvertSelfRef : public IRGraphMutator {
using IRGraphMutator::visit;
const string &func;
const vector<Expr> &args;
// If that function has multiple values, which value does this
// call node refer to?
const int value_index;
const vector<string> &op_x_names;
Expr visit(const Call *op) override {
if (!is_solvable) {
return op;
}
Expr expr = IRGraphMutator::visit(op);
op = expr.as<Call>();
internal_assert(op);
if ((op->call_type == Call::Halide) && (func == op->name)) {
internal_assert(args.size() == op->args.size())
<< "Self-reference should have the same number of args as the original\n";
for (size_t i = 0; i < op->args.size(); i++) {
if (!graph_equal(op->args[i], args[i])) {
debug(5) << "Self-reference of " << op->name
<< " with different args from the LHS. Operation is not associative\n";
is_solvable = false;
return expr;
}
}
// Substitute the call
internal_assert(op->value_index < (int)op_x_names.size());
debug(5) << " Substituting Call " << op->name << " at value index "
<< op->value_index << " with " << op_x_names[op->value_index] << "\n";
expr = Variable::make(op->type, op_x_names[op->value_index]);
if (op->value_index == value_index) {
x_part = op;
} else {
x_dependencies.insert(op->value_index);
}
}
return expr;
}
public:
ConvertSelfRef(const string &f, const vector<Expr> &args, int idx,
const vector<string> &x_names)
: func(f), args(args), value_index(idx), op_x_names(x_names) {
}
bool is_solvable = true;
set<int> x_dependencies; // Contains dependencies on self-reference at different tuple indices
Expr x_part; // Undefined if there is no self-reference at value_index
};
bool associative_op_pattern_match(const Expr &e,
const Expr &op,
const vector<string> &x_names,
const vector<string> &y_names,
const Scope<> &x_scope,
map<string, Expr> &match) {
internal_assert(e.type() == op.type())
<< "Expr has type " << e.type() << ", while pattern has type " << op.type() << "\n";
map<string, Expr> result;
if (expr_match(op, e, result)) {
debug(5) << "Found associative ops for " << e << " -> " << op
<< ", y_part: " << result["y0"] << "\n";
for (size_t i = 0; i < x_names.size(); ++i) {
const auto &iter = result.find("x" + std::to_string(i));
if (iter != result.end()) {
const Variable *xvar = iter->second.as<Variable>();
if ((xvar == nullptr) || (xvar->name != x_names[i])) {
debug(5) << "...Skipping match since the x_part is different than expected. "
<< "Expect: " << x_names[i] << "; get: " << iter->second << "\n";
return false;
}
}
}
for (size_t i = 0; i < y_names.size(); ++i) {
const auto &iter = result.find("y" + std::to_string(i));
if (iter != result.end()) {
// Make sure that y_part should not depend on x vars
if (expr_uses_vars(iter->second, x_scope)) {
debug(5) << "...Skipping match since the y_part depends on x vars\n";
return false;
}
}
}
for (size_t i = 0; i < x_names.size(); ++i) {
const auto &iter = result.find("k" + std::to_string(i));
if (iter != result.end()) {
// Make sure that k_part is constant
if (!is_const(iter->second)) {
debug(5) << "...Skipping match since the k_part is not constant\n";
return false;
}
}
}
// Make sure that the new matches are in agreement with any previous matches
for (const auto &iter : result) {
const auto &match_iter = match.find(iter.first);
if (match_iter == match.end()) {
debug(5) << "Adding result: " << iter.first << " -> " << iter.second << "\n";
match.emplace(iter.first, iter.second);
} else {
if (!equal(iter.first, match_iter->first) || !equal(iter.second, match_iter->second)) {
return false;
}
}
}
return true;
}
return false;
}
// Return true if we are able to find a match in the table (i.e. the op can be
// proven associative) and update 'assoc_op'.
bool find_match(const vector<AssociativePattern> &table, const vector<string> &op_x_names,
const vector<string> &op_y_names, const vector<Expr> &x_parts,
const vector<Expr> &exprs, AssociativeOp &assoc_op) {
internal_assert(op_x_names.size() == op_y_names.size());
internal_assert(op_x_names.size() == x_parts.size());
internal_assert(op_x_names.size() == exprs.size());
internal_assert(op_x_names.size() == assoc_op.size());
Scope<> x_scope;
for (const auto &x : op_x_names) {
x_scope.push(x);
}
for (const AssociativePattern &pattern : table) {
internal_assert(pattern.size() == op_x_names.size());
map<string, Expr> pattern_match;
bool matched = true;
// If any of element in 'pattern' does not match, try the next thing in
// the table.
for (size_t i = 0; i < pattern.size(); ++i) {
if (!associative_op_pattern_match(exprs[i], pattern.ops[i], op_x_names,
op_y_names, x_scope, pattern_match)) {
matched = false;
break;
}
}
if (!matched) {
continue;
}
vector<pair<Expr, Expr>> replacement; // find -> replacement
for (size_t index = 0; index < op_y_names.size(); ++index) {
const auto &y_iter = pattern_match.find("y" + std::to_string(index));
if (y_iter == pattern_match.end()) {
// Didn't find y{index} during pattern matching. Try next pattern.
matched = false;
break;
}
Expr y_part = y_iter->second;
debug(5) << "Pattern at index " << index << ":\n " << op_x_names[index]
<< " -> " << x_parts[index] << "\n " << op_y_names[index]
<< " -> " << y_part << "\n";
assoc_op.xs[index] = {op_x_names[index], x_parts[index]};
assoc_op.ys[index] = {op_y_names[index], y_part};
replacement.emplace_back(y_part, Variable::make(y_part.type(), op_y_names[index]));
}
if (!matched) {
continue;
}
for (size_t index = 0; index < exprs.size(); ++index) {
Expr e = exprs[index];
// Order of substitution matters, e.g. in the argmin case, _y_0 -> g(rx)[0]
// and _y_1 -> rx. If we substitute the 2nd element rx first, substitution
// of g(rx)[0] will fail.
for (const auto &iter : replacement) {
e = substitute(iter.first, iter.second, e);
}
assoc_op.pattern.ops[index] = e;
assoc_op.pattern.identities[index] = pattern.identities[index];
}
assoc_op.pattern.is_commutative = pattern.is_commutative;
return true;
}
return false;
}
// Return a pair of booleans indicating if an operator is associative.
// 'assoc_op' contains the the equivalent associative binary/unary operator
// for that operator. If the operator is non-associative, 'assoc_op' is not valid.
bool extract_associative_op(const vector<Expr> &exprs, const vector<string> &op_x_names,
const vector<string> &op_y_names, const vector<Expr> &x_parts,
AssociativeOp &assoc_op) {
if (exprs.size() == 1) {
Type t = exprs[0].type();
if (!x_parts[0].defined()) {
// Update with no self-recurrence is associative and the identity
// can be anything since it's going to be replaced anyway, but it's
// not commutative
assoc_op.pattern.ops[0] = Variable::make(t, op_y_names[0]);
assoc_op.pattern.identities[0] = make_const(t, 0);
assoc_op.pattern.is_commutative = false;
assoc_op.xs[0] = {"", Expr()};
assoc_op.ys[0] = {op_y_names[0], exprs[0]};
return true;
} else if (equal(exprs[0], Variable::make(t, op_x_names[0]))) {
// Self assignment, f(x) = f(x), is both associative
// and commutative. The identity can be anything since it's
// going to be replaced by itself.
debug(5) << "Self assignment: " << x_parts[0] << " = " << x_parts[0] << "\n";
assoc_op.pattern.ops[0] = Variable::make(t, op_x_names[0]);
assoc_op.pattern.identities[0] = make_const(t, 0);
assoc_op.pattern.is_commutative = true;
assoc_op.xs[0] = {op_x_names[0], x_parts[0]};
assoc_op.ys[0] = {"", Expr()};
return true;
}
}
return find_match(get_ops_table(exprs), op_x_names, op_y_names,
x_parts, exprs, assoc_op);
}
void add_transitive_dependencies(vector<set<int>> &dependencies) {
// TODO(psuriana): there might be a better way to find all the transitive dependencies
bool change = true;
while (change) {
change = false;
for (size_t i = 0; i < dependencies.size(); ++i) {
for (size_t j = 0; j < dependencies.size(); ++j) {
if (i == j) {
continue;
}
if (dependencies[i].count(j)) {
for (const auto &idx : dependencies[j]) {
if (dependencies[i].count(idx) == 0) {
dependencies[i].insert(idx);
change = true;
}
}
}
}
}
}
}
// Given dependencies of each tuple element, compute the set of subgraphs:
// all vertices that are reachable from a given vertex. If a subgraph is fully
// contained in another subgraph, remove it from the final output.
vector<set<int>> compute_subgraphs(vector<set<int>> dependencies) {
vector<set<int>> subgraphs(dependencies.size());
for (size_t i = 0; i < dependencies.size(); ++i) {
// Check if the current subgraph is a subset of another
const auto ¤t = dependencies[i];
if (current.empty()) {
continue;
}
bool should_remove = false;
for (size_t j = 0; j < dependencies.size(); ++j) {
const auto &other = dependencies[j];
if ((i == j) || (current.size() > other.size()) || (j < i && subgraphs[i].empty())) {
continue;
}
vector<int> diff;
// Compute the vertices in the current set that are not contained in the other
std::set_difference(current.begin(), current.end(), other.begin(), other.end(),
std::inserter(diff, diff.begin()));
if (diff.empty()) {
// 'current' is fully contained in 'other'
should_remove = true;
break;
}
}
if (!should_remove) {
subgraphs[i] = current;
}
}
return subgraphs;
}
} // anonymous namespace
AssociativeOp prove_associativity(const string &f, vector<Expr> args, vector<Expr> exprs) {
AssociativeOp assoc_op(exprs.size());
for (Expr &arg : args) {
// Undo the existing CSE pass done at function definition time
// to ensure things like += are in the expected form. Make no
// further transformations so that the LHS and RHS don't
// diverge.
arg = substitute_in_all_lets(arg);
}
vector<string> op_x_names(exprs.size()), op_y_names(exprs.size());
for (size_t idx = 0; idx < exprs.size(); ++idx) {
op_x_names[idx] = unique_name("_x_" + std::to_string(idx));
op_y_names[idx] = unique_name("_y_" + std::to_string(idx));
}
vector<set<int>> dependencies(exprs.size());
vector<Expr> x_parts(exprs.size());
bool all_independent = true;
// For a Tuple of exprs to be associative, each element of the Tuple
// has to be associative.
for (int idx = exprs.size() - 1; idx >= 0; --idx) {
// Undo the existing CSE pass done at function definition time.
exprs[idx] = substitute_in_all_lets(exprs[idx]);
// Replace any self-reference to Func 'f' with a Var
ConvertSelfRef csr(f, args, idx, op_x_names);
exprs[idx] = csr.mutate(exprs[idx]);
if (!csr.is_solvable) {
return AssociativeOp();
}
if (!csr.x_dependencies.empty()) {
all_independent = false;
}
x_parts[idx] = csr.x_part;
dependencies[idx] = csr.x_dependencies;
// Add dependency on itself (regardless whether it actually depends on
// its previous values) for the purpose of computing the subgraph
dependencies[idx].insert(idx);
exprs[idx] = common_subexpression_elimination(exprs[idx]);
exprs[idx] = simplify(exprs[idx]);
exprs[idx] = solve_expression(exprs[idx], op_x_names[idx]).result; // Move 'x' to the left as possible
exprs[idx] = substitute_in_all_lets(exprs[idx]);
}
internal_assert((exprs.size() != 1) || all_independent) << "1D tuple should be all independent\n";
vector<set<int>> subgraphs;
if (!all_independent) {
debug(5) << "There are cross-dependencies. Need to prove associativity in bulk.\n";
// Find all transitive dependencies and add them to the graph
add_transitive_dependencies(dependencies);
// Decompose the tuple into subgraphs and solve for each separately
subgraphs = compute_subgraphs(dependencies);
} else {
debug(5) << "All tuple elements are independent. Try proving associativity of "
<< "each element separately.\n";
// If all elements are independent, the subgraph is equal to the dependencies graph
subgraphs = dependencies;
}
internal_assert(subgraphs.size() == exprs.size());
for (size_t i = 0; i < subgraphs.size(); ++i) {
if (subgraphs[i].empty()) {
debug(5) << "Empty subgraph " << i << "\n";
continue;
}
if (subgraphs[i].size() > 2) {
// TODO(psuriana): Currently only support max of 2 tuple elements
debug(5) << "Subgraph " << i << " size is " << subgraphs[i].size() << " which is bigger than 2\n";
return AssociativeOp();
}
vector<Expr> sub_exprs = get_subvector(exprs, subgraphs[i]);
vector<string> sub_op_x_names = get_subvector(op_x_names, subgraphs[i]);
vector<string> sub_op_y_names = get_subvector(op_y_names, subgraphs[i]);
vector<Expr> sub_x_parts = get_subvector(x_parts, subgraphs[i]);
AssociativeOp sub_assoc_op(sub_exprs.size());
// TODO(psuriana): In general, if we fail to find a match for the
// set of initial subgraphs, we need to consider other possible
// grouping of those initial subgraphs. Since only the 'x' is
// apparent from the Halide update definition, the compute_subgraphs
// method over-partitions the graph (e.g. 2x2 matrix multiplication
// written as a four-dimensional reduction).
if (!extract_associative_op(sub_exprs, sub_op_x_names, sub_op_y_names,
sub_x_parts, sub_assoc_op)) {
debug(5) << "Cannot find matching associative ops\n";
return AssociativeOp();
}
debug(5) << "...Proving associativity of subgraph " << i << "\n";
const set<int> &indices = subgraphs[i];
for (auto iter = indices.begin(); iter != indices.end(); ++iter) {
int index = *iter;
int j = std::distance(indices.begin(), iter);
// If the ops/x/y have been extracted previously, we have to make sure
// they are consistent with the new extracted values.
if (assoc_op.pattern.ops[index].defined()) {
if (!equal(assoc_op.pattern.ops[index], sub_assoc_op.pattern.ops[j]) ||
!equal(assoc_op.pattern.identities[index], sub_assoc_op.pattern.identities[j])) {
debug(5) << "Conflicting associative ops/identities from different subgraphs\n";
return AssociativeOp();
}
}
if (assoc_op.xs[index].expr.defined()) {
if (assoc_op.xs[index] != sub_assoc_op.xs[j]) {
debug(5) << "Conflicting associative x-replacements from different subgraphs\n";
return AssociativeOp();
}
}
if (assoc_op.ys[index].expr.defined()) {
if (assoc_op.ys[index] != sub_assoc_op.ys[j]) {
debug(5) << "Conflicting associative y-replacements from different subgraphs\n";
return AssociativeOp();
}
}
assoc_op.pattern.ops[index] = sub_assoc_op.pattern.ops[j];
assoc_op.pattern.identities[index] = sub_assoc_op.pattern.identities[j];
assoc_op.pattern.is_commutative = sub_assoc_op.pattern.is_commutative;
assoc_op.xs[index] = sub_assoc_op.xs[j];
assoc_op.ys[index] = sub_assoc_op.ys[j];
}
}
assoc_op.is_associative = true;
debug(5) << "Found associative ops:\n"
<< assoc_op << "\n";
return assoc_op;
}
namespace {
std::string print_args(const string &f, const vector<Expr> &args, const vector<Expr> &exprs) {
std::ostringstream stream;
stream << f << "(";
for (size_t i = 0; i < args.size(); ++i) {
stream << args[i];
if (i != args.size() - 1) {
stream << ", ";
}
}
stream << ") = ";
if (exprs.size() == 1) {
stream << exprs[0];
} else if (exprs.size() > 1) {
stream << "Tuple(";
for (size_t i = 0; i < exprs.size(); ++i) {
stream << exprs[i];
if (i != exprs.size() - 1) {
stream << ", ";
}
}
stream << ")";
}
return stream.str();
}
void check_associativity(const string &f, const vector<Expr> &args, const vector<Expr> &exprs,
const AssociativeOp &assoc_op) {
auto result = prove_associativity(f, args, exprs);
internal_assert(result.associative() == assoc_op.associative())
<< "Checking associativity: " << print_args(f, args, exprs) << "\n"
<< " Expect is associative: " << assoc_op.associative() << "\n"
<< " instead of " << result.associative() << "\n";
if (assoc_op.associative()) {
map<string, Expr> replacement;
for (size_t i = 0; i < assoc_op.size(); ++i) {
internal_assert(equal(result.pattern.identities[i], assoc_op.pattern.identities[i]))
<< "Checking associativity: " << print_args(f, args, exprs) << "\n"
<< " Index: " << i << "\n"
<< " Expect identity: " << assoc_op.pattern.identities[i] << "\n"
<< " instead of " << result.pattern.identities[i] << "\n";
internal_assert(equal(result.xs[i].expr, assoc_op.xs[i].expr))
<< "Checking associativity: " << print_args(f, args, exprs) << "\n"
<< " Index: " << i << "\n"
<< " Expect x: " << assoc_op.xs[i].expr << "\n"
<< " instead of " << result.xs[i].expr << "\n";
internal_assert(equal(result.ys[i].expr, assoc_op.ys[i].expr))
<< "Checking associativity: " << print_args(f, args, exprs) << "\n"
<< " Index: " << i << "\n"
<< " Expect y: " << assoc_op.ys[i].expr << "\n"
<< " instead of " << result.ys[i].expr << "\n";
if (result.xs[i].expr.defined()) {
replacement.emplace(assoc_op.xs[i].var, Variable::make(result.xs[i].expr.type(), result.xs[i].var));
}
if (result.ys[i].expr.defined()) {
replacement.emplace(assoc_op.ys[i].var, Variable::make(result.ys[i].expr.type(), result.ys[i].var));
}
}
for (size_t i = 0; i < assoc_op.size(); ++i) {
Expr expected_op = substitute(replacement, assoc_op.pattern.ops[i]);
internal_assert(equal(result.pattern.ops[i], expected_op))
<< "Checking associativity: " << print_args(f, args, exprs) << "\n"
<< " Index: " << i << "\n"
<< " Expect bin op: " << expected_op << "\n"
<< " instead of " << result.pattern.ops[i] << "\n";
debug(5) << "\nExpected op: " << expected_op << "\n";
debug(5) << "Operator: " << result.pattern.ops[i] << "\n";
debug(5) << " identity: " << result.pattern.identities[i] << "\n";
debug(5) << " x: " << result.xs[i].var << " -> " << result.xs[i].expr << "\n";
debug(5) << " y: " << result.ys[i].var << " -> " << result.ys[i].expr << "\n";
}
}
}
} // anonymous namespace
void associativity_test() {
typedef AssociativeOp::Replacement Replacement;
{
// Tests for saturating addition
Type t = UInt(8);
Expr x = Variable::make(t, "x");
Expr y = Variable::make(t, "y");
Expr x_idx = Variable::make(Int(32), "x_idx");
Expr f_call_0 = Call::make(t, "f", {x_idx}, Call::CallType::Halide, FunctionPtr(), 0);
// f(x) = uint8(uint16(x + y), 255)
check_associativity("f", {x_idx}, {Cast::make(UInt(8), min(Cast::make(UInt(16), y + f_call_0), make_const(t, 255)))},
AssociativeOp(
AssociativePattern(Cast::make(UInt(8), min(Cast::make(UInt(16), x + y), make_const(t, 255))), make_const(t, 0), true),
{Replacement("x", f_call_0)},
{Replacement("y", y)},
true));
// f(x) = uint8(uint16(x + y), uint16(255))
check_associativity("f", {x_idx}, {Cast::make(UInt(8), min(Cast::make(UInt(16), y + f_call_0), Cast::make(UInt(16), make_const(t, 255))))},
AssociativeOp(
AssociativePattern(Cast::make(UInt(8), min(Cast::make(UInt(16), x + y), make_const(t, 255))), make_const(t, 0), true),
{Replacement("x", f_call_0)},
{Replacement("y", y)},
true));
// f(x) = select(x > 255 - y, 255, y)
check_associativity("f", {x_idx}, {select(f_call_0 > make_const(t, 255) - y, make_const(t, 255), y)},
AssociativeOp(
AssociativePattern(select(x > make_const(t, 255) - y, make_const(t, 255), y), make_const(t, 0), true),
{Replacement("x", f_call_0)},
{Replacement("y", y)},
true));
// f(x) = select(x >= -y, 255, y)
check_associativity("f", {x_idx}, {select(f_call_0 >= -y, make_const(t, 255), y)},
AssociativeOp(
AssociativePattern(select(x < -y, y, make_const(t, 255)), make_const(t, 0), true),
{Replacement("x", f_call_0)},
{Replacement("y", y)},
true));
}
{
// Tests for logical And/Or
Type t = UInt(1);
Expr x = Variable::make(t, "x");
Expr y = Variable::make(t, "y");
Expr x_idx = Variable::make(Int(32), "x_idx");
Expr f_call_0 = Call::make(t, "f", {x_idx}, Call::CallType::Halide, FunctionPtr(), 0);
// f(x) = y && f(x)
check_associativity("f", {x_idx}, {And::make(y, f_call_0)},
AssociativeOp(
AssociativePattern(And::make(x, y), const_true(), true),
{Replacement("x", f_call_0)},
{Replacement("y", y)},
true));
// f(x) = y || f(x)
check_associativity("f", {x_idx}, {Or::make(y, f_call_0)},
AssociativeOp(
AssociativePattern(Or::make(x, y), const_false(), true),
{Replacement("x", f_call_0)},
{Replacement("y", y)},
true));
}
{
// Tests for 1D reduction
Type t = Int(32);
Expr x = Variable::make(t, "x");
Expr y = Variable::make(t, "y");
Expr z = Variable::make(t, "z");
Expr rx = Variable::make(t, "rx");
Expr f_call_0 = Call::make(t, "f", {x}, Call::CallType::Halide, FunctionPtr(), 0);
Expr g_call_0 = Call::make(t, "g", {rx}, Call::CallType::Halide, FunctionPtr(), 0);
// f(x) = f(x)
check_associativity("f", {x}, {f_call_0},
AssociativeOp(
AssociativePattern(x, make_const(t, 0), true),
{Replacement("x", f_call_0)},
{Replacement("", Expr())},
true));
// f(x) = min(f(x), y + int16(z))
check_associativity("f", {x}, {min(f_call_0, y + Cast::make(Int(16), z))},
AssociativeOp(
AssociativePattern(min(x, y), t.max(), true),
{Replacement("x", f_call_0)},
{Replacement("y", y + Cast::make(Int(16), z))},
true));
// f(x) = f(x) + g(rx) + y + z
check_associativity("f", {x}, {y + z + f_call_0},
AssociativeOp(
AssociativePattern(x + y, make_const(t, 0), true),
{Replacement("x", f_call_0)},
{Replacement("y", y + z)},
true));
// f(x) = max(y, f(x))
check_associativity("f", {x}, {max(y, f_call_0)},
AssociativeOp(
AssociativePattern(max(x, y), t.min(), true),
{Replacement("x", f_call_0)},
{Replacement("y", y)},
true));
// f(x) = max(f(x) + g(rx), g(rx)) -> not associative
check_associativity("f", {x}, {max(f_call_0 + g_call_0, g_call_0)}, AssociativeOp());
// f(x) = max(f(x) + g(rx), f(x) - 3) -> f(x) + max(g(rx) - 3)
check_associativity("f", {x}, {max(f_call_0 + g_call_0, f_call_0 - 3)},
AssociativeOp(
AssociativePattern(x + y, 0, true),
{Replacement("x", f_call_0)},
{Replacement("y", max(g_call_0, -3))},
true));
// f(x) = min(4, g(rx)) -> trivially associative
check_associativity("f", {x}, {min(4, g_call_0)},
AssociativeOp(
AssociativePattern(y, make_const(t, 0), true),
{Replacement("", Expr())},
{Replacement("y", min(g_call_0, 4))},
true));
// f(x) = max(max(min(f(x), g(rx) + 2), f(x)), g(rx) + 2) -> can be simplified into max(f(x), g(rx) + 2)
check_associativity("f", {x}, {max(max(min(f_call_0, g_call_0 + 2), f_call_0), g_call_0 + 2)},
AssociativeOp(
AssociativePattern(max(x, y), t.min(), true),
{Replacement("x", f_call_0)},
{Replacement("y", g_call_0 + 2)},
true));
// f(x) = max(x0, f(x)) -> x0 may conflict with the wildcard associative op pattern
Expr x0 = Variable::make(t, "x0");
check_associativity("f", {x}, {max(x0, f_call_0)},
AssociativeOp(
AssociativePattern(max(x, y), t.min(), true),
{Replacement("x", f_call_0)},
{Replacement("y", x0)},
true));
}
{
// Tests for multi-dimensional reduction (with mixed types)
Type t = Int(32);
Expr x = Variable::make(t, "x");
Expr y = Variable::make(t, "y");
Expr z = Variable::make(t, "z");
Expr rx = Variable::make(t, "rx");
vector<Type> ts = {Int(32), Int(32), Float(32)};
vector<Expr> xs(3), ys(3), zs(3);
for (size_t i = 0; i < xs.size(); ++i) {
xs[i] = Variable::make(ts[i], "x" + std::to_string(i));
ys[i] = Variable::make(ts[i], "y" + std::to_string(i));
zs[i] = Variable::make(ts[i], "z" + std::to_string(i));
}
Expr f_call_0 = Call::make(ts[0], "f", {x}, Call::CallType::Halide, FunctionPtr(), 0);
Expr f_call_1 = Call::make(ts[1], "f", {x}, Call::CallType::Halide, FunctionPtr(), 1);
Expr f_call_2 = Call::make(ts[2], "f", {x}, Call::CallType::Halide, FunctionPtr(), 2);
Expr g_call_0 = Call::make(ts[0], "g", {rx}, Call::CallType::Halide, FunctionPtr(), 0);
Expr g_call_1 = Call::make(ts[1], "g", {rx}, Call::CallType::Halide, FunctionPtr(), 1);
// f(x) = Tuple(f(x)[0], 3, f(x)[2] + z)
check_associativity("f", {x}, {f_call_0, make_const(ts[1], 3), f_call_2 + cast(ts[2], z)},
AssociativeOp(
AssociativePattern({xs[0], ys[1], xs[2] + ys[2]},
{make_const(ts[0], 0), make_const(ts[1], 0), make_const(ts[2], 0)},
true),
{Replacement("x0", f_call_0), Replacement("", Expr()), Replacement("x2", f_call_2)},
{Replacement("", Expr()), Replacement("y1", make_const(ts[1], 3)), Replacement("y2", cast(ts[2], z))},
true));
// f(x) = Tuple(2, 3, f(x)[2] + z)
check_associativity("f", {x}, {make_const(ts[0], 2), make_const(ts[1], 3), f_call_2 + cast(ts[2], z)},
AssociativeOp(
AssociativePattern({ys[0], ys[1], xs[2] + ys[2]},
{make_const(ts[0], 0), make_const(ts[1], 0), make_const(ts[2], 0)},
true),
{Replacement("", Expr()), Replacement("", Expr()), Replacement("x2", f_call_2)},
{Replacement("y0", make_const(ts[0], 2)), Replacement("y1", make_const(ts[1], 3)), Replacement("y2", cast(ts[2], z))},
true));
// f(x) = Tuple(min(f(x)[0], g(rx)), f(x)[1]*g(x)*2, f(x)[2] + z)
check_associativity("f", {x}, {min(f_call_0, g_call_0), f_call_1 * g_call_0 * 2, f_call_2 + cast(ts[2], z)},
AssociativeOp(
AssociativePattern(
{min(xs[0], ys[0]), xs[1] * ys[1], xs[2] + ys[2]},
{ts[0].max(), make_const(ts[1], 1), make_const(ts[2], 0)},
true),
{Replacement("x0", f_call_0), Replacement("x1", f_call_1), Replacement("x2", f_call_2)},
{Replacement("y0", g_call_0), Replacement("y1", g_call_0 * 2), Replacement("y2", cast(ts[2], z))},
true));
// Complex multiplication: f(x) = Tuple(f(x)[0]*g(r.x)[0] - f(x)[1]*g(r.x)[1], f(x)[0]*g(r.x)[1] + f(x)[1]*g(r.x)[0])
check_associativity("f", {x}, {f_call_0 * g_call_0 - f_call_1 * g_call_1, f_call_0 * g_call_1 + f_call_1 * g_call_0},
AssociativeOp(
AssociativePattern(
{xs[0] * ys[0] - ys[1] * xs[1], xs[1] * ys[0] + ys[1] * xs[0]},
{make_const(ts[0], 1), make_const(ts[1], 0)},
true),
{Replacement("x0", f_call_0), Replacement("x1", f_call_1)},
{Replacement("y0", g_call_0), Replacement("y1", g_call_1)},
true));
// 1D argmin: f(x) = Tuple(min(f(x)[0], g(r.x)[0]), select(f(x)[0] < g(r.x)[0], f(x)[1], g(r.x)[1])
check_associativity("f", {x}, {min(f_call_0, g_call_0), select(f_call_0 < g_call_0, f_call_1, g_call_1)},
AssociativeOp(
AssociativePattern(
{min(xs[0], ys[0]), select(xs[0] < ys[0], xs[1], ys[1])},
{ts[0].max(), make_const(ts[1], 0)},
true),
{Replacement("x0", f_call_0), Replacement("x1", f_call_1)},
{Replacement("y0", g_call_0), Replacement("y1", g_call_1)},
true));
}
{
Type t = Int(32);
Expr x = Variable::make(t, "x");
Expr y = Variable::make(t, "y");
Expr rx = Variable::make(t, "rx");
Expr ry = Variable::make(t, "ry");
vector<Type> ts = {UInt(8), Int(32), Int(16), Float(32)};
vector<Expr> xs(4), ys(4), zs(4);
for (size_t i = 0; i < xs.size(); ++i) {
xs[i] = Variable::make(ts[i], "x" + std::to_string(i));
ys[i] = Variable::make(ts[i], "y" + std::to_string(i));
zs[i] = Variable::make(ts[i], "z" + std::to_string(i));
}
Expr f_xy_call_0 = Call::make(ts[0], "f", {x, y}, Call::CallType::Halide, FunctionPtr(), 0);
Expr f_xy_call_1 = Call::make(ts[1], "f", {x, y}, Call::CallType::Halide, FunctionPtr(), 1);
Expr f_xy_call_2 = Call::make(ts[2], "f", {x, y}, Call::CallType::Halide, FunctionPtr(), 2);
Expr f_xy_call_3 = Call::make(ts[3], "f", {x, y}, Call::CallType::Halide, FunctionPtr(), 3);
Expr g_xy_call_0 = Call::make(ts[0], "g", {rx, ry}, Call::CallType::Halide, FunctionPtr(), 0);
// 2D argmin + trivial update (with mixed types):
// f(x, y) = Tuple(min(f(x, y)[0], g(r.x, r.y)[0]),
// r.x + r.y,
// select(f(x, y)[0] < g(r.x, r.y)[0], f(x)[2], r.x),
// select(f(x, y)[0] < g(r.x, r.y)[0], f(x)[3], r.y))
check_associativity("f", {x, y},
{min(f_xy_call_0, g_xy_call_0),
rx + ry,
select(f_xy_call_0 < g_xy_call_0, f_xy_call_2, cast(Int(16), rx)),
select(f_xy_call_0 < g_xy_call_0, f_xy_call_3, cast(Float(32), ry))},
AssociativeOp(
AssociativePattern(
{min(xs[0], ys[0]), ys[1], select(xs[0] < ys[0], xs[2], ys[2]), select(xs[0] < ys[0], xs[3], ys[3])},
{ts[0].max(), make_const(ts[1], 0), make_const(ts[2], 0), make_const(ts[3], 0)},
true),
{Replacement("x0", f_xy_call_0), Replacement("", Expr()),
Replacement("x2", f_xy_call_2), Replacement("x3", f_xy_call_3)},
{Replacement("y0", g_xy_call_0), Replacement("y1", rx + ry),
Replacement("y2", cast(Int(16), rx)), Replacement("y3", cast(Float(32), ry))},
true));
}
std::cout << "Associativity test passed" << std::endl;
}
} // namespace Internal
} // namespace Halide
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...