Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[MemoryBanking] Bank memref::copy logically using memref::subView #8074

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

jiahanxie353
Copy link
Contributor

This patch banks memref::copy op logically when the operation is not used in parallel ops (which implies that we don't need to create new physical memories) by using memref::subView op.

return failure();

ArrayRef<int64_t> dstShape = dstType.getShape();
int64_t chunkSize = dstShape[bankingDimension] / (int64_t)bankingFactor;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please avoid implicit casts (here and above and below).

Suggested change
int64_t chunkSize = dstShape[bankingDimension] / (int64_t)bankingFactor;
int64_t chunkSize = dstShape[bankingDimension] / static_cast<int64_t>(bankingFactor);


ArrayRef<int64_t> srcShape = srcType.getShape();
// We'll assume static shapes and exact divisibility.
int64_t chunkSize = srcShape[bankingDimension] / (int64_t)bankingFactor;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same thing here, prefer explicit cast vs implicit.

/// Copy from unbanked source memref `srcMem` to destination memory references
/// `dstBanks`, where each bank has the same MemRefType, and `srcMem` has the
/// shape of the "original" unbanked array.
LogicalResult copyFromUnbankedToBanks(PatternRewriter &rewriter, Location loc,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems to be almost exactly the same as the function above. Can we add additional arguments or template parameters to avoid having two nearly duplicate functions? WDYT?

Value dstMem = copyOp.getTarget();

bool srcIsBanked = memoryToBanks.count(srcMem) != 0;
bool dstIsBanked = memoryToBanks.count(dstMem) != 0;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Again, a good habit here to avoid double lookups is

auto srcIt = memoryToBanks.find(srcMem);
auto dstIt = memoryToBanks.find(dstMem);
bool srcIsBanked = srcIt != memoryToBanks.end();
bool dstIsBanked = dstIt != memoryToBanks.end();

...

auto& srcBanks = srcIt->second;

if (d == bankingDimension) {
// Round-robin partition: offset = bankIndex, size = chunkSize, stride =
// bankingFactor
int64_t chunkSize = dimSize / (int64_t)bankingFactor;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here

@@ -409,6 +409,151 @@ struct BankAffineStorePattern
DenseSet<Value> &oldMemRefVals;
};

struct BankMemrefCopyPattern : public OpRewritePattern<memref::CopyOp> {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add a comment here please describing the purpose/goal of this rewrite.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants