Changeset View
Changeset View
Standalone View
Standalone View
mlir/lib/Dialect/Linalg/Utils/Utils.cpp
| Show First 20 Lines • Show All 519 Lines • ▼ Show 20 Lines | SmallVector<Value, 4> makeTiledShapes(OpBuilder &b, Location loc, | ||||
| using namespace edsc::op; | using namespace edsc::op; | ||||
| // Construct (potentially temporary) mins and maxes on which to apply maps | // Construct (potentially temporary) mins and maxes on which to apply maps | ||||
| // that define tile subshapes. | // that define tile subshapes. | ||||
| SmallVector<Value, 8> lbs, subShapeSizes; | SmallVector<Value, 8> lbs, subShapeSizes; | ||||
| for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) { | for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) { | ||||
| LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n"); | LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n"); | ||||
| bool isTiled = !isZero(tileSizes[idx]); | bool isTiled = !isZero(tileSizes[idx]); | ||||
| lbs.push_back(isTiled ? ivs[idxIvs++] : (Value)std_constant_index(0)); | lbs.push_back(isTiled ? ivs[idxIvs++] | ||||
| : (Value)b.create<ConstantIndexOp>(loc, 0)); | |||||
| // Before composing, we need to make range a closed interval. | // Before composing, we need to make range a closed interval. | ||||
| Value size = isTiled ? tileSizes[idx] : sizeBounds[idx]; | Value size = isTiled ? tileSizes[idx] : sizeBounds[idx]; | ||||
| subShapeSizes.push_back(size - std_constant_index(1)); | subShapeSizes.push_back(size - b.create<ConstantIndexOp>(loc, 1)); | ||||
| LLVM_DEBUG(llvm::dbgs() << "lb: " << lbs.back() << "\n"); | LLVM_DEBUG(llvm::dbgs() << "lb: " << lbs.back() << "\n"); | ||||
| LLVM_DEBUG(llvm::dbgs() << "size: " << subShapeSizes.back() << "\n"); | LLVM_DEBUG(llvm::dbgs() << "size: " << subShapeSizes.back() << "\n"); | ||||
| } | } | ||||
| MLIRContext *context = b.getContext(); | MLIRContext *context = b.getContext(); | ||||
| SmallVector<Value, 4> tiledShapes; | SmallVector<Value, 4> tiledShapes; | ||||
| tiledShapes.reserve(tiledOperands.size()); | tiledShapes.reserve(tiledOperands.size()); | ||||
| for (auto en : llvm::enumerate(tiledOperands)) { | for (auto en : llvm::enumerate(tiledOperands)) { | ||||
| Show All 15 Lines | for (auto en : llvm::enumerate(tiledOperands)) { | ||||
| SmallVector<OpFoldResult, 4> offsets, sizes, strides; | SmallVector<OpFoldResult, 4> offsets, sizes, strides; | ||||
| offsets.reserve(rank); | offsets.reserve(rank); | ||||
| sizes.reserve(rank); | sizes.reserve(rank); | ||||
| strides.reserve(rank); | strides.reserve(rank); | ||||
| for (unsigned r = 0; r < rank; ++r) { | for (unsigned r = 0; r < rank; ++r) { | ||||
| LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for dim#" << r); | LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for dim#" << r); | ||||
| if (!isTiled(map.getSubMap({r}), tileSizes)) { | if (!isTiled(map.getSubMap({r}), tileSizes)) { | ||||
| offsets.push_back(b.getIndexAttr(0)); | offsets.push_back(b.getIndexAttr(0)); | ||||
| Value dim = memref_dim(shapedOp, r).value; | Value dim = b.createOrFold<memref::DimOp>(loc, shapedOp, r); | ||||
| sizes.push_back(dim); | sizes.push_back(dim); | ||||
| strides.push_back(b.getIndexAttr(1)); | strides.push_back(b.getIndexAttr(1)); | ||||
| LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n"); | LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n"); | ||||
| continue; | continue; | ||||
| } | } | ||||
| LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n"); | LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n"); | ||||
| // Tiling creates a new slice at the proper index, the slice step is 1 | // Tiling creates a new slice at the proper index, the slice step is 1 | ||||
| // (i.e. the op does not subsample, stepping occurs in the loop). | // (i.e. the op does not subsample, stepping occurs in the loop). | ||||
| auto m = map.getSubMap({r}); | auto m = map.getSubMap({r}); | ||||
| LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: submap: " << map << "\n"); | LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: submap: " << map << "\n"); | ||||
| auto offset = applyMapToValues(b, loc, m, lbs).front(); | auto offset = applyMapToValues(b, loc, m, lbs).front(); | ||||
| offsets.push_back(offset); | offsets.push_back(offset); | ||||
| auto closedIntSize = applyMapToValues(b, loc, m, subShapeSizes).front(); | auto closedIntSize = applyMapToValues(b, loc, m, subShapeSizes).front(); | ||||
| // Resulting size needs to be made half open interval again. | // Resulting size needs to be made half open interval again. | ||||
| auto size = closedIntSize + std_constant_index(1); | auto size = closedIntSize + b.create<ConstantIndexOp>(loc, 1); | ||||
| LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: raw size: " << size << "\n"); | LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: raw size: " << size << "\n"); | ||||
| // The size of the subview / subtensor should be trimmed to avoid | // The size of the subview / subtensor should be trimmed to avoid | ||||
| // out-of-bounds accesses, unless we statically know the subshape size | // out-of-bounds accesses, unless we statically know the subshape size | ||||
| // divides the shape size evenly. | // divides the shape size evenly. | ||||
| int64_t shapeSize = shapedType.getDimSize(r); | int64_t shapeSize = shapedType.getDimSize(r); | ||||
| auto sizeCst = size.getDefiningOp<ConstantIndexOp>(); | auto sizeCst = size.getDefiningOp<ConstantIndexOp>(); | ||||
| if (ShapedType::isDynamic(shapeSize) || !sizeCst || | if (ShapedType::isDynamic(shapeSize) || !sizeCst || | ||||
| Show All 37 Lines | |||||