Skip to content

Commit 9119826

Browse files
committed
[MLIR][OpenMP] Fix and simplify bounds offset calculation for 1-D GEP offsets
Currently this is being calculated incorrectly and will result in incorrect index offsets in more complicated array slices. This PR tries to address it by refactoring and changing the calculation.
1 parent 0a66411 commit 9119826

File tree

1 file changed

+15
-33
lines changed

1 file changed

+15
-33
lines changed

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Lines changed: 15 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -4230,46 +4230,28 @@ calculateBoundsOffset(LLVM::ModuleTranslation &moduleTranslation,
42304230
// with a pointer that's being treated like an array and we have the
42314231
// underlying type e.g. an i32, or f64 etc, e.g. a fortran descriptor base
42324232
// address (pointer pointing to the actual data) so we must caclulate the
4233-
// offset using a single index which the following two loops attempts to
4234-
// compute.
4235-
4236-
// Calculates the size offset we need to make per row e.g. first row or
4237-
// column only needs to be offset by one, but the next would have to be
4238-
// the previous row/column offset multiplied by the extent of current row.
4233+
// offset using a single index which the following loop attempts to
4234+
// compute using the standard column-major algorihtm e.g for a 3D array:
42394235
//
4240-
// For example ([1][10][100]):
4236+
// ((((c-idx * b-len) + b-idx) * a-len) + a_idx)
42414237
//
4242-
// - First row/column we move by 1 for each index increment
4243-
// - Second row/column we move by 1 (first row/column) * 10 (extent/size of
4244-
// current) for 10 for each index increment
4245-
// - Third row/column we would move by 10 (second row/column) *
4246-
// (extent/size of current) 100 for 1000 for each index increment
4247-
std::vector<llvm::Value *> dimensionIndexSizeOffset{builder.getInt64(1)};
4248-
for (size_t i = 1; i < bounds.size(); ++i) {
4249-
if (auto boundOp = dyn_cast_if_present<omp::MapBoundsOp>(
4250-
bounds[i].getDefiningOp())) {
4251-
dimensionIndexSizeOffset.push_back(builder.CreateMul(
4252-
moduleTranslation.lookupValue(boundOp.getExtent()),
4253-
dimensionIndexSizeOffset[i - 1]));
4254-
}
4255-
}
4256-
4257-
// Now that we have calculated how much we move by per index, we must
4258-
// multiply each lower bound offset in indexes by the size offset we
4259-
// have calculated in the previous and accumulate the results to get
4260-
// our final resulting offset.
4238+
// It is of note that it's doing column-major rather than row-major at the
4239+
// moment, but having a way for the frontend to indicate which major format
4240+
// to use or standardizing/canonicalizing the order of the bounds to compute
4241+
// the offset may be useful in the future when there's other frontends with
4242+
// different formats.
4243+
std::vector<llvm::Value *> dimensionIndexSizeOffset;
42614244
for (int i = bounds.size() - 1; i >= 0; --i) {
42624245
if (auto boundOp = dyn_cast_if_present<omp::MapBoundsOp>(
42634246
bounds[i].getDefiningOp())) {
4264-
if (idx.empty())
4265-
idx.emplace_back(builder.CreateMul(
4266-
moduleTranslation.lookupValue(boundOp.getLowerBound()),
4267-
dimensionIndexSizeOffset[i]));
4247+
if (i == bounds.size() - 1)
4248+
idx.emplace_back(
4249+
moduleTranslation.lookupValue(boundOp.getLowerBound()));
42684250
else
42694251
idx.back() = builder.CreateAdd(
4270-
idx.back(), builder.CreateMul(moduleTranslation.lookupValue(
4271-
boundOp.getLowerBound()),
4272-
dimensionIndexSizeOffset[i]));
4252+
builder.CreateMul(idx.back(), moduleTranslation.lookupValue(
4253+
boundOp.getExtent())),
4254+
moduleTranslation.lookupValue(boundOp.getLowerBound()));
42734255
}
42744256
}
42754257
}

0 commit comments

Comments
 (0)