Skip to content

Commit

Permalink
[CPU] Improve distribution tile sizes selection for mmt4d ops. (iree-…
Browse files Browse the repository at this point in the history 
…org#15448)

Perfect distribution is a really hard problem, especially for the whole
program/model. As a general rule workgroups like these should be
processing thousands to tens of thousands of operations in order to
amortize the scheduling overheads. Ideally each workgroup would roughly
take 100us-500us; less than that and overhead starts to dominate and
more than that needs to be balanced with total count: ~200 workgroups is
more than enough to hide latencies and more than that adds overhead with
diminishing returns.

The basic heuristic we want is to spend enough time inside each
workgroup to justify the overhead involved in launching it but not so
much time and so few total that variance in the system can't be hidden
(4 x 400ms workgroups on 3 threads will always have 400ms of work happen
on one thread with the other 2 idle, or 4 x 400ms on 4 threads where 1
thread gets delayed 50ms will cause the total latency to increase 50ms,
etc).

In the context, `maxTileSizes` is critical because it is one of factors
about number of workgroups. In IREE CPU, we've been using
`defaultDistTileSize` to model tile sizes. It is used by matmul and
other generic ops. In `mmt4d` semantics, some data are already tiled and
moved to inner loops. To model distribution with the same factor, we
should scale it down with corresponding inner tile sizes. Then the
distribution of all the dispatches is modeled by the same factor.

Fixes iree-org#15391
  • Loading branch information
@hanhanW
hanhanW authored Nov 9, 2023
1 parent 7be3992 commit 3ef0ea1
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 29 deletions.
54 changes: 27 additions & 27 deletions compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1216,6 +1216,13 @@ setRootConfig(func::FuncOp entryPointFn,
static LogicalResult setRootConfig(func::FuncOp entryPointFn,
linalg::Mmt4DOp mmt4dOp) {
assert(!getLoweringConfig(mmt4dOp) && "expected lowering_config is not set");

auto lhsShape = cast<ShapedType>(mmt4dOp.getInputs()[0].getType()).getShape();
auto rhsShape = cast<ShapedType>(mmt4dOp.getInputs()[1].getType()).getShape();
int M0 = lhsShape[2];
int N0 = rhsShape[2];
int K0 = lhsShape[3];

auto getDistTileSizes = [&]() -> SmallVector<int64_t> {
if (!clMmt4dDistributionTileSizes.empty()) {
return SmallVector<int64_t>(clMmt4dDistributionTileSizes.begin(),
Expand All @@ -1224,10 +1231,12 @@ static LogicalResult setRootConfig(func::FuncOp entryPointFn,
unsigned numLoops = mmt4dOp.getNumLoops();
SmallVector<int64_t> minTileSizes(numLoops, 0);
SmallVector<int64_t> maxTileSizes(numLoops, 0);
minTileSizes[0] = 4;
minTileSizes[1] = 4;
maxTileSizes[0] = 48;
maxTileSizes[1] = 32;
minTileSizes[0] = minTileSizes[1] = 1;
// Scale default distribution tile size down because it is already in packed
// domain. With outer M dim size=X means that there will be `X * M0`
// elements to process. Same for N dimension.
maxTileSizes[0] = llvm::divideCeil(defaultDistTileSize, M0);
maxTileSizes[1] = llvm::divideCeil(defaultDistTileSize, N0);
SmallVector<int64_t> distTileSizes = getDefaultDistributedLevelTileSizes(
mmt4dOp, minTileSizes, maxTileSizes);
return distTileSizes;
Expand All @@ -1238,14 +1247,6 @@ static LogicalResult setRootConfig(func::FuncOp entryPointFn,
return SmallVector<int64_t>(mmt4dL1TileSizes.begin(),
mmt4dL1TileSizes.end());
}
auto lhsShape =
llvm::cast<ShapedType>(mmt4dOp.getInputs()[0].getType()).getShape();
auto rhsShape =
llvm::cast<ShapedType>(mmt4dOp.getInputs()[1].getType()).getShape();
int M0 = lhsShape[2];
int N0 = rhsShape[2];
int K0 = lhsShape[3];

return {1, 1, 1, M0, N0, K0};
};

Expand All @@ -1272,6 +1273,13 @@ static LogicalResult setRootConfig(func::FuncOp entryPointFn,
linalg::BatchMmt4DOp batchMmt4dOp) {
assert(!getLoweringConfig(batchMmt4dOp) &&
"expected lowering_config is not set");
auto lhsShape =
cast<ShapedType>(batchMmt4dOp.getInputs()[0].getType()).getShape();
auto rhsShape =
cast<ShapedType>(batchMmt4dOp.getInputs()[1].getType()).getShape();
int M0 = lhsShape[3];
int N0 = rhsShape[3];
int K0 = lhsShape[4];
auto getDistTileSizes = [&]() -> SmallVector<int64_t> {
if (!clMmt4dDistributionTileSizes.empty()) {
SmallVector<int64_t> tileSizes;
Expand All @@ -1285,12 +1293,13 @@ static LogicalResult setRootConfig(func::FuncOp entryPointFn,
unsigned numLoops = batchMmt4dOp.getNumLoops();
SmallVector<int64_t> minTileSizes(numLoops, 0);
SmallVector<int64_t> maxTileSizes(numLoops, 0);
minTileSizes[0] = 1;
minTileSizes[1] = 4;
minTileSizes[2] = 4;
maxTileSizes[0] = 1;
maxTileSizes[1] = 48;
maxTileSizes[2] = 32;
minTileSizes[0] = maxTileSizes[0] = 1; // Force batch dim being 1.
minTileSizes[1] = minTileSizes[2] = 1;
// Scale default distribution tile size down because it is already in packed
// domain. With outer M dim size=X means that there will be `X * M0`
// elements to process. Same for N dimension.
maxTileSizes[1] = llvm::divideCeil(defaultDistTileSize, M0);
maxTileSizes[2] = llvm::divideCeil(defaultDistTileSize, N0);
SmallVector<int64_t> distTileSizes = getDefaultDistributedLevelTileSizes(
batchMmt4dOp, minTileSizes, maxTileSizes);
return distTileSizes;
Expand All @@ -1307,15 +1316,6 @@ static LogicalResult setRootConfig(func::FuncOp entryPointFn,
return tileSizes;
}

auto lhsShape =
llvm::cast<ShapedType>(batchMmt4dOp.getInputs()[0].getType())
.getShape();
auto rhsShape =
llvm::cast<ShapedType>(batchMmt4dOp.getInputs()[1].getType())
.getShape();
int M0 = lhsShape[3];
int N0 = rhsShape[3];
int K0 = lhsShape[4];
tileSizes.append({1, 1, 1, M0, N0, K0});
return tileSizes;
};
Expand Down
4 changes: 2 additions & 2 deletions compiler/src/iree/compiler/Codegen/LLVMCPU/test/test_config_mmt4d.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ hal.executable private @mmt4d_384x384x512_4x1x4_dispatch_0 {
}
}

// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[32, 32, 0, 0, 0, 0], [1, 1, 0, 4, 4, 0], [0, 0, 1, 0, 0, 1]{{\]}}
// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[16, 16, 0, 0, 0, 0], [1, 1, 0, 4, 4, 0], [0, 0, 1, 0, 0, 1]{{\]}}
// CHECK: func.func @mmt4d_384x384x512_4x1x4_dispatch_0()
// CHECK: linalg.mmt4d
// CHECK-SAME: lowering_config = #[[CONFIG]]
Expand Down Expand Up @@ -90,7 +90,7 @@ hal.executable private @batch_mmt4d {
}
}

// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 10, 20, 0, 0, 0, 0], [1, 1, 1, 0, 8, 4, 0], [0, 0, 0, 1, 0, 0, 1]{{\]}}>
// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 5, 16, 0, 0, 0, 0], [1, 1, 1, 0, 8, 4, 0], [0, 0, 0, 1, 0, 0, 1]{{\]}}>
// CHECK: func.func @batch_mmt4d()
// CHECK: linalg.batch_mmt4d
// CHECK-SAME: lowering_config = #[[CONFIG]]

0 comments on commit 3ef0ea1

Please sign in to comment.