diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h --- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h +++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h @@ -215,6 +215,42 @@ Value materializeOpFoldResult(ImplicitLocOpBuilder &builder, OpFoldResult opFoldResult); +/// A struct containg offsets-sizes-strides arguments of the tiled shape. +struct SliceParameters { + SmallVector offsets; + SmallVector sizes; + SmallVector strides; +}; + +/// Computes SliceParameters for a single `valueToTile`. `omitPartialTileCheck` +/// controls whether to omit the partial/boundary tile condition check in cases +/// where we statically know that it is unnecessary. +SliceParameters computeSliceParameters(OpBuilder &builder, Location loc, + Value valueToTile, ValueRange tileSizes, + AffineMap map, ValueRange lbs, + ValueRange ubs, ValueRange subShapeSizes, + bool omitPartialTileCheck); + +/// Computes SliceParamaters for all `valuesToTile` of the given +/// `linalgOp`, assuming `linalgOp` is being fused into a loop +/// nest for tiling with the given induction variables `ivs` and tile sizes +/// `tileSizes`. `sizeBounds` are the iteration space bounds for *all* the +/// implicit loops in `linalgOp`. `omitPartialTileCheck` controls whether to +/// omit the partial/boundary tile condition check in cases where we statically +/// know that it is unnecessary. +/// +/// Note that a constant zero in `tileSizes` means no tiling at that implicit +/// loop. The number of non-zero values in `tileSizes` should be equal to the +/// number of values in `ivs`. +/// +/// Some of the `valuesToTile` won't be affected by tiling. For these values, +/// llvm::None will be returned. +SmallVector, 4> +computeAllSliceParameters(OpBuilder &builder, Location loc, LinalgOp linalgOp, + ArrayRef valuesToTile, ValueRange ivs, + ValueRange tileSizes, ArrayRef sizeBounds, + bool omitPartialTileCheck); + /// Creates an extract_slice/subview op for a single `valueToTile` with /// `builder`. This new operation extracts a tile of `valueToTile`, starting /// at offsets `lbs` and with sizes `subShapeSizes`. `omitPartialTileCheck` diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp --- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp @@ -796,27 +796,59 @@ return b.createOrFold(loc, map, normalizedOperands); } +static Value materializeTiledShape(OpBuilder &builder, Location loc, + Value valueToTile, + const SliceParameters &sliceParams) { + auto shapedType = valueToTile.getType().dyn_cast(); + auto *sliceOp = TypeSwitch(shapedType) + .Case([&](MemRefType) { + return builder.create( + loc, valueToTile, sliceParams.offsets, + sliceParams.sizes, sliceParams.strides); + }) + .Case([&](RankedTensorType) { + return makeComposedExtractSliceOp( + builder, loc, valueToTile, sliceParams.offsets, + sliceParams.sizes, sliceParams.strides); + }) + .Default([](ShapedType) -> Operation * { + llvm_unreachable("Unexpected shaped type"); + }); + return sliceOp->getResult(0); +} + Value makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile, ValueRange tileSizes, AffineMap map, ValueRange lbs, ValueRange ubs, ValueRange subShapeSizes, bool omitPartialTileCheck) { + SliceParameters sliceParams = + computeSliceParameters(builder, loc, valueToTile, tileSizes, map, lbs, + ubs, subShapeSizes, omitPartialTileCheck); + return materializeTiledShape(builder, loc, valueToTile, sliceParams); +} + +SliceParameters computeSliceParameters(OpBuilder &builder, Location loc, + Value valueToTile, ValueRange tileSizes, + AffineMap map, ValueRange lbs, + ValueRange ubs, ValueRange subShapeSizes, + bool omitPartialTileCheck) { auto shapedType = valueToTile.getType().dyn_cast(); assert(shapedType && "only shaped types can be tiled"); ArrayRef shape = shapedType.getShape(); int64_t rank = shapedType.getRank(); // Construct a new subview / extract_slice for the tile. - SmallVector offsets, sizes, strides; - offsets.reserve(rank); - sizes.reserve(rank); - strides.reserve(rank); + SliceParameters sliceParams; + sliceParams.offsets.reserve(rank); + sliceParams.sizes.reserve(rank); + sliceParams.strides.reserve(rank); for (unsigned r = 0; r < rank; ++r) { - LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: for dim#" << r); + LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: for dim#" << r); if (!isTiled(map.getSubMap({r}), tileSizes)) { - offsets.push_back(builder.getIndexAttr(0)); + sliceParams.offsets.push_back(builder.getIndexAttr(0)); Value dim = createOrFoldDimOp(builder, loc, valueToTile, r); - sizes.push_back(getAsOpFoldResult(dim)); - strides.push_back(builder.getIndexAttr(1)); + sliceParams.sizes.push_back(getAsOpFoldResult(dim)); + sliceParams.strides.push_back(builder.getIndexAttr(1)); LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n"); continue; } @@ -825,25 +857,27 @@ // 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). auto m = map.getSubMap({r}); - LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: submap: " << m << "\n"); + LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: submap: " << m << "\n"); auto offset = applyMapToValues(builder, loc, m, lbs).front(); - offsets.push_back(getAsOpFoldResult(offset)); + sliceParams.offsets.push_back(getAsOpFoldResult(offset)); auto closedIntSize = applyMapToValues(builder, loc, m, subShapeSizes).front(); // Resulting size needs to be made half open interval again. AffineExpr s0 = getAffineSymbolExpr(0, builder.getContext()); Value size = fullyComposeAndAffineApply(builder, loc, s0 + 1, closedIntSize); - LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: raw size: " << size << "\n"); LLVM_DEBUG(llvm::dbgs() - << "makeTiledShape: new offset: " << offset << "\n"); - strides.push_back(builder.getIndexAttr(1)); + << "computeSliceParameters: raw size: " << size << "\n"); + LLVM_DEBUG(llvm::dbgs() + << "computeSliceParameters: new offset: " << offset << "\n"); + sliceParams.strides.push_back(builder.getIndexAttr(1)); if (omitPartialTileCheck) { // We statically know that the partial/boundary tile condition is // unnecessary. - LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: new size: " << size << "\n"); - sizes.push_back(getAsOpFoldResult(size)); + LLVM_DEBUG(llvm::dbgs() + << "computeSliceParameters: new size: " << size << "\n"); + sliceParams.sizes.push_back(getAsOpFoldResult(size)); continue; } @@ -895,22 +929,9 @@ operands); } LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: new size: " << size << "\n"); - sizes.push_back(getAsOpFoldResult(size)); + sliceParams.sizes.push_back(getAsOpFoldResult(size)); } - - auto *sliceOp = TypeSwitch(shapedType) - .Case([&](MemRefType) { - return builder.create( - loc, valueToTile, offsets, sizes, strides); - }) - .Case([&](RankedTensorType) { - return makeComposedExtractSliceOp( - builder, loc, valueToTile, offsets, sizes, strides); - }) - .Default([](ShapedType) -> Operation * { - llvm_unreachable("Unexpected shaped type"); - }); - return sliceOp->getResult(0); + return sliceParams; } Value createSlice(OpBuilder &builder, Location loc, Value value, @@ -1001,12 +1022,11 @@ return builder.create(attr.getValue().getSExtValue()); } -SmallVector makeTiledShapes(OpBuilder &b, Location loc, - LinalgOp linalgOp, - ArrayRef valuesToTile, - ValueRange ivs, ValueRange tileSizes, - ArrayRef sizeBounds, - bool omitPartialTileCheck) { +SmallVector, 4> +computeAllSliceParameters(OpBuilder &builder, Location loc, LinalgOp linalgOp, + ArrayRef valuesToTile, ValueRange ivs, + ValueRange tileSizes, ArrayRef sizeBounds, + bool omitPartialTileCheck) { assert(ivs.size() == static_cast(llvm::count_if( llvm::make_range(tileSizes.begin(), tileSizes.end()), [](Value v) { return !isZero(v); })) && @@ -1014,15 +1034,15 @@ // Construct (potentially temporary) mins and maxes on which to apply maps // that define tile subshapes. - SmallVector lbs = computeTileOffsets(b, loc, ivs, tileSizes); + SmallVector lbs = computeTileOffsets(builder, loc, ivs, tileSizes); SmallVector subShapeSizes = - computeTileSizes(b, loc, tileSizes, sizeBounds); + computeTileSizes(builder, loc, tileSizes, sizeBounds); assert(static_cast(valuesToTile.size()) == linalgOp.getNumInputsAndOutputs() && "expected one value to tile for every operand"); - SmallVector tiledShapes; - tiledShapes.reserve(valuesToTile.size()); + SmallVector, 4> allSliceParams; + allSliceParams.reserve(valuesToTile.size()); for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) { Value shapedOp = valuesToTile[opOperand->getOperandNumber()]; LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp); @@ -1033,18 +1053,39 @@ // extract/insert slice pairs make the accessed iteration argument // subdomains explicit. if (!isTiled(map, tileSizes) && !linalgOp.isOutputTensor(opOperand)) { - tiledShapes.push_back(shapedOp); + allSliceParams.push_back(llvm::None); LLVM_DEBUG(llvm::dbgs() << ": not tiled: use shape: " << opOperand->get().getType() << "\n"); continue; } LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subshape...\n"); - tiledShapes.push_back(makeTiledShape(b, loc, shapedOp, tileSizes, map, lbs, - sizeBounds, subShapeSizes, - omitPartialTileCheck)); + allSliceParams.push_back(computeSliceParameters( + builder, loc, shapedOp, tileSizes, map, lbs, sizeBounds, subShapeSizes, + omitPartialTileCheck)); } + return allSliceParams; +} + +SmallVector makeTiledShapes(OpBuilder &b, Location loc, + LinalgOp linalgOp, + ArrayRef valuesToTile, + ValueRange ivs, ValueRange tileSizes, + ArrayRef sizeBounds, + bool omitPartialTileCheck) { + SmallVector, 4> allSliceParameter = + computeAllSliceParameters(b, loc, linalgOp, valuesToTile, ivs, tileSizes, + sizeBounds, omitPartialTileCheck); + SmallVector tiledShapes; + for (auto item : llvm::zip(valuesToTile, allSliceParameter)) { + Value valueToTile = std::get<0>(item); + Optional sliceParams = std::get<1>(item); + tiledShapes.push_back( + sliceParams.hasValue() + ? materializeTiledShape(b, loc, valueToTile, *sliceParams) + : valueToTile); + } return tiledShapes; } diff --git a/mlir/test/Dialect/Linalg/tile-and-distribute.mlir b/mlir/test/Dialect/Linalg/tile-and-distribute.mlir --- a/mlir/test/Dialect/Linalg/tile-and-distribute.mlir +++ b/mlir/test/Dialect/Linalg/tile-and-distribute.mlir @@ -16,12 +16,12 @@ // CHECK-DAG: %[[BIDX:.*]] = gpu.block_id x // CHECK: scf.for %[[ARG3:.*]] = // CHECK: %[[OFFSETY:.*]] = affine.apply #[[MAP0]]()[%[[BIDY]]] -// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[OFFSETY]], %[[ARG3]]] // CHECK: %[[OFFSETX:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]] -// CHECK: %[[SV2:.*]] = memref.subview %[[ARG1]][%[[ARG3]], %[[OFFSETX]]] // CHECK: %[[OFFSETY_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDY]]] -// CHECK: %[[OFFSETX:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]] -// CHECK: %[[SV3:.*]] = memref.subview %[[ARG2]][%[[OFFSETY_2]], %[[OFFSETX]]] +// CHECK: %[[OFFSETX_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]] +// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[OFFSETY]], %[[ARG3]]] +// CHECK: %[[SV2:.*]] = memref.subview %[[ARG1]][%[[ARG3]], %[[OFFSETX]]] +// CHECK: %[[SV3:.*]] = memref.subview %[[ARG2]][%[[OFFSETY_2]], %[[OFFSETX_2]]] // CHECK: linalg.matmul ins(%[[SV1]], %[[SV2]]{{.*}} outs(%[[SV3]] // ----- @@ -48,11 +48,11 @@ // CHECK: scf.if %[[INBOUNDS]] // CHECK: scf.for %[[ARG3:.*]] = // CHECK: %[[OFFSETY:.*]] = affine.apply #[[MAP0]]()[%[[BIDY]]] -// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[OFFSETY]], %[[ARG3]]] // CHECK: %[[OFFSETX:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]] -// CHECK: %[[SV2:.*]] = memref.subview %[[ARG1]][%[[ARG3]], %[[OFFSETX]]] // CHECK: %[[OFFSETY_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDY]]] // CHECK: %[[OFFSETX_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]] +// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[OFFSETY]], %[[ARG3]]] +// CHECK: %[[SV2:.*]] = memref.subview %[[ARG1]][%[[ARG3]], %[[OFFSETX]]] // CHECK: %[[SV3:.*]] = memref.subview %[[ARG2]][%[[OFFSETY_2]], %[[OFFSETX_2]]] // CHECK: linalg.matmul ins(%[[SV1]], %[[SV2]]{{.*}} outs(%[[SV3]] @@ -106,11 +106,11 @@ // CHECK: scf.if %[[INBOUNDS]] // CHECK: scf.for %[[ARG3:.*]] = // CHECK: %[[OFFSETY:.*]] = affine.apply #[[MAP0]]()[%[[BIDY]]] -// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[OFFSETY]], %[[ARG3]]] // CHECK: %[[OFFSETX:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]] -// CHECK: %[[SV2:.*]] = memref.subview %[[ARG1]][%[[ARG3]], %[[OFFSETX]]] // CHECK: %[[OFFSETY_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDY]]] // CHECK: %[[OFFSETX_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]] +// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[OFFSETY]], %[[ARG3]]] +// CHECK: %[[SV2:.*]] = memref.subview %[[ARG1]][%[[ARG3]], %[[OFFSETX]]] // CHECK: %[[SV3:.*]] = memref.subview %[[ARG2]][%[[OFFSETY_2]], %[[OFFSETX_2]]] // CHECK: linalg.matmul ins(%[[SV1]], %[[SV2]]{{.*}} outs(%[[SV3]] @@ -139,9 +139,9 @@ // CHECK: scf.parallel (%[[ARG3:.*]]) = (%[[LBX]]) to (%{{.*}}) step (%[[STEPX]]) // CHECK: scf.for %[[ARG4:.*]] = // CHECK: %[[OFFSETY:.*]] = affine.apply #[[MAP0]]()[%[[BIDY]]] +// CHECK: %[[OFFSETY_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDY]]] // CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[OFFSETY]], %[[ARG4]]] // CHECK: %[[SV2:.*]] = memref.subview %[[ARG1]][%[[ARG4]], %[[ARG3]]] -// CHECK: %[[OFFSETY_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDY]]] // CHECK: %[[SV3:.*]] = memref.subview %[[ARG2]][%[[OFFSETY_2]], %[[ARG3]]] // CHECK: linalg.matmul ins(%[[SV1]], %[[SV2]]{{.*}} outs(%[[SV3]] @@ -166,10 +166,10 @@ // CHECK: %[[STEPY:.*]] = affine.apply #[[MAP0]]()[%[[NBLOCKSY]]] // CHECK: scf.parallel (%[[ARG3:.*]]) = (%[[LBY]]) to (%{{.*}}) step (%[[STEPY]]) // CHECK: scf.for %[[ARG4:.*]] = -// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[ARG3]], %[[ARG4]]] // CHECK: %[[OFFSETX:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]] -// CHECK: %[[SV2:.*]] = memref.subview %[[ARG1]][%[[ARG4]], %[[OFFSETX]]] // CHECK: %[[OFFSETX_2:.*]] = affine.apply #[[MAP0]]()[%[[BIDX]]] +// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[ARG3]], %[[ARG4]]] +// CHECK: %[[SV2:.*]] = memref.subview %[[ARG1]][%[[ARG4]], %[[OFFSETX]]] // CHECK: %[[SV3:.*]] = memref.subview %[[ARG2]][%[[ARG3]], %[[OFFSETX_2]]] // CHECK: linalg.matmul ins(%[[SV1]], %[[SV2]]{{.*}} outs(%[[SV3]] diff --git a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir --- a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir +++ b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir @@ -241,9 +241,9 @@ // CHECK-NEXT: %[[SIZE_ELEM_OC:.+]] = affine.min #[[BOUND2_MAP]](%[[IV2]])[%[[ELEM_OC]]] // CHECK-NEXT: %[[OFFSET_OW:.+]] = affine.apply #[[X2_MAP]](%[[IV2]]) // CHECK-NEXT: %[[SIZE_INPUT_W:.+]] = affine.min #[[INPUT_BOUND]](%[[SIZE_ELEM_OW]], %[[IV2]])[%[[FILTER_W]], %[[FILL_W]]] +// CHECK-NEXT: %[[SIZE_ELEM_OW_2:.+]] = affine.min #[[BOUND4_MAP_2]](%[[IV2]])[%[[FILL_W]], %[[ELEM_OW]]] // CHECK-NEXT: %[[ST_INPUT:.+]] = tensor.extract_slice %[[INPUT]][%[[IV0]], %[[OFFSET_OH]], %[[OFFSET_OW]], 0] // CHECK-SAME: [%[[SIZE_INPUT_N]], %[[SIZE_INPUT_H]], %[[SIZE_INPUT_W]], %[[INPUT_C]]] -// CHECK-NEXT: %[[SIZE_ELEM_OW_2:.+]] = affine.min #[[BOUND4_MAP_2]](%[[IV2]])[%[[FILL_W]], %[[ELEM_OW]]] // CHECK-NEXT: scf.for %[[IV3:.+]] = %{{.+}} to %[[ELEM_OC]] step %{{.+}} iter_args(%[[ARG:[a-z0-9]+]] // CHECK-NEXT: %[[ST_ELEM:.+]] = tensor.extract_slice %[[ELEM]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] // CHECK-SAME: [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]] diff --git a/mlir/test/Dialect/Linalg/tile-conv.mlir b/mlir/test/Dialect/Linalg/tile-conv.mlir --- a/mlir/test/Dialect/Linalg/tile-conv.mlir +++ b/mlir/test/Dialect/Linalg/tile-conv.mlir @@ -26,9 +26,9 @@ // CHECK: scf.for %[[ARG4:.*]] = %[[C0]] to %[[T3]] step %[[C3]] // CHECK: %[[T4:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T0]], %[[T2]]] // CHECK: %[[T5:.*]] = affine.min #[[MAP1]](%[[ARG4]])[%[[T1]], %[[T3]]] -// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[ARG3]], %[[ARG4]]] [%[[T4]], %[[T5]]] // CHECK: %[[T6:.*]] = affine.min #[[MAP2]](%[[ARG3]])[%[[T2]] // CHECK: %[[T7:.*]] = affine.min #[[MAP3]](%[[ARG4]])[%[[T3]]] +// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[ARG3]], %[[ARG4]]] [%[[T4]], %[[T5]]] // CHECK: %[[SV2:.*]] = memref.subview %[[ARG2]][%[[ARG3]], %[[ARG4]]] [%[[T6]], %[[T7]]] // CHECK: linalg.conv_2d // CHECK-SAME: ins(%[[SV1]], %[[ARG1]] diff --git a/mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir b/mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir --- a/mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir +++ b/mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir @@ -64,12 +64,12 @@ // CHECK-NOT: affine.min // CHECK-NOT: affine.max // CHECK: %[[LB0:.+]] = affine.apply #[[$map2]](%[[IV0]]) - // CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] : // CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]]) + // CHECK: %[[LB0_1:.+]] = affine.apply #[[$map2]](%[[IV0]]) + // CHECK: %[[LB1_1:.+]] = affine.apply #[[$map3]](%[[IV1]]) + // CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] : // CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] : - // CHECK: %[[LB0:.+]] = affine.apply #[[$map2]](%[[IV0]]) - // CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]]) - // CHECK: %[[tC:.+]] = tensor.extract_slice %[[C]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] : + // CHECK: %[[tC:.+]] = tensor.extract_slice %[[C]][%[[LB0_1]], %[[LB1_1]]] [10, %[[TS]]] [1, 1] : // CHECK: linalg.matmul // CHECK: scf.foreach_thread.perform_concurrently // CHECK-NEXT: tensor.parallel_insert_slice @@ -155,12 +155,12 @@ // CHECK-NOT: affine.max // CHECK-NOT: affine.min // CHECK: %[[LB0:.+]] = affine.apply #[[$map2]](%[[IV0]]) - // CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] : // CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]]) + // CHECK: %[[LB0_1:.+]] = affine.apply #[[$map2]](%[[IV0]]) + // CHECK: %[[LB1_1:.+]] = affine.apply #[[$map3]](%[[IV1]]) + // CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] : // CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] : - // CHECK: %[[LB0:.+]] = affine.apply #[[$map2]](%[[IV0]]) - // CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]]) - // CHECK: %[[tC:.+]] = tensor.extract_slice %[[C]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] : + // CHECK: %[[tC:.+]] = tensor.extract_slice %[[C]][%[[LB0_1]], %[[LB1_1]]] [10, %[[TS]]] [1, 1] : // CHECK: linalg.matmul // CHECK: scf.foreach_thread.perform_concurrently // CHECK-NEXT: tensor.parallel_insert_slice diff --git a/mlir/test/Dialect/Linalg/tile.mlir b/mlir/test/Dialect/Linalg/tile.mlir --- a/mlir/test/Dialect/Linalg/tile.mlir +++ b/mlir/test/Dialect/Linalg/tile.mlir @@ -40,9 +40,9 @@ // TILE-2: scf.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { // TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[M]]] // TILE-2: %[[K:.*]] = memref.dim %{{.*}}, %c1 : memref -// TILE-2: %[[sAi:.*]] = memref.subview %{{.*}}[%[[I]], 0] [%[[szM]], %[[K]]] [1, 1] : memref to memref // TILE-2: %[[szK:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[M]]] // TILE-2: %[[N:.*]] = memref.dim %{{.*}}, %c1 : memref +// TILE-2: %[[sAi:.*]] = memref.subview %{{.*}}[%[[I]], 0] [%[[szM]], %[[K]]] [1, 1] : memref to memref // TILE-2: %[[sCi:.*]] = memref.subview %{{.*}}[%[[I]], 0] [%[[szK]], %[[N]]] [1, 1] : memref to memref // TILE-2: linalg.matmul ins(%[[sAi]]{{.*}} outs(%[[sCi]] @@ -53,9 +53,9 @@ // TILE-02: scf.for %[[J:.*]] = %{{.*}} to %[[N]] step %{{.*}} { // TILE-02: %[[K:.*]] = memref.dim %{{.*}}, %c0 : memref // TILE-02: %[[szN:.*]] = affine.min #[[$bound_map]](%[[J]])[%[[N]]] -// TILE-02: %[[sBj:.*]] = memref.subview %{{.*}}[0, %[[J]]] [%[[K]], %[[szN]]] [1, 1] : memref to memref // TILE-02: %[[M:.*]] = memref.dim %{{.*}}, %c0 : memref // TILE-02: %[[szK:.*]] = affine.min #[[$bound_map]](%[[J]])[%[[N]]] +// TILE-02: %[[sBj:.*]] = memref.subview %{{.*}}[0, %[[J]]] [%[[K]], %[[szN]]] [1, 1] : memref to memref // TILE-02: %[[sCj:.*]] = memref.subview %{{.*}}[0, %[[J]]] [%[[M]], %[[szK]]] [1, 1] : memref to memref // TILE-02: linalg.matmul ins(%{{.*}}, %[[sBj]]{{.*}} outs(%[[sCj]] @@ -66,10 +66,10 @@ // TILE-002: scf.for %[[K:.*]] = %{{.*}}{{.*}} to %[[ubK]] step %{{.*}} { // TILE-002: %[[M:.*]] = memref.dim %{{.*}}, %c0 : memref // TILE-002: %[[szK:.*]] = affine.min #[[$bound_map]](%[[K]])[%[[ubK]]] -// TILE-002: %[[sAj:.*]] = memref.subview %{{.*}}[0, %[[K]]] [%[[M]], %[[szK]]] [1, 1] : memref to memref -// TILE-002: %[[szK:.*]] = affine.min #[[$bound_map]](%[[K]])[%[[ubK]]] +// TILE-002: %[[szK_1:.*]] = affine.min #[[$bound_map]](%[[K]])[%[[ubK]]] // TILE-002: %[[N:.*]] = memref.dim %{{.*}}, %c1 : memref -// TILE-002: %[[sBj:.*]] = memref.subview %{{.*}}[%[[K]], 0] [%[[szK]], %[[N]]] [1, 1] : memref to memref +// TILE-002: %[[sAj:.*]] = memref.subview %{{.*}}[0, %[[K]]] [%[[M]], %[[szK]]] [1, 1] : memref to memref +// TILE-002: %[[sBj:.*]] = memref.subview %{{.*}}[%[[K]], 0] [%[[szK_1]], %[[N]]] [1, 1] : memref to memref // TILE-002: linalg.matmul ins(%[[sAj]], %[[sBj]]{{.*}} outs(%{{.*}} // TILE-234-LABEL: func @matmul( @@ -85,13 +85,13 @@ // TILE-234: scf.for %[[K:.*]] = %{{.*}}{{.*}} to %[[ubK]] step %{{.*}} { // TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[ubM]]] // TILE-234: %[[szK:.*]] = affine.min #[[$bound_map_4]](%[[K]])[%[[ubK]]] -// TILE-234: %[[sAik:.*]] = memref.subview %{{.*}}[%[[I]], %[[K]]] [%[[szM]], %[[szK]]] [1, 1] : memref to memref -// TILE-234: %[[szK:.*]] = affine.min #[[$bound_map_4]](%[[K]])[%[[ubK]]] +// TILE-234: %[[szK_1:.*]] = affine.min #[[$bound_map_4]](%[[K]])[%[[ubK]]] // TILE-234: %[[szN:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[ubN]]] -// TILE-234: %[[sBkj:.*]] = memref.subview %{{.*}}[%[[K]], %[[J]]] [%[[szK]], %[[szN]]] [1, 1] : memref to memref -// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[ubM]]] -// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[ubN]]] -// TILE-234: %[[sCij:.*]] = memref.subview %{{.*}}[%[[I]], %[[J]]] [%[[szM]], %[[szN]]] [1, 1] : memref to memref +// TILE-234: %[[szM_1:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[ubM]]] +// TILE-234: %[[szN_1:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[ubN]]] +// TILE-234: %[[sAik:.*]] = memref.subview %{{.*}}[%[[I]], %[[K]]] [%[[szM]], %[[szK]]] [1, 1] : memref to memref +// TILE-234: %[[sBkj:.*]] = memref.subview %{{.*}}[%[[K]], %[[J]]] [%[[szK_1]], %[[szN]]] [1, 1] : memref to memref +// TILE-234: %[[sCij:.*]] = memref.subview %{{.*}}[%[[I]], %[[J]]] [%[[szM_1]], %[[szN_1]]] [1, 1] : memref to memref // // TILE-234: linalg.matmul ins(%[[sAik]], %[[sBkj]]{{.*}} outs(%[[sCij]] @@ -172,8 +172,8 @@ // TILE-2: scf.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { // TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[M]]] // TILE-2: %[[N:.*]] = memref.dim %{{.*}}, %c1 : memref -// TILE-2: %[[sAi:.*]] = memref.subview %{{.*}}[%[[I]], 0] [%[[szM]], %[[N]]] [1, 1] : memref to memref // TILE-2: %[[szN:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[M]]] +// TILE-2: %[[sAi:.*]] = memref.subview %{{.*}}[%[[I]], 0] [%[[szM]], %[[N]]] [1, 1] : memref to memref // TILE-2: %[[sCi:.*]] = memref.subview %{{.*}}[%[[I]]] [%[[szN]]] [1] : memref to memref // TILE-2: linalg.matvec ins(%[[sAi]], %{{.*}} outs(%[[sCi]] @@ -187,9 +187,9 @@ // TILE-02: scf.for %[[J:.*]] = %{{.*}}{{.*}} to %[[K]] step %{{.*}} { // TILE-02: %[[M:.*]] = memref.dim %{{.*}}, %c0 : memref // TILE-02: %[[szN:.*]] = affine.min #[[$bound_map]](%[[J]])[%[[K]]] +// TILE-02: %[[szN_1:.*]] = affine.min #[[$bound_map]](%[[J]])[%[[K]]] // TILE-02: %[[sAj:.*]] = memref.subview %{{.*}}[0, %[[J]]] [%[[M]], %[[szN]]] [1, 1] : memref to memref -// TILE-02: %[[szN:.*]] = affine.min #[[$bound_map]](%[[J]])[%[[K]]] -// TILE-02: %[[sBj:.*]] = memref.subview %{{.*}}[%[[J]]] [%[[szN]]] [1] : memref to memref +// TILE-02: %[[sBj:.*]] = memref.subview %{{.*}}[%[[J]]] [%[[szN_1]]] [1] : memref to memref // TILE-02: linalg.matvec ins(%[[sAj]], %[[sBj]]{{.*}} outs(%{{.*}} // TILE-002-LABEL: func @matvec( @@ -211,11 +211,11 @@ // TILE-234: scf.for %[[J:.*]] = %{{.*}}{{.*}} to %[[K]] step %{{.*}} { // TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[M]]] // TILE-234: %[[szN:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[K]]] +// TILE-234: %[[szN_1:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[K]]] +// TILE-234: %[[szM_1:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[M]]] // TILE-234: %[[sAij:.*]] = memref.subview %{{.*}}[%[[I]], %[[J]]] [%[[szM]], %[[szN]]] [1, 1] : memref to memref -// TILE-234: %[[szN:.*]] = affine.min #[[$bound_map_3]](%[[J]])[%[[K]]] -// TILE-234: %[[sBj:.*]] = memref.subview %{{.*}}[%[[J]]] [%[[szN]]] [1] : memref to memref -// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[M]]] -// TILE-234: %[[sCi:.*]] = memref.subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref +// TILE-234: %[[sBj:.*]] = memref.subview %{{.*}}[%[[J]]] [%[[szN_1]]] [1] : memref to memref +// TILE-234: %[[sCi:.*]] = memref.subview %{{.*}}[%[[I]]] [%[[szM_1]]] [1] : memref to memref // // TILE-234: linalg.matvec ins(%[[sAij]], %[[sBj]]{{.*}} outs(%[[sCi]] @@ -231,9 +231,9 @@ // TILE-2: %[[M:.*]] = memref.dim %{{.*}}, %c0 : memref // TILE-2: scf.for %[[I:.*]] = %{{.*}}{{.*}} to %[[M]] step %{{.*}} { // TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[M]]] +// TILE-2: %[[szM_1:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[M]]] // TILE-2: %[[sAi:.*]] = memref.subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref -// TILE-2: %[[szM:.*]] = affine.min #[[$bound_map]](%[[I]])[%[[M]]] -// TILE-2: %[[sBi:.*]] = memref.subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref +// TILE-2: %[[sBi:.*]] = memref.subview %{{.*}}[%[[I]]] [%[[szM_1]]] [1] : memref to memref // TILE-2: linalg.dot ins(%[[sAi]], %[[sBi]]{{.*}} outs( // TILE-02-LABEL: func @dot( @@ -248,9 +248,9 @@ // TILE-234: %[[ubK:.*]] = memref.dim %{{.*}}, %c0 : memref // TILE-234: scf.for %[[I:.*]] = %{{.*}} to %[[ubK]] step %{{.*}} { // TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[ubK]]] +// TILE-234: %[[szM_1:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[ubK]]] // TILE-234: %[[sAi:.*]] = memref.subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref -// TILE-234: %[[szM:.*]] = affine.min #[[$bound_map_2]](%[[I]])[%[[ubK]]] -// TILE-234: %[[sBi:.*]] = memref.subview %{{.*}}[%[[I]]] [%[[szM]]] [1] : memref to memref +// TILE-234: %[[sBi:.*]] = memref.subview %{{.*}}[%[[I]]] [%[[szM_1]]] [1] : memref to memref // TILE-234: linalg.dot ins(%[[sAi]], %[[sBi]]{{.*}} outs( func.func @fill_static(%arg0: memref<127x99xf32>, %arg1: f32) {