Changeset View
Changeset View
Standalone View
Standalone View
mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
| Show All 10 Lines | |||||
| //===----------------------------------------------------------------------===// | //===----------------------------------------------------------------------===// | ||||
| #include "PassDetail.h" | #include "PassDetail.h" | ||||
| #include "mlir/Dialect/Affine/EDSC/Intrinsics.h" | #include "mlir/Dialect/Affine/EDSC/Intrinsics.h" | ||||
| #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" | #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" | ||||
| #include "mlir/Dialect/Linalg/Passes.h" | #include "mlir/Dialect/Linalg/Passes.h" | ||||
| #include "mlir/Dialect/Linalg/Transforms/Transforms.h" | #include "mlir/Dialect/Linalg/Transforms/Transforms.h" | ||||
| #include "mlir/Dialect/Linalg/Utils/Utils.h" | #include "mlir/Dialect/Linalg/Utils/Utils.h" | ||||
| #include "mlir/Dialect/MemRef/EDSC/Intrinsics.h" | |||||
| #include "mlir/Dialect/MemRef/IR/MemRef.h" | #include "mlir/Dialect/MemRef/IR/MemRef.h" | ||||
| #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h" | #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h" | ||||
| #include "mlir/Dialect/Tensor/IR/Tensor.h" | #include "mlir/Dialect/Tensor/IR/Tensor.h" | ||||
| #include "mlir/IR/AffineExpr.h" | #include "mlir/IR/AffineExpr.h" | ||||
| #include "mlir/IR/AffineMap.h" | #include "mlir/IR/AffineMap.h" | ||||
| #include "mlir/Transforms/FoldUtils.h" | #include "mlir/Transforms/FoldUtils.h" | ||||
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" | #include "mlir/Transforms/GreedyPatternRewriteDriver.h" | ||||
| ▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines | if (isZero(tileSizes[idx - zerosCount])) { | ||||
| continue; | continue; | ||||
| } | } | ||||
| loopIndexToRangeIndex[idx] = idx - zerosCount; | loopIndexToRangeIndex[idx] = idx - zerosCount; | ||||
| } | } | ||||
| // Create a new range with the applied tile sizes. | // Create a new range with the applied tile sizes. | ||||
| SmallVector<Range, 4> res; | SmallVector<Range, 4> res; | ||||
| for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx) | for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx) | ||||
| res.push_back( | res.push_back(Range{b.create<ConstantIndexOp>(loc, 0), shapeSizes[idx], | ||||
| Range{std_constant_index(0), shapeSizes[idx], tileSizes[idx]}); | tileSizes[idx]}); | ||||
| return std::make_tuple(res, loopIndexToRangeIndex); | return std::make_tuple(res, loopIndexToRangeIndex); | ||||
| } | } | ||||
| // All indices returned by IndexOp should be invariant with respect to tiling. | // All indices returned by IndexOp should be invariant with respect to tiling. | ||||
| // Therefore, if an operation is tiled, we have to transform the indices | // Therefore, if an operation is tiled, we have to transform the indices | ||||
| // accordingly, i.e. offset them by the values of the corresponding induction | // accordingly, i.e. offset them by the values of the corresponding induction | ||||
| // variables that are captured implicitly in the body of the op. | // variables that are captured implicitly in the body of the op. | ||||
| // | // | ||||
| ▲ Show 20 Lines • Show All 237 Lines • ▼ Show 20 Lines | Optional<TiledLinalgOp> static tileLinalgOpImpl( | ||||
| // Enforce the convention that "tiling by zero" skips tiling a particular | // Enforce the convention that "tiling by zero" skips tiling a particular | ||||
| // dimension. This convention is significantly simpler to handle instead of | // dimension. This convention is significantly simpler to handle instead of | ||||
| // adjusting affine maps to account for missing dimensions. | // adjusting affine maps to account for missing dimensions. | ||||
| auto nLoops = op.getNumLoops(); | auto nLoops = op.getNumLoops(); | ||||
| SmallVector<Value, 4> tileSizeVector = | SmallVector<Value, 4> tileSizeVector = | ||||
| options.tileSizeComputationFunction(b, op); | options.tileSizeComputationFunction(b, op); | ||||
| if (tileSizeVector.size() < nLoops) { | if (tileSizeVector.size() < nLoops) { | ||||
| auto zero = std_constant_index(0); | auto zero = b.create<ConstantIndexOp>(op.getLoc(), 0); | ||||
| tileSizeVector.append(nLoops - tileSizeVector.size(), zero); | tileSizeVector.append(nLoops - tileSizeVector.size(), zero); | ||||
| } | } | ||||
| return tileLinalgOpImpl<LoopTy>(b, op, tileSizeVector, options); | return tileLinalgOpImpl<LoopTy>(b, op, tileSizeVector, options); | ||||
| } | } | ||||
| Optional<TiledLinalgOp> | Optional<TiledLinalgOp> | ||||
| mlir::linalg::tileLinalgOp(OpBuilder &b, LinalgOp op, | mlir::linalg::tileLinalgOp(OpBuilder &b, LinalgOp op, | ||||
| ▲ Show 20 Lines • Show All 163 Lines • Show Last 20 Lines | |||||