Skip to content

Commit

Permalink
Merge branch 'array-part' into main
Browse files Browse the repository at this point in the history
  • Loading branch information
@kumasento
kumasento committed Oct 9, 2021
2 parents 95bc22a + 3f1d895 commit c2a2469
Show file tree
Hide file tree
Showing 2 changed files with 0 additions and 370 deletions.
346 changes: 0 additions & 346 deletions lib/mlir/Transforms/ArrayPartition.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,349 +51,6 @@ struct ArrayPartitionPipelineOptions
};
} // namespace

/// -------------------------- Dependence analysis ---------------------------

static FlatAffineValueConstraints getOpIndexSet(Operation *op) {
FlatAffineValueConstraints cst;
SmallVector<Operation *, 4> ops;
getEnclosingAffineForAndIfOps(*op, &ops);
getIndexSet(ops, &cst);
return cst;
}

/// Get the access domain from all the provided operations.
static FlatAffineValueConstraints getDomain(ArrayRef<Operation *> ops) {
FlatAffineValueConstraints cst;
for (Operation *op : ops) {
MemRefAccess access(op);

AffineValueMap accessMap;
access.getAccessMap(&accessMap);

SmallSetVector<Value, 4> ivs; // used to access memref.
for (Value operand : accessMap.getOperands())
ivs.insert(operand);

FlatAffineValueConstraints domain = getOpIndexSet(op);

// project out those IDs that are not in the accessMap.
SmallVector<Value> values;
domain.getValues(0, domain.getNumDimIds(), &values);

SmallVector<Value> toProject;
for (Value value : values)
if (!ivs.count(value))
toProject.push_back(value);

for (Value id : toProject) {
unsigned pos;
domain.findId(id, &pos);
domain.projectOut(pos);
}

cst.mergeAndAlignIdsWithOther(0, &domain);
cst.append(domain);
}

cst.removeTrivialRedundancy();
return cst;
}

struct MemRefAccessInfo {
CallOp caller;
FlatAffineValueConstraints readCst;
FlatAffineValueConstraints writeCst;

MemRefAccessInfo(CallOp caller, FlatAffineValueConstraints readCst,
FlatAffineValueConstraints writeCst)
: caller(caller), readCst(readCst), writeCst(writeCst) {}

bool isReadOnly() const {
return writeCst.getNumConstraints() == 0 && readCst.getNumConstraints() > 0;
}

bool isWriteOnly() const {
return writeCst.getNumConstraints() > 0 && readCst.getNumConstraints() == 0;
}

bool isEmpty() const {
return writeCst.getNumConstraints() == 0 &&
readCst.getNumConstraints() == 0;
}

bool isReadWrite() const {
return writeCst.getNumConstraints() > 0 && readCst.getNumConstraints() > 0;
}

unsigned getNumDims() const {
return isReadOnly() ? readCst.getNumDimIds() : writeCst.getNumDimIds();
}
};

static FlatAffineValueConstraints
getArrayPartition(const FlatAffineValueConstraints &cst, ArrayRef<int64_t> inds,
FuncOp callee, ArrayRef<Value> ivs,
SmallDenseMap<Value, unsigned> &ivIds) {
FlatAffineValueConstraints cur{cst};
for (auto ind : enumerate(inds)) {
Value id = callee.getArgument(ivIds[ivs[ind.index()]]);
unsigned pos;
cur.findId(id, &pos);
cur.setAndEliminate(pos, {ind.value()});
}
cur.removeRedundantConstraints();

return cur;
}

/// TODO: this function tries to find whether the FORMs of two constraints are
/// the same, not exactly the domain they are covering.
static bool isSame(const FlatAffineValueConstraints &cst1,
const FlatAffineValueConstraints &cst2) {
if (cst1.getNumCols() != cst2.getNumCols())
return false;
if (cst1.getNumConstraints() != cst2.getNumConstraints())
return false;
if (cst1.getNumEqualities() != cst2.getNumEqualities())
return false;
for (unsigned i = 0; i < (unsigned)cst1.getNumEqualities(); ++i)
for (unsigned j = 0; j < (unsigned)cst2.getNumCols(); ++j)
if (cst1.atEq(i, j) != cst2.atEq(i, j))
return false;
for (unsigned i = 0; i < (unsigned)cst1.getNumInequalities(); ++i)
for (unsigned j = 0; j < (unsigned)cst2.getNumCols(); ++j)
if (cst1.atIneq(i, j) != cst2.atIneq(i, j))
return false;

return true;
}

/// Map from memrefs to their accesses from all PE callers.
static auto getMemRefAccessInfo(ArrayRef<Operation *> callers, ModuleOp m) {
SmallDenseMap<Value, std::vector<MemRefAccessInfo>> accesses;

// Initialise the map from memref to PE caller accesses.
for (Operation *op : callers) {
mlir::CallOp caller = cast<mlir::CallOp>(op);
FuncOp callee = cast<FuncOp>(m.lookupSymbol(caller.getCallee()));

SmallVector<std::pair<Value, unsigned>> memrefs;
for (auto arg : enumerate(caller.getArgOperands()))
if (arg.value().getType().isa<MemRefType>())
memrefs.push_back({arg.value(), arg.index()});

// Iterate every memref being accessed by the current caller.
for (auto memref : memrefs) {
Value arg = callee.getArgument(memref.second);

// Get all the read/write accesses.
SmallVector<Operation *> loadOps, storeOps;
copy_if(arg.getUsers(), std::back_inserter(loadOps),
[](Operation *op) { return isa<mlir::AffineLoadOp>(op); });
copy_if(arg.getUsers(), std::back_inserter(storeOps),
[](Operation *op) { return isa<mlir::AffineStoreOp>(op); });

// Union all the read constraints from all the load operations.
FlatAffineValueConstraints readCst = getDomain(loadOps);
FlatAffineValueConstraints writeCst = getDomain(storeOps);
accesses[memref.first].push_back({caller, readCst, writeCst});
}
}

return accesses;
}

static bool
getLoopIVsAndBounds(Operation *op, SmallVectorImpl<Value> &ivs,
SmallVectorImpl<std::pair<int64_t, int64_t>> &bounds) {
SmallVector<AffineForOp> forOps;
getLoopIVs(*op, &forOps);

for (AffineForOp forOp : forOps)
if (forOp.hasConstantUpperBound() && forOp.hasConstantLowerBound()) {
bounds.push_back(
{forOp.getConstantLowerBound(), forOp.getConstantUpperBound()});
ivs.push_back(forOp.getInductionVar());
}

// If there is no bound, or not all bounds are constant, return false.
if (bounds.empty() || bounds.size() != forOps.size())
return false;

return true;
}

static auto getArgIndexMap(CallOp caller) {
SmallDenseMap<Value, unsigned> argId;
for (auto arg : enumerate(caller.getArgOperands()))
argId[arg.value()] = arg.index();
return argId;
}

// Get the combinations of indices within the provided bounds.
static auto getIndexCombinations(ArrayRef<std::pair<int64_t, int64_t>> bounds) {
std::vector<std::vector<int64_t>> indices{{}};
for (auto bound : bounds) {
int64_t lb, ub;
std::tie(lb, ub) = bound;
std::vector<std::vector<int64_t>> newIndices;

for (auto index : indices)
for (int64_t value = lb; value < ub; ++value) {
std::vector<int64_t> newIndex{index};
newIndex.push_back(value);
newIndices.push_back(newIndex);
}

std::swap(indices, newIndices);
}
return indices;
}

using Partition = std::vector<std::pair<int64_t, int64_t>>;

static bool checkAccessOverlap(MemRefAccessInfo &info, ModuleOp m,
std::set<Partition> &partitions) {
// Get the IVs and constant bounds for the loops surrounding the caller.
SmallVector<Value> ivs;
SmallVector<std::pair<int64_t, int64_t>> bounds;
if (!getLoopIVsAndBounds(info.caller.getOperation(), ivs, bounds))
return true;

// Loop induction variable to argument ID in the caller.
auto argId = getArgIndexMap(info.caller);

// Get every partition.
// TODO: Can we make this part less memory intensive?
std::vector<std::vector<int64_t>> indices = getIndexCombinations(bounds);

FlatAffineValueConstraints cst =
info.isWriteOnly() ? info.writeCst : info.readCst;
FuncOp callee = cast<FuncOp>(m.lookupSymbol(info.caller.getCallee()));

// Iterate every possible partitions and check if they would overlap.
std::vector<FlatAffineValueConstraints> parts; // temporary results;
for (auto inds1 : indices) {
FlatAffineValueConstraints cur1 =
getArrayPartition(cst, inds1, callee, ivs, argId);
if (cur1.isEmpty())
continue;

for (auto inds2 : indices) {
if (inds1 == inds2)
continue;

FlatAffineValueConstraints cur2 =
getArrayPartition(cst, inds2, callee, ivs, argId);

// If cur1 and cur2 are exactly the same, then it shouldn't be
// considered as overlapping.
if (isSame(cur1, cur2))
continue;

FlatAffineValueConstraints tmp{cur1};
tmp.append(cur2);
if (!tmp.isEmpty())
return true;
}
parts.push_back(cur1);
}

// De-duplicate
// TODO: make it more efficient.
for (unsigned i = 0; i < parts.size(); ++i) {
std::vector<std::pair<int64_t, int64_t>> partition;

for (unsigned pos = 0; pos < parts[i].getNumDimIds(); ++pos) {
auto lb = parts[i].getConstantBound(
FlatAffineValueConstraints::BoundType::LB, pos);
auto ub = parts[i].getConstantBound(
FlatAffineValueConstraints::BoundType::UB, pos);
if (lb.hasValue() && ub.hasValue())
partition.push_back({lb.getValue(), ub.getValue()});
}

if (partition.size() != parts[i].getNumDimIds()) {
llvm::errs()
<< "The number of constant partitions are less than the dim Ids\n";
return true;
}

partitions.insert(partition);
}

return false;
}

static void arrayPartition(FuncOp f, ModuleOp m, OpBuilder &b) {
// Get all the PE callers.
SmallVector<Operation *> callers;
f.walk([&](CallOp caller) {
if (caller->hasAttr("scop.pe"))
callers.push_back(caller);
});
if (callers.empty())
return;

// Get MemRef accesses.
auto accesses = getMemRefAccessInfo(callers, m);

for (auto &access : accesses) {
Value memref;
std::vector<MemRefAccessInfo> accessInfos;
std::tie(memref, accessInfos) = access;

// For now, we only look at those memrefs that are only accessed by one PE.
if (accessInfos.size() > 1)
continue;

// Check if the only access is read or write. We cannot deal with read-write
// access at the moment.
// TODO: deal with read-write.
MemRefAccessInfo info = accessInfos.front();
if (info.isEmpty() || info.isReadWrite())
continue;

// Check if there any overlap for the memref being accessed. If not, then we
// can find valid array partitions from it.
std::set<Partition> partitions;
if (checkAccessOverlap(info, m, partitions))
continue;

llvm::errs() << "Partitions: \n";
for (auto it : partitions) {
for (auto bound : enumerate(it)) {
llvm::errs() << bound.index() << " -> " << bound.value().first << ' '
<< bound.value().second << "\n";
}
llvm::errs() << "-----------------\n";
}

// Next, decide how to update the type of the memref.
// First determine what the partition size is for each dimension, then
// decide new memref type.
SmallVector<int64_t> dimSizes(info.getNumDims(), 0);
for (unsigned i = 0; i < info.getNumDims(); ++i)
for (auto p : partitions)
dimSizes[i] = std::max((int64_t)dimSizes[i],
(int64_t)(p[i].second - p[i].first + 1));
}
}

namespace {
struct ArrayPartitionPass
: public mlir::PassWrapper<ArrayPartitionPass, OperationPass<ModuleOp>> {

void runOnOperation() override {
ModuleOp m = getOperation();
OpBuilder b(m.getContext());

m.walk([&](FuncOp f) { arrayPartition(f, m, b); });
}
};

} // namespace

/// ---------------- Simple array partition ---------------------------
/// Find partitioning opportunities by affine expressions.

Expand Down Expand Up @@ -1403,9 +1060,6 @@ struct SimpleArrayPartitionPass
} // namespace

void phism::registerArrayPartitionPasses() {
// PassRegistration<ArrayPartitionPass>("array-partition", "Partition
// arrays");

PassPipelineRegistration<ArrayPartitionPipelineOptions>(
"simple-array-partition", "Partition arrays",
[&](OpPassManager &pm, const ArrayPartitionPipelineOptions &options) {
Expand Down
24 changes: 0 additions & 24 deletions scripts/fix-cosim.py
View file

This file was deleted.

0 comments on commit c2a2469

Please sign in to comment.