Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Brgemm register tiling for bf16 type #1005

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Brgemm register tiling for bf16 type #1005

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
7943e1e
code re-factoring
Feb 1, 2025
310a21a
code re-factoring and test cases for bf16
Feb 3, 2025
7653434
code re-factoring - adding comments, space etc.,.
Feb 3, 2025
09d9ae7
Added a integration test-case for bf16 type
Feb 3, 2025
ef08d66
Added extra validation, made m,n,k tile size as mandatory
Feb 13, 2025
974c3eb
Code re-factoring, changes to test-cases
Feb 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
222 changes: 109 additions & 113 deletions lib/TPP/Transforms/BrgemmLinalgTiling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
//===- BrgemmLinalgTiling.cpp -----------------------------------------*- C++-*-===//
//===- BrgemmLinalgTiling.cpp -----------------------------------------*-
//C++-*-===//
arun-thmn marked this conversation as resolved.
Show resolved Hide resolved
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
Expand Down Expand Up @@ -43,160 +44,152 @@ using namespace mlir::tpp;

namespace mlir {
namespace tpp {
struct LinalgOpTiling : OpRewritePattern<linalg::BatchReduceMatmulOp> {
using OpRewritePattern<linalg::BatchReduceMatmulOp>::OpRewritePattern;

template <typename BrgemmOp>
struct LinalgOpTiling : OpRewritePattern<BrgemmOp> {
using OpRewritePattern<BrgemmOp>::OpRewritePattern;

LinalgOpTiling(MLIRContext *ctx, BrgemmLinalgTilingOptions tilingoptions)
: OpRewritePattern(ctx), options(tilingoptions) {}
: OpRewritePattern<BrgemmOp>(ctx), options(tilingoptions) {}

LogicalResult matchAndRewrite(linalg::BatchReduceMatmulOp brgemmOp,
LogicalResult matchAndRewrite(BrgemmOp brgemmOp,
PatternRewriter &rewriter) const override {

if (!brgemmOp.hasPureBufferSemantics())
return failure();
// Get the register blocking tile shape from the user input
SmallVector<int64_t> tileShapeM(options.registerTileShape.begin(),
options.registerTileShape.end());

if (tileShapeM.size() != 2)
return failure();
// Check whether the tile sizes are valid
if (options.registerTileShape.size() != 3 &&
options.registerTileShape.size() != 2)
return failure();
arun-thmn marked this conversation as resolved.
Show resolved Hide resolved

SmallVector<int64_t> tileShapeN(2);
tileShapeN[0] = 1;
tileShapeN[1] = tileShapeM[1];
tileShapeM[1] = 1;
// Check the whether the operation is brmatmul fp32 or bf16 type using
// reduction count
SmallVector<utils::IteratorType> brgemmIteratorTypes =
brgemmOp.getIteratorTypesArray();
int reductionCount =
std::count(brgemmIteratorTypes.begin(), brgemmIteratorTypes.end(),
utils::IteratorType::reduction);
if (reductionCount != 2 && reductionCount != 3)
return failure();

// Stores the M, N, and K Tile Sizes
// Get the register blocking tile shape from the user input
SmallVector<int64_t> mxnxkTile(3);
// Stores the M, and N Tile Sizes
SmallVector<int64_t> mxnTile(2);
for (size_t i = 0; i < options.registerTileShape.size(); i++) {
mxnxkTile[i] = options.registerTileShape[i];
}

// Set the K tile to 1, if the user not provided (it is fp32 target)
if (options.registerTileShape.size() == 2)
mxnxkTile[2] = 1;
arun-thmn marked this conversation as resolved.
Show resolved Hide resolved

// k-tile size adjusted based on the vnni layout for bf16 type
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This has baked-in assumptions that are not verified.
As the pass now operates on generic, we need to strictly filter ops that are accepted. I think you need to at least ensure it is a VNNI contraction first - there should be some suitable helpers in VnniUtils.
If f32 generic should be supported as well, it might need some extra checks there too.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing out this.
Added more checks for f32 and bf16 type. Used a check from vnniutils.

auto tensorShape =
dyn_cast<MemRefType>(brgemmOp.getOperand(0).getType()).getShape();
if (tensorShape.size() == 4 && options.registerTileShape.size() == 3) {
mxnxkTile[2] = mxnxkTile[2] / tensorShape[3];
}

mxnxkTile[0] = tileShapeM[0];
mxnxkTile[1] = tileShapeN[1];
mxnxkTile[2] = tileShapeM[1];
mxnTile[0] = tileShapeM[0];
mxnTile[1] = tileShapeN[1];

// To assist in calculating the argument and step values for the tiled loop.
SmallVector<int64_t> boundariesOne{1,
static_cast<long>(tileShapeM.size() - 1),
static_cast<long>(mxnxkTile.size() - 1)};

SmallVector<int64_t> tileSizesIndex{static_cast<long>(tileShapeM.size()),
static_cast<long>(tileShapeN.size()),
static_cast<long>(mxnTile.size())};
SmallVector<SmallVector<int64_t>> tileshapes{tileShapeM, tileShapeN, mxnTile};
SmallVector<int> swap_i = {0, 2, 1};
size_t i = 0;
SmallVector<int> swap_i = {0, 2, 1};
std::map<int, std::map<int, Value>> inductionVars;

// For M, N, and K loops
scf::ForOp innermostForLoop;
// For brgemm reduction loop
scf::ForOp reductionForLoop;

// Creating the tiled loops
for (auto itrShapeM = mxnxkTile.begin(); itrShapeM != mxnxkTile.end();
itrShapeM++, i++) {
int index = swap_i[i] / boundariesOne[swap_i[i]];
int offset = swap_i[i] / (mxnxkTile.size() - 1);

int operandSize =
dyn_cast<MemRefType>(brgemmOp.getOperand(index).getType())
.getShape()
.size();
int effectiveOffset = operandSize - tileSizesIndex[index] + offset;
for (auto itrShapeMNK = mxnxkTile.begin(); itrShapeMNK != mxnxkTile.end();
itrShapeMNK++, i++) {
arun-thmn marked this conversation as resolved.
Show resolved Hide resolved
auto upperBound =
dyn_cast<MemRefType>(brgemmOp.getOperand(index).getType())
.getShape()[effectiveOffset];
dyn_cast<MemRefType>(brgemmOp.getOperand(swap_i[i]).getType())
.getShape()[1];
// Tile size should not be greater than the upperBound
if ((*itrShapeMNK) > upperBound)
return failure();

Location loc = brgemmOp.getLoc();
Value zeroCst = rewriter.create<arith::ConstantIndexOp>(loc, 0);
Value ubCstTiledLoop = rewriter.create<arith::ConstantIndexOp>(loc, upperBound);
//Tile size should not be greater than the upperBound
if ((*itrShapeM) > upperBound)
return failure();
Value stepCstTiledLoop = rewriter.create<arith::ConstantIndexOp>(loc, *itrShapeM);
Value ubCstTiledLoop =
rewriter.create<arith::ConstantIndexOp>(loc, upperBound);
Value stepCstTiledLoop =
rewriter.create<arith::ConstantIndexOp>(loc, *itrShapeMNK);
// Creates M, N, and K tile loops
scf::ForOp loopOp = rewriter.create<scf::ForOp>(brgemmOp.getLoc(),
zeroCst, ubCstTiledLoop, stepCstTiledLoop);
scf::ForOp loopOp = rewriter.create<scf::ForOp>(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I am understanding right, this transform is meant to operate on linalg ops. As I expect all the ops you want to support will implement TilingInterface, would it be possible to just use the TileUsingFor transform instead of manually implementing tiling?

Copy link
Contributor Author

@arun-thmn arun-thmn Feb 13, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, it a target for linalg ops. But, now we are focusing only on batch reduce matmul. As far as I know, the plan is to make tiling, hoisting, vector-to-fma, and vector-to-amx as a one register blocking transform schedule in the upstream.
Also, we extended this pass to support bf16 type, so not taught of using TileUsingFor transform schedule.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not one register blocking transform, but one schedule of multiple transforms that implement register blocking at the vector level. These should still be separate transforms and upstreamed not as "vector tiling" transforms, but "register blocking on vector" transforms.

brgemmOp.getLoc(), zeroCst, ubCstTiledLoop, stepCstTiledLoop);
rewriter.setInsertionPointToStart(loopOp.getBody());
int indexTwo = offset;
int operandSizeTwo =
dyn_cast<MemRefType>(brgemmOp.getOperand(indexTwo).getType())
.getShape()
.size();
int effectiveOffsetTwo = operandSizeTwo - tileSizesIndex[index] + index;

inductionVars[index][effectiveOffset] = loopOp.getInductionVar();

inductionVars[indexTwo][effectiveOffsetTwo] = loopOp.getInductionVar();
int indexThree = mxnTile.size();
int effectiveOffsetThree =
index +
dyn_cast<MemRefType>(brgemmOp.getOperand(indexThree).getType())
.getShape()
.size() -
tileSizesIndex[indexThree];
if (inductionVars[indexThree][effectiveOffsetThree] == NULL) {
inductionVars[indexThree][effectiveOffsetThree] =
loopOp.getInductionVar();
}

innermostForLoop = loopOp;
if ((mxnxkTile.size() - 1) == (i + 1)) {
//Creates the brgemm reduction loop

// Stores the induction variable with respect to the operands mapping it's
// subview.
if (i == 0) { // Stores iv for M loop
inductionVars[0][1] = loopOp.getInductionVar();
inductionVars[2][0] = loopOp.getInductionVar();
} else if (i == 1) { //stores iv for N loop, creates batch loop, and maps iv of batch loop
inductionVars[1][2] = loopOp.getInductionVar();
inductionVars[2][1] = loopOp.getInductionVar();
// Creates reduction loop after the N loop
Value ubCstReduction = rewriter.create<arith::ConstantIndexOp>(
loc, dyn_cast<MemRefType>(brgemmOp.getOperand(0).getType())
.getShape()[0]);
Value stepCstReduction = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value stepCstReduction =
rewriter.create<arith::ConstantIndexOp>(loc, 1);
scf::ForOp redloopOp = rewriter.create<scf::ForOp>(
brgemmOp.getLoc(), zeroCst, ubCstReduction, stepCstReduction);
rewriter.setInsertionPointToStart(redloopOp.getBody());
reductionForLoop = redloopOp;
inductionVars[0][0] = redloopOp.getInductionVar();
inductionVars[1][0] = redloopOp.getInductionVar();

} else if (i == 2) { // stores iv for k-loop
inductionVars[0][2] = loopOp.getInductionVar();
inductionVars[1][1] = loopOp.getInductionVar();
}
}

// DS to assist while creating new subviews with correct indices and shapes
SmallVector<int64_t> mxkTile(2);
arun-thmn marked this conversation as resolved.
Show resolved Hide resolved
SmallVector<int64_t> kxnTile(2);
SmallVector<int64_t> mxnTile(2);
mxkTile[0] = mxnxkTile[0];
mxkTile[1] = mxnxkTile[2];
kxnTile[0] = mxnxkTile[2];
kxnTile[1] = mxnxkTile[1];
mxnTile[0] = mxnxkTile[0];
mxnTile[1] = mxnxkTile[1];

SmallVector<SmallVector<int64_t>> tileshapes{mxkTile, kxnTile, mxnTile};
// Creating subviews
SmallVector<SmallVector<int64_t>> tiles = {tileShapeM, tileShapeN};
for (size_t i = 0; i < brgemmOp.getNumOperands(); i++) {
SmallVector<int64_t> indices;
auto input = brgemmOp.getOperand(i);
auto operandType = input.getType();
SmallVector<OpFoldResult> offsets;
size_t k = 0;
auto tileItr = tileshapes[i].begin();
auto tensorShape = dyn_cast<MemRefType>(operandType).getShape();
SmallVector<int64_t> indices;
SmallVector<OpFoldResult> shape;
SmallVector<OpFoldResult> strides;

auto input = brgemmOp.getOperand(i);
auto tensorShape = dyn_cast<MemRefType>(input.getType()).getShape();
auto tileItr = tileshapes[i].begin();

// Iterates over the shape of each tensor and update its offsets, indices,
// shapes, strides with respect to tile sizes
for (size_t j = 0; j < tensorShape.size(); j++) {
if (j < tensorShape.size() - tileSizesIndex[i]) {
if (j == ((tensorShape.size() - tileSizesIndex[i]) - 1) &&
i < (brgemmOp.getNumOperands() - 1)) {
offsets.push_back(reductionForLoop.getInductionVar());
indices.push_back(tensorShape[j] / tensorShape[j]);
shape.push_back(rewriter.getIndexAttr(tensorShape[j] / tensorShape[j]));
strides.push_back(rewriter.getIndexAttr(1));

} else {
offsets.push_back(rewriter.getIndexAttr(0));
indices.push_back(tensorShape[j]);
shape.push_back(rewriter.getIndexAttr(tensorShape[j]));
strides.push_back(rewriter.getIndexAttr(1));
}
} else {
shape.push_back(rewriter.getIndexAttr(*tileItr));
if (j == 0 && (i < 2)) { // Updates the batch dimension
offsets.push_back(inductionVars[i][j]);
indices.push_back(1);
shape.push_back(rewriter.getIndexAttr(1));
strides.push_back(rewriter.getIndexAttr(1));
} else if (j < 3) { // Updates the M, N, and K dimensions
offsets.push_back(inductionVars[i][j]);
indices.push_back((*tileItr));
shape.push_back(rewriter.getIndexAttr(*tileItr));
strides.push_back(rewriter.getIndexAttr(1));
offsets.push_back(
inductionVars[i][tensorShape.size() - tileSizesIndex[i] + k]);
k++;
tileItr++;
} else { // Just copies the vnni layout dimensions
offsets.push_back(rewriter.getIndexAttr(0));
indices.push_back(tensorShape[j]);
shape.push_back(rewriter.getIndexAttr(tensorShape[j]));
strides.push_back(rewriter.getIndexAttr(1));
}
}

auto subview = rewriter.create<memref::SubViewOp>(
brgemmOp.getLoc(), MemRefType(),
input, offsets, shape, strides);
brgemmOp.getLoc(), MemRefType(), input, offsets, shape, strides);
arun-thmn marked this conversation as resolved.
Show resolved Hide resolved
brgemmOp.setOperand(i, subview);
}

Expand All @@ -214,11 +207,14 @@ struct LinalgOpTiling : OpRewritePattern<linalg::BatchReduceMatmulOp> {
};

void populateBrgemmLinalgTilingPatterns(RewritePatternSet &patterns,
BrgemmLinalgTilingOptions options) {
patterns.add<LinalgOpTiling>(patterns.getContext(), options);
BrgemmLinalgTilingOptions options) {
patterns.add<LinalgOpTiling<linalg::GenericOp>,
LinalgOpTiling<linalg::BatchReduceMatmulOp>>(
patterns.getContext(), options);
}

struct BrgemmLinalgTiling : public tpp::impl::BrgemmLinalgTilingBase<BrgemmLinalgTiling> {
struct BrgemmLinalgTiling
: public tpp::impl::BrgemmLinalgTilingBase<BrgemmLinalgTiling> {

using BrgemmLinalgTilingBase::BrgemmLinalgTilingBase;

Expand Down
30 changes: 30 additions & 0 deletions test/Integration/tile-brgemm-linalg-matmul-bf16.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
// RUN: tpp-run -e register_tile_bf16 --entry-point-result=void -print %s > %t.1
// RUN: tpp-opt %s --tile-brgemm-linalg="registerBlocking=32,32,32" -convert-linalg-to-xsmm | tpp-run -e register_tile_bf16 --entry-point-result=void -print > %t.2
// RUN: diff %t.1 %t.2
arun-thmn marked this conversation as resolved.
Show resolved Hide resolved
// RUN: rm %t.1 %t.2
arun-thmn marked this conversation as resolved.
Show resolved Hide resolved

#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d4, d1)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3, d1)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)>
module {
memref.global "private" constant @__constant_32x16x32x2xbf16 : memref<32x16x32x2xbf16> = dense<1.000000e+00> {alignment = 64 : i64}
func.func @register_tile_bf16(%arg0: memref<8x32x32x32xbf16>) -> memref<8x32x32x32xbf16> {
%cst = arith.constant 0.000000e+00 : bf16
%0 = memref.get_global @__constant_32x16x32x2xbf16 : memref<32x16x32x2xbf16>
%alloc = memref.alloc() {alignment = 64 : i64} : memref<8x32x32x32xbf16>
%expand_shape = memref.expand_shape %arg0 [[0], [1], [2], [3, 4]] output_shape [8, 32, 32, 16, 2] : memref<8x32x32x32xbf16> into memref<8x32x32x16x2xbf16>
scf.forall (%arg1, %arg2) in (8, 32) {
%subview = memref.subview %alloc[%arg1, %arg2, 0, 0] [1, 1, 32, 32] [1, 1, 1, 1] : memref<8x32x32x32xbf16> to memref<32x32xbf16, strided<[32, 1], offset: ?>>
linalg.fill ins(%cst : bf16) outs(%subview : memref<32x32xbf16, strided<[32, 1], offset: ?>>)
%subview_0 = memref.subview %expand_shape[%arg1, 0, 0, 0, 0] [1, 32, 32, 16, 2] [1, 1, 1, 1, 1] : memref<8x32x32x16x2xbf16> to memref<32x32x16x2xbf16, strided<[1024, 32, 2, 1], offset: ?>>
linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["reduction", "reduction", "parallel", "parallel", "reduction"]} ins(%subview_0, %0 : memref<32x32x16x2xbf16, strided<[1024, 32, 2, 1], offset: ?>>, memref<32x16x32x2xbf16>) outs(%subview : memref<32x32xbf16, strided<[32, 1], offset: ?>>) {
^bb0(%in: bf16, %in_1: bf16, %out: bf16):
%1 = arith.mulf %in, %in_1 : bf16
%2 = arith.addf %out, %1 : bf16
linalg.yield %2 : bf16
}
}
return %alloc : memref<8x32x32x32xbf16>
}
}

Loading