colin
/
onnx-mlir
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Set alignment attribute to AllocOp of insertAllocAndDealloc() (#195) 

* Added alignment in AllocOp of insertAllocAndDealloc()

* Added comments

Co-authored-by: Tian Jin <tjingrant@gmail.com>
This commit is contained in:
Haruki Imai 2020-07-02 16:21:01 +09:00 committed by GitHub
parent 5c6d85e6f3
commit 4d96247327
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.cpp
View File

@ -45,7 +45,8 @@ MemRefType convertToMemRefType(Type type) {
/// Insert an allocation and deallocation for the given MemRefType. /// Insert an allocation and deallocation for the given MemRefType.
Value insertAllocAndDealloc(MemRefType type, Location loc, Value insertAllocAndDealloc(MemRefType type, Location loc,
PatternRewriter &rewriter, bool insertDealloc, ArrayRef<Value> operands) { PatternRewriter &rewriter, bool insertDealloc, ArrayRef<Value> operands,
int64_t alignment) {
// Put together alloc operands for any dynamic dimensions of the memref. // Put together alloc operands for any dynamic dimensions of the memref.
AllocOp alloc; AllocOp alloc;
if (!operands.empty()) { if (!operands.empty()) {
@ -87,10 +88,27 @@ Value insertAllocAndDealloc(MemRefType type, Location loc,
for (int i = 0; i < rank; ++i) for (int i = 0; i < rank; ++i)
if (memRefShape[i] < 0) if (memRefShape[i] < 0)
allocOperands.push_back(fromOperands[i]); allocOperands.push_back(fromOperands[i]);
// Set alignment attribute. Default value is `-1`, which does not set
// alignment.
if (alignment >= 0) {
IntegerAttr constAlignAttr = rewriter.getI64IntegerAttr(alignment);
alloc =
rewriter.create<AllocOp>(loc, type, allocOperands, constAlignAttr);
} else {
alloc = rewriter.create<AllocOp>(loc, type, allocOperands); alloc = rewriter.create<AllocOp>(loc, type, allocOperands);
}
} else {
// Set alignment attribute. Default value is `-1`, which does not set
// alignment.
if (alignment >= 0) {
SmallVector<Value, 4> allocOperandsEmpty;
IntegerAttr constAlignAttr = rewriter.getI64IntegerAttr(alignment);
alloc = rewriter.create<AllocOp>(
loc, type, allocOperandsEmpty, constAlignAttr);
} else { } else {
alloc = rewriter.create<AllocOp>(loc, type); alloc = rewriter.create<AllocOp>(loc, type);
} }
}
// Make sure to allocate at the beginning of the block if // Make sure to allocate at the beginning of the block if
// all dimensions are known. // all dimensions are known.

2
src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp
View File

@ -45,7 +45,7 @@ MemRefType convertToMemRefType(Type type);
/// Insert an allocation and deallocation for the given MemRefType. /// Insert an allocation and deallocation for the given MemRefType.
Value insertAllocAndDealloc(MemRefType type, Location loc, Value insertAllocAndDealloc(MemRefType type, Location loc,
PatternRewriter &rewriter, bool insertDealloc, PatternRewriter &rewriter, bool insertDealloc,
ArrayRef<Value> operands = {}); ArrayRef<Value> operands = {}, int64_t alignment = -1);
// Determine if current function returns the result value of the // Determine if current function returns the result value of the
// current op being lowered. If it does then dealloc should not be // current op being lowered. If it does then dealloc should not be