[mlir] [VectorOps] Correctly account for rank-0 affine-map result in vector.contract

Summary:
Now that, thanks to ntv, we have the ability to parse and represent an affine
map with rank-0 results, viz. (i,j) -> (), we can pay off some engineering debt
in special casing the verification of such affine maps in dot-product flavored
vector.contract operations.

Reviewers: nicolasvasilache, andydavis1, rriddle

Reviewed By: nicolasvasilache

Subscribers: mehdi_amini, rriddle, jpienaar, burmako, shauheen, antiagainst, nicolasvasilache, arpith-jacob, mgester, lucyrfox, liufengdb, Joonsoo, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D76028
This commit is contained in:
aartbik 2020-03-11 14:26:00 -07:00
parent 70c0a9675d
commit 27dfb2257a
2 changed files with 5 additions and 11 deletions

View File

@ -302,13 +302,8 @@ static LogicalResult verify(ContractionOp op) {
<< index << " to have no symbols";
auto vectorType = op.getOperand(index).getType().dyn_cast<VectorType>();
unsigned rank = vectorType ? vectorType.getShape().size() : 0;
// Since (...) -> () is parsed into an empty map, we need to add
// a special case for this situation: continue the verification
// of an empty map if the resulting rank is indeed zero, i.e. this
// is a reduction into a scalar.
if (map.getNumDims() == 0 && map.getNumResults() == 0 && rank == 0)
continue;
// Verify that the map has the right number of inputs, outputs, and indices.
// This also correctly accounts for (..) -> () for rank-0 results.
if (map.getNumDims() != numIterators)
return op.emitOpError("expected indexing map ")
<< index << " to have " << numIterators << " number of inputs";

View File

@ -1077,6 +1077,7 @@ private:
// Helper to construct an affine map with one index removed.
static AffineMap adjustMap(AffineMap map, int64_t index,
PatternRewriter &rewriter) {
auto *ctx = rewriter.getContext();
SmallVector<AffineExpr, 4> results;
for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
int64_t idx = map.getResult(i).cast<AffineDimExpr>().getPosition();
@ -1084,13 +1085,11 @@ private:
continue;
// Re-insert remaining indices, but renamed when occurring
// after the removed index.
auto targetExpr =
getAffineDimExpr(idx < index ? idx : idx - 1, rewriter.getContext());
auto targetExpr = getAffineDimExpr(idx < index ? idx : idx - 1, ctx);
results.push_back(targetExpr);
}
// Since (...) -> () cannot be represented properly,
// we resort to an empty map when this situation happens.
return results.empty() ? AffineMap::get(rewriter.getContext())
// The (...) -> () affine map has its own factory method.
return results.empty() ? AffineMap::get(map.getNumDims() - 1, 0, ctx)
: AffineMap::get(map.getNumDims() - 1, 0, results);
}