Skip to content

Commit f4b735c

Browse files
lyuchuny3lyuchuny3
authored
[Conversion] support VarMeanConverter (#19)
1 parent e93fb46 commit f4b735c

3 files changed

Lines changed: 343 additions & 0 deletions

File tree

‎include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp‎

Lines changed: 341 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1399,6 +1399,347 @@ struct ReduceConverter : public OpConversionPattern<triton::ReduceOp> {
13991399
}
14001400
};
14011401

1402+
class VarMeanConverter : public OpConversionPattern<triton::ReduceOp> {
1403+
using OpConversionPattern<triton::ReduceOp>::OpConversionPattern;
1404+
// We're looking for an op that looks like this:
1405+
//
1406+
// %26:3 = "tt.reduce"(%25#1, %25#2, %25#0) <{axis = 1 : i32}> ({
1407+
// ^bb0(%arg6: f32, %arg7: f32, %arg8: f32, %arg9: f32, %arg10: f32,
1408+
// %arg11: f32):
1409+
// %33 = arith.addf %arg7, %arg10 : f32
1410+
// %34 = arith.maxnumf %33, %cst : f32
1411+
// %35 = arith.mulf %arg6, %arg7 : f32
1412+
// %36 = arith.mulf %arg9, %arg10 : f32
1413+
// %37 = arith.addf %35, %36 : f32
1414+
// %38 = arith.divf %37, %34 : f32
1415+
// %39 = arith.mulf %35, %arg6 : f32
1416+
// %40 = arith.addf %arg8, %39 : f32
1417+
// %41 = arith.addf %40, %arg11 : f32
1418+
// %42 = arith.mulf %36, %arg9 : f32
1419+
// %43 = arith.addf %41, %42 : f32
1420+
// %44 = arith.mulf %33, %38 : f32
1421+
// %45 = arith.mulf %44, %38 : f32
1422+
// %46 = arith.subf %43, %45 : f32
1423+
// tt.reduce.return %38, %33, %46 : f32, f32, f32
1424+
// }) : (tensor<8x2048xf32>, tensor<8x2048xf32>, tensor<8x2048xf32>) ->
1425+
// (tensor<8xf32>, tensor<8xf32>, tensor<8xf32>)
1426+
//
1427+
// The above mlir code is lowered from this combinator in triton's
1428+
// standard.py:
1429+
//
1430+
// def welford_func(mean_x, count_x, M_x, mean_y, count_y, M_y):
1431+
// count = count_x + count_y
1432+
// _count = tl.maximum(count, 1)
1433+
// mc_x = mean_x * count_x
1434+
// mc_y = mean_y * count_y
1435+
// mean = (mc_x + mc_y) / _count
1436+
// M = M_x + mc_x * mean_x + M_y + mc_y * mean_y - count * mean * mean
1437+
// return mean, count, M
1438+
1439+
Value getInitTensor(ConversionPatternRewriter &rewriter,
1440+
ArrayRef<int64_t> shape, Value fillValue,
1441+
Location loc) const {
1442+
Value initTensor =
1443+
rewriter.create<tensor::EmptyOp>(loc, shape, fillValue.getType());
1444+
return rewriter
1445+
.create<linalg::FillOp>(loc, ValueRange{fillValue},
1446+
ValueRange{initTensor})
1447+
.result();
1448+
}
1449+
1450+
LogicalResult checkConstFloat(Value value, float val) const {
1451+
if (auto constOp = dyn_cast<arith::ConstantOp>(value.getDefiningOp())) {
1452+
if (auto floatAttr = dyn_cast<FloatAttr>(constOp.getValue())) {
1453+
if (floatAttr.getValueAsDouble() == val) {
1454+
return success();
1455+
}
1456+
}
1457+
}
1458+
return failure();
1459+
}
1460+
1461+
LogicalResult matchVarMeanBody(Value mean_x, Value count_x, Value M_x,
1462+
Value mean_y, Value count_y, Value M_y,
1463+
mlir::Block::iterator &it,
1464+
Operation *block_teminator) const {
1465+
1466+
// %33 = arith.addf %arg7, %arg10 : f32 // count = count_x + count_y
1467+
// %34 = arith.maxnumf %33, %cst : f32 // _count = tl.maximum(count, 1)
1468+
1469+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1470+
auto addOp0 = dyn_cast<arith::AddFOp>(*it++);
1471+
if (addOp0) {
1472+
if (count_x != addOp0.getLhs() || count_y != addOp0.getRhs()) {
1473+
return failure();
1474+
}
1475+
} else {
1476+
return failure();
1477+
}
1478+
1479+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1480+
auto maxOp0 = dyn_cast<arith::MaxNumFOp>(*it++);
1481+
if (maxOp0) {
1482+
if (maxOp0.getLhs() != addOp0) {
1483+
return failure();
1484+
}
1485+
if (failed(checkConstFloat(maxOp0.getRhs(), 1.f))) {
1486+
return failure();
1487+
}
1488+
} else {
1489+
return failure();
1490+
}
1491+
1492+
// %35 = arith.mulf %arg6, %arg7 : f32 //mc_x = mean_x * count_x
1493+
// %36 = arith.mulf %arg9, %arg10 : f32 //mc_y = mean_y * count_y
1494+
1495+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1496+
auto mulOp0 = dyn_cast<arith::MulFOp>(*it++);
1497+
if (mulOp0) {
1498+
if (mean_x != mulOp0.getLhs() || count_x != mulOp0.getRhs()) {
1499+
return failure();
1500+
}
1501+
} else {
1502+
return failure();
1503+
}
1504+
1505+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1506+
auto mulOp1 = dyn_cast<arith::MulFOp>(*it++);
1507+
if (mulOp1) {
1508+
if (mean_y != mulOp1.getLhs() || count_y != mulOp1.getRhs()) {
1509+
return failure();
1510+
}
1511+
} else {
1512+
return failure();
1513+
}
1514+
1515+
// mean = (mc_x + mc_y) / _count
1516+
//
1517+
// %37 = arith.addf %35, %36 : f32 // sum_mc = mc_x + mc_y
1518+
// %38 = arith.divf %37, %34 : f32 // mean = sum_mc / _count
1519+
1520+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1521+
auto addOp1 = dyn_cast<arith::AddFOp>(*it++);
1522+
if (addOp1) {
1523+
if (addOp1.getLhs() != mulOp0 || addOp1.getRhs() != mulOp1) {
1524+
return failure();
1525+
}
1526+
} else {
1527+
return failure();
1528+
}
1529+
1530+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1531+
auto divOp0 = dyn_cast<arith::DivFOp>(*it++);
1532+
if (divOp0) {
1533+
if (divOp0.getLhs() != addOp1 || divOp0.getRhs() != maxOp0) {
1534+
return failure();
1535+
}
1536+
} else {
1537+
return failure();
1538+
}
1539+
1540+
// M = M_x + mc_x * mean_x + M_y + mc_y * mean_y - count * mean * mean
1541+
//
1542+
// %39 = arith.mulf %35, %arg6 : f32 // part_x = mc_x * mean_x
1543+
// %40 = arith.addf %arg8, %39 : f32 // item_1 = M_x + part_x
1544+
// %41 = arith.addf %40, %arg11 : f32 // item_2 = item_1 + M_y
1545+
// %42 = arith.mulf %36, %arg9 : f32 // part_y = mc_y * mean_y
1546+
// %43 = arith.addf %41, %42 : f32 // item_3 = iterm_2 + part_y
1547+
// %44 = arith.mulf %33, %38 : f32 // mean_0 = count * mean
1548+
// %45 = arith.mulf %44, %38 : f32 // mean_1 = mean_0 * mean
1549+
// %46 = arith.subf %43, %45 : f32 // M = item_3 - mean_1
1550+
1551+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1552+
auto mulOp2 = dyn_cast<arith::MulFOp>(*it++);
1553+
if (mulOp2) {
1554+
if (mulOp2.getLhs() != mulOp0 || mulOp2.getRhs() != mean_x) {
1555+
return failure();
1556+
}
1557+
} else {
1558+
return failure();
1559+
}
1560+
1561+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1562+
auto addOp2 = dyn_cast<arith::AddFOp>(*it++);
1563+
if (addOp2) {
1564+
if (addOp2.getLhs() != M_x || addOp2.getRhs() != mulOp2) {
1565+
return failure();
1566+
}
1567+
} else {
1568+
return failure();
1569+
}
1570+
1571+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1572+
auto addOp3 = dyn_cast<arith::AddFOp>(*it++);
1573+
if (addOp3) {
1574+
if (addOp3.getLhs() != addOp2 || addOp3.getRhs() != M_y) {
1575+
return failure();
1576+
}
1577+
} else {
1578+
return failure();
1579+
}
1580+
1581+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1582+
auto mulOp3 = dyn_cast<arith::MulFOp>(*it++);
1583+
if (mulOp3) {
1584+
if (mulOp3.getLhs() != mulOp1 || mulOp3.getRhs() != mean_y) {
1585+
return failure();
1586+
}
1587+
} else {
1588+
return failure();
1589+
}
1590+
1591+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1592+
auto addOp4 = dyn_cast<arith::AddFOp>(*it++);
1593+
if (addOp4) {
1594+
if (addOp4.getLhs() != addOp3 || addOp4.getRhs() != mulOp3) {
1595+
return failure();
1596+
}
1597+
} else {
1598+
return failure();
1599+
}
1600+
1601+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1602+
auto mulOp4 = dyn_cast<arith::MulFOp>(*it++);
1603+
if (mulOp4) {
1604+
if (mulOp4.getLhs() != addOp0 || mulOp4.getRhs() != divOp0) {
1605+
return failure();
1606+
}
1607+
} else {
1608+
return failure();
1609+
}
1610+
1611+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1612+
auto mulOp5 = dyn_cast<arith::MulFOp>(*it++);
1613+
if (mulOp5) {
1614+
if (mulOp5.getLhs() != mulOp4 || mulOp5.getRhs() != divOp0) {
1615+
return failure();
1616+
}
1617+
} else {
1618+
return failure();
1619+
}
1620+
1621+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1622+
auto subOp = dyn_cast<arith::SubFOp>(*it++);
1623+
if (subOp) {
1624+
if (subOp.getLhs() != addOp4 || subOp.getRhs() != mulOp5) {
1625+
return failure();
1626+
}
1627+
} else {
1628+
return failure();
1629+
}
1630+
1631+
// tt.reduce.return %38, %33, %46 : f32, f32, f32 //return mean, count, M
1632+
1633+
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n");
1634+
auto termOp = dyn_cast<triton::ReduceReturnOp>(*it++);
1635+
if (termOp && termOp == block_teminator) {
1636+
auto opnds = termOp.getOperands();
1637+
if (opnds != ArrayRef<Value>{divOp0, addOp0, subOp}) {
1638+
return failure();
1639+
}
1640+
} else {
1641+
return failure();
1642+
}
1643+
1644+
return success();
1645+
}
1646+
1647+
public:
1648+
VarMeanConverter(MLIRContext *context) : OpConversionPattern(context) {}
1649+
1650+
LogicalResult
1651+
matchAndRewrite(ReduceOp op, OpAdaptor adaptor,
1652+
ConversionPatternRewriter &rewriter) const override final {
1653+
// check num_args = 6 : mean_x, count_x, M_x, mean_y, count_y, M_y
1654+
if (op.getBody()->getNumArguments() != 6) {
1655+
return failure();
1656+
}
1657+
1658+
auto block = op.getBody();
1659+
auto ops = block->without_terminator();
1660+
1661+
Value mean_x = block->getArgument(0);
1662+
Value count_x = block->getArgument(1);
1663+
Value M_x = block->getArgument(2);
1664+
Value mean_y = block->getArgument(3);
1665+
Value count_y = block->getArgument(4);
1666+
Value M_y = block->getArgument(5);
1667+
1668+
auto opsIt = ops.begin();
1669+
if (failed(matchVarMeanBody(mean_x, count_x, M_x, mean_y, count_y, M_y,
1670+
opsIt, block->getTerminator()))) {
1671+
return failure();
1672+
}
1673+
auto loc = op.getLoc();
1674+
auto elemTypes = op.getElementTypes();
1675+
1676+
auto meanType = elemTypes[0];
1677+
auto countType = elemTypes[1];
1678+
auto MType = elemTypes[2];
1679+
1680+
Value zeroMean = rewriter.create<arith::ConstantOp>(
1681+
loc, meanType, rewriter.getFloatAttr(meanType, 0.f));
1682+
Value zeroCount = rewriter.create<arith::ConstantOp>(
1683+
loc, countType, rewriter.getFloatAttr(countType, 0.f));
1684+
Value zeroM = rewriter.create<arith::ConstantOp>(
1685+
loc, MType, rewriter.getFloatAttr(MType, 0.f));
1686+
1687+
auto valueResultType = dyn_cast<RankedTensorType>(op.getType(0));
1688+
const auto isScalarReduce = valueResultType == nullptr;
1689+
SmallVector<int64_t> reductionResultShape{
1690+
isScalarReduce ? SmallVector<int64_t>{}
1691+
: SmallVector<int64_t>(valueResultType.getShape())};
1692+
1693+
auto initTensorMean =
1694+
getInitTensor(rewriter, reductionResultShape, zeroMean, loc);
1695+
auto initTensorCount =
1696+
getInitTensor(rewriter, reductionResultShape, zeroCount, loc);
1697+
auto initTensorM =
1698+
getInitTensor(rewriter, reductionResultShape, zeroM, loc);
1699+
1700+
SmallVector<Value> outputs = {initTensorMean, initTensorCount, initTensorM};
1701+
1702+
auto linalgOp = rewriter.create<linalg::ReduceOp>(
1703+
loc, adaptor.getOperands(), outputs,
1704+
SmallVector<int64_t>{adaptor.getAxis()},
1705+
[&](OpBuilder &b, Location loc, ValueRange inputs) {
1706+
assert(inputs.size() == 6 &&
1707+
"Expected 6 inputs to varmean reduce block");
1708+
1709+
auto tritonReduceBlock = op.getBody();
1710+
IRMapping mapping;
1711+
mapping.map(tritonReduceBlock->getArguments(), inputs);
1712+
1713+
for (auto &op : tritonReduceBlock->without_terminator()) {
1714+
b.clone(op, mapping);
1715+
}
1716+
1717+
auto tritonYield = tritonReduceBlock->getTerminator();
1718+
auto results =
1719+
llvm::map_to_vector(tritonYield->getOperands(), [&](Value val) {
1720+
return mapping.lookup(val);
1721+
});
1722+
b.create<linalg::YieldOp>(loc, results);
1723+
});
1724+
1725+
if (isScalarReduce) {
1726+
SmallVector<Value> reduceResults{
1727+
rewriter.create<tensor::ExtractOp>(
1728+
loc, meanType, linalgOp.getResults()[0], ValueRange{}),
1729+
rewriter.create<tensor::ExtractOp>(
1730+
loc, countType, linalgOp.getResults()[1], ValueRange{}),
1731+
rewriter.create<tensor::ExtractOp>(
1732+
loc, MType, linalgOp.getResults()[2], ValueRange{}),
1733+
};
1734+
rewriter.replaceOp(op, reduceResults);
1735+
} else {
1736+
rewriter.replaceOp(op, linalgOp);
1737+
}
1738+
1739+
return success();
1740+
}
1741+
};
1742+
14021743
template <typename T>
14031744
class ArgMinMaxBaseConverter : public OpConversionPattern<triton::ReduceOp> {
14041745
using OpConversionPattern<triton::ReduceOp>::OpConversionPattern;

‎lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ void mlir::triton::populateTritonArithToLinalgConversionPatterns(
9494
// the first elements along the reduction axis and perform the reduction on
9595
// the remaining elements. However, this results in creatings sub-tensors that
9696
// aren't always multiple of 2s, which are sub-optimal for certain hardwares.
97+
patterns.add<VarMeanConverter>(patterns.getContext());
9798
patterns.add<ArgMinConverter>(patterns.getContext());
9899
patterns.add<ArgMaxConverter>(patterns.getContext());
99100
patterns.add<ReduceConverter>(patterns.getContext());

‎lib/Conversion/TritonToLinalg/TritonToLinalg.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ void mlir::triton::populateTritonToLinalgConversionPatterns(
8080
// the first elements along the reduction axis and perform the reduction on
8181
// the remaining elements. However, this results in creatings sub-tensors that
8282
// aren't always multiple of 2s, which are sub-optimal for certain hardwares.
83+
patterns.add<VarMeanConverter>(patterns.getContext());
8384
patterns.add<ArgMinConverter>(patterns.getContext());
8485
patterns.add<ArgMaxConverter>(patterns.getContext());
8586
patterns.add<ReduceConverter>(patterns.getContext());

0 commit comments

Comments
 (0)