Changeset View
Changeset View
Standalone View
Standalone View
mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
| Show First 20 Lines • Show All 266 Lines • ▼ Show 20 Lines | HoistingAnalysis::dropNonIndexDependencies(PadTensorOp padTensorOp, | ||||
| // %ubi = affine.min #map(%i) | // %ubi = affine.min #map(%i) | ||||
| // %ubj = affine.min #map(%j) | // %ubj = affine.min #map(%j) | ||||
| // %slice = tensor.extract_slice %source [%i, %j] [%ubi, %ubj] | // %slice = tensor.extract_slice %source [%i, %j] [%ubi, %ubj] | ||||
| // %padded_slice = linalg.pad_tensor %slice | // %padded_slice = linalg.pad_tensor %slice | ||||
| // ``` | // ``` | ||||
| // After iterating `backwardSlice` we obtain: | // After iterating `backwardSlice` we obtain: | ||||
| // indexEdges = [%i, %j, %ubi, %ubj] | // indexEdges = [%i, %j, %ubi, %ubj] | ||||
| // backwardSlice = backwardSlice / [linalg.fill(%cst, %arg1), scf.for %k] | // backwardSlice = backwardSlice / [linalg.fill(%cst, %arg1), scf.for %k] | ||||
| SetVector<Operation *> operationsToRemove; | |||||
| for (Operation *op : llvm::reverse(backwardSlice)) { | for (Operation *op : llvm::reverse(backwardSlice)) { | ||||
| // Add the index operands of `padTensorOp` and `sliceOp` to start the | // Add the index operands of `padTensorOp` and `sliceOp` to start the | ||||
| // exploration of the index computation. | // exploration of the index computation. | ||||
| if (op == padTensorOp || op == sliceOp) { | if (op == padTensorOp || op == sliceOp) { | ||||
| addIndexOperandsToIndexEdges(op); | addIndexOperandsToIndexEdges(op); | ||||
| continue; | continue; | ||||
| } | } | ||||
| // Add the index operands of the loop if its induction variable is | // Add the index operands of the loop if its induction variable is | ||||
| Show All 20 Lines | if (hasIndexResult(op)) { | ||||
| bool hasMemoryEffect = effectInterface && !effectInterface.hasNoEffect(); | bool hasMemoryEffect = effectInterface && !effectInterface.hasNoEffect(); | ||||
| if (hasMemoryEffect || op->getNumRegions() != 0) { | if (hasMemoryEffect || op->getNumRegions() != 0) { | ||||
| LLVM_DEBUG(DBGS() << "Unsupported op with region or memory effect: " | LLVM_DEBUG(DBGS() << "Unsupported op with region or memory effect: " | ||||
| << op << " -> skip\n"); | << op << " -> skip\n"); | ||||
| return failure(); | return failure(); | ||||
| } | } | ||||
| continue; | continue; | ||||
| } | } | ||||
| // Remove all other operation not used by the index computation except for | // Remove all other operations not used by the index computation. An | ||||
| // constant operations that may be padding values used by `padTensorOp`. | // exception are constant operations that may be used by `padTensorOp`. | ||||
| if (!isa<arith::ConstantOp>(op)) | if (!isa<arith::ConstantOp>(op)) | ||||
| backwardSlice.remove(op); | operationsToRemove.insert(op); | ||||
| } | } | ||||
| backwardSlice.set_subtract(operationsToRemove); | |||||
| return success(); | return success(); | ||||
| } | } | ||||
| SmallVector<Value> | SmallVector<Value> | ||||
| HoistingAnalysis::getPackedTensorSizes(ImplicitLocOpBuilder &b) { | HoistingAnalysis::getPackedTensorSizes(ImplicitLocOpBuilder &b) { | ||||
| SmallVector<Value> dynamicTensorSizes; | SmallVector<Value> dynamicTensorSizes; | ||||
| // Upper bound the packing loop lengths to size the packed tensor. Taking | // Upper bound the packing loop lengths to size the packed tensor. Taking | ||||
| ▲ Show 20 Lines • Show All 186 Lines • Show Last 20 Lines | |||||