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

[linalgExt::Attention] add AttentionOP boolean useExp2 to choose either exp or exp2 #19820

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

[linalgExt::Attention] add AttentionOP boolean useExp2 to choose either exp or exp2 #19820

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
76a805b
add exp
ita9naiwa Jan 23, 2025
f3ce0d5
lint
ita9naiwa Jan 27, 2025
fb114fa
add Attention TMTensorToLinalgExt Lowering
ita9naiwa Feb 24, 2025
a1f9be1
lint
ita9naiwa Feb 25, 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
4 changes: 4 additions & 0 deletions compiler/plugins/input/Torch/InputConversion/ConvertTMTensorToLinalgExt.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,10 +186,14 @@ struct AttentionOpConversion
map = map.insertResult(rewriter.getAffineDimExpr(batch), batch);
}
}
auto useExp2 = op->getAttr("use_exp2");
if (!useExp2)
useExp2 = rewriter.getBoolAttr(false);

auto attention = rewriter.create<IREE::LinalgExt::AttentionOp>(
loc, result.getType(), query, key, value, scale, result,
rewriter.getAffineMapArrayAttr(indexingMaps), optionalMask);
attention->setAttr("use_exp2", useExp2);

{
auto *block = rewriter.createBlock(&attention.getRegion());
Expand Down
36 changes: 32 additions & 4 deletions compiler/plugins/input/Torch/InputConversion/test/attention.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ func.func @attention(%arg0: tensor<5x2x3x4xf32>, %arg1: tensor<5x2x3x4xf32>, %ar
// CHECK-SAME: %[[ARG3:.*]]: tensor<5x2x3x4xf32>) -> tensor<5x2x3x4xf32> {
// CHECK: %[[SCALE:.*]] = arith.constant 5.000000e-01 : f32
// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<5x2x3x4xf32>
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]]} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<5x2x3x4xf32>, tensor<5x2x3x4xf32>, tensor<5x2x3x4xf32>, f32) outs(%[[EMPTY]] : tensor<5x2x3x4xf32>) {
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]], use_exp2 = false} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<5x2x3x4xf32>, tensor<5x2x3x4xf32>, tensor<5x2x3x4xf32>, f32) outs(%[[EMPTY]] : tensor<5x2x3x4xf32>) {
// CHECK: ^[[BLOCK:.+]](%[[SCORE:.+]]: f32):
// CHECK: linalg_ext.yield %[[SCORE]]
// CHECK: } -> tensor<5x2x3x4xf32>
Expand All @@ -39,7 +39,7 @@ func.func @attention(%arg0: tensor<5x2x8x4xf32>, %arg1: tensor<5x2x3x4xf32>, %ar
// CHECK-SAME: %[[ARG3:.*]]: tensor<5x2x8x4xf32>) -> tensor<5x2x8x4xf32> {
// CHECK: %[[SCALE:.*]] = arith.constant 5.000000e-01 : f32
// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<5x2x8x4xf32>
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]]} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<5x2x8x4xf32>, tensor<5x2x3x4xf32>, tensor<5x2x3x4xf32>, f32) outs(%[[EMPTY]] : tensor<5x2x8x4xf32>) {
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]], use_exp2 = false} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<5x2x8x4xf32>, tensor<5x2x3x4xf32>, tensor<5x2x3x4xf32>, f32) outs(%[[EMPTY]] : tensor<5x2x8x4xf32>) {
// CHECK: ^[[BLOCK:.+]](%[[SCORE:.+]]: f32):
// CHECK: linalg_ext.yield %[[SCORE]]
// CHECK: } -> tensor<5x2x8x4xf32>
Expand All @@ -62,7 +62,7 @@ func.func @attention(%arg0: tensor<1x3x4xf32>, %arg1: tensor<1x3x4xf32>, %arg2:
// CHECK: %[[ARG3:.*]]: tensor<1x3x4xf32>) -> tensor<1x3x4xf32> {
// CHECK: %[[SCALE:.*]] = arith.constant 5.000000e-01 : f32
// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1x3x4xf32>
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]]} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<1x3x4xf32>, tensor<1x3x4xf32>, tensor<1x3x4xf32>, f32) outs(%[[EMPTY]] : tensor<1x3x4xf32>) {
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]], use_exp2 = false} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<1x3x4xf32>, tensor<1x3x4xf32>, tensor<1x3x4xf32>, f32) outs(%[[EMPTY]] : tensor<1x3x4xf32>) {
// CHECK: ^[[BLOCK:.+]](%[[SCORE:.+]]: f32):
// CHECK: linalg_ext.yield %[[SCORE]]
// CHECK: } -> tensor<1x3x4xf32>
Expand All @@ -89,7 +89,35 @@ func.func @attention_dyn(%arg0: tensor<?x?x4xf32>, %arg1: tensor<?x?x4xf32>, %ar
// CHECK-DAG: %[[DIM0:.*]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK-DAG: %[[DIM1:.*]] = tensor.dim %[[ARG0]], %[[C1]]
// CHECK-DAG: %[[EMPTY:.*]] = tensor.empty(%[[DIM0]], %[[DIM1]]) : tensor<?x?x4xf32>
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]]} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<?x?x4xf32>, tensor<?x?x4xf32>, tensor<?x?x4xf32>, f32) outs(%[[EMPTY]] : tensor<?x?x4xf32>) {
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]], use_exp2 = false} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<?x?x4xf32>, tensor<?x?x4xf32>, tensor<?x?x4xf32>, f32) outs(%[[EMPTY]] : tensor<?x?x4xf32>) {
// CHECK: ^[[BLOCK:.+]](%[[SCORE:.+]]: f32):
// CHECK: linalg_ext.yield %[[SCORE]]
// CHECK: } -> tensor<?x?x4xf32>
// CHECK: return %[[ATTN]] : tensor<?x?x4xf32>


// -----
func.func @attention_use_exp2_attr(%arg0: tensor<?x?x4xf32>, %arg1: tensor<?x?x4xf32>, %arg2: tensor<?x?x4xf32>, %arg3: tensor<?x?x4xf32>) -> (tensor<?x?x4xf32>) {
%0 = tm_tensor.attention {use_exp2 = true} ins(%arg0, %arg1, %arg2 : tensor<?x?x4xf32>, tensor<?x?x4xf32>, tensor<?x?x4xf32>) outs(%arg3: tensor<?x?x4xf32>) -> tensor<?x?x4xf32>
return %0 : tensor<?x?x4xf32>
}

// CHECK-DAG: #[[$MAP_Q:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3)>
// CHECK-DAG: #[[$MAP_K:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3)>
// CHECK-DAG: #[[$MAP_V:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d2)>
// CHECK-DAG: #[[$MAP_S:.+]] = affine_map<(d0, d1, d2, d3, d4) -> ()>
// CHECK-DAG: #[[$MAP_O:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>

// CHECK-LABEL: func.func @attention_use_exp2_attr(
// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?x4xf32>, %[[ARG1:.*]]: tensor<?x?x4xf32>, %[[ARG2:.*]]: tensor<?x?x4xf32>,
// CHECK: %arg3: tensor<?x?x4xf32>) -> tensor<?x?x4xf32> {
// CHECK-DAG: %[[SCALE:.*]] = arith.constant 5.000000e-01 : f32
// CHECK-DAG: %[[C0:.*]] = arith.constant 0
// CHECK-DAG: %[[C1:.*]] = arith.constant 1
// CHECK-DAG: %[[DIM0:.*]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK-DAG: %[[DIM1:.*]] = tensor.dim %[[ARG0]], %[[C1]]
// CHECK-DAG: %[[EMPTY:.*]] = tensor.empty(%[[DIM0]], %[[DIM1]]) : tensor<?x?x4xf32>
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]], use_exp2 = true} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<?x?x4xf32>, tensor<?x?x4xf32>, tensor<?x?x4xf32>, f32) outs(%[[EMPTY]] : tensor<?x?x4xf32>) {
// CHECK: ^[[BLOCK:.+]](%[[SCORE:.+]]: f32):
// CHECK: linalg_ext.yield %[[SCORE]]
// CHECK: } -> tensor<?x?x4xf32>
Expand Down
29 changes: 22 additions & 7 deletions compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,9 +256,9 @@ static Value applyMask(OpBuilder &builder, Location loc, AffineMap qkMap,
}

// Compute output = exp2(output - input)
static Value computeSubAndExp2(OpBuilder &builder, Location loc,
AffineMap inputMap, AffineMap outputMap,
Value input, Value output) {
static Value computeSubAndExp(OpBuilder &builder, Location loc,
AffineMap inputMap, AffineMap outputMap,
Value input, Value output, bool useExp2) {
SmallVector<AffineMap> compressedMaps =
compressUnusedDims(SmallVector<AffineMap>{inputMap, outputMap});
inputMap = compressedMaps[0];
Expand All @@ -274,7 +274,11 @@ static Value computeSubAndExp2(OpBuilder &builder, Location loc,
Value in = convertScalarToDtype(b, loc, args[0], args[1].getType(),
/*isUnsignedCast=*/false);
Value diff = b.create<arith::SubFOp>(loc, args[1], in);
Value weight = b.create<math::Exp2Op>(loc, diff);
Value weight;
if (useExp2)
weight = b.create<math::Exp2Op>(loc, diff);
else
weight = b.create<math::ExpOp>(loc, diff);
b.create<linalg::YieldOp>(loc, weight);
});
return genericOp.getResult(0);
Expand Down Expand Up @@ -405,10 +409,15 @@ FailureOr<SmallVector<Value>> AttentionOp::decomposeOperation(OpBuilder &b) {
std::optional<Value> mask = getMask();
DictionaryAttr config = getDecompositionConfigAttr();

bool useExp2 = true;
DictionaryAttr qkAttrs, pvAttrs;
if (config) {
qkAttrs = config.getAs<DictionaryAttr>(getQKAttrStr());
pvAttrs = config.getAs<DictionaryAttr>(getPVAttrStr());
if (mlir::BoolAttr useExp2Attr = mlir::dyn_cast_or_null<mlir::BoolAttr>(
config.get(getUseExp2AttrStr()))) {
useExp2 = useExp2Attr.getValue();
}
}
Value output = getOutput();

Expand Down Expand Up @@ -475,7 +484,7 @@ FailureOr<SmallVector<Value>> AttentionOp::decomposeOperation(OpBuilder &b) {

// P = exp2(S - max)
AffineMap pMap = sMap;
Value p = computeSubAndExp2(b, loc, maxMap, sMap, max, s);
Value p = computeSubAndExp(b, loc, maxMap, sMap, max, s, useExp2);

// sum = rowSum(P)
Value sum = reduce<arith::AddFOp>(b, loc, pMap, sumMap, p, sumFill);
Expand Down Expand Up @@ -524,10 +533,15 @@ OnlineAttentionOp::decomposeOperation(OpBuilder &b) {
Type elementType = getElementTypeOrSelf(getOutput().getType());
DictionaryAttr config = getDecompositionConfigAttr();

bool useExp2 = true;
DictionaryAttr qkAttrs, pvAttrs;
if (config) {
qkAttrs = config.getAs<DictionaryAttr>(getQKAttrStr());
pvAttrs = config.getAs<DictionaryAttr>(getPVAttrStr());
if (mlir::BoolAttr useExp2Attr = mlir::dyn_cast_or_null<mlir::BoolAttr>(
config.get(getUseExp2AttrStr()))) {
useExp2 = useExp2Attr.getValue();
}
}

FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get(
Expand Down Expand Up @@ -561,7 +575,8 @@ OnlineAttentionOp::decomposeOperation(OpBuilder &b) {
// norm = exp2(oldMax - newMax)
// normMap = maxMap
AffineMap normMap = getMaxMap();
Value norm = computeSubAndExp2(b, loc, maxMap, normMap, newMax, oldMax);
Value norm =
computeSubAndExp(b, loc, maxMap, normMap, newMax, oldMax, useExp2);

// normSum = norm * oldSum
AffineMap sumMap = getSumMap();
Expand All @@ -571,7 +586,7 @@ OnlineAttentionOp::decomposeOperation(OpBuilder &b) {
// P = exp2(S - newMax)
// PMap = SMap
AffineMap pMap = sMap;
Value p = computeSubAndExp2(b, loc, maxMap, sMap, newMax, s);
Value p = computeSubAndExp(b, loc, maxMap, sMap, newMax, s, useExp2);

// newSum = normSum + rowSum(P)
Value newSum = reduce<arith::AddFOp>(b, loc, pMap, sumMap, p, normSum);
Expand Down
2 changes: 2 additions & 0 deletions compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -614,6 +614,7 @@ def IREELinalgExt_AttentionOp : IREELinalgExt_PureOp<"attention",
// Attributes to set on QK and PV matmul after decomposition.
static StringRef getQKAttrStr() { return "qk_attrs"; }
static StringRef getPVAttrStr() { return "pv_attrs"; }
static StringRef getUseExp2AttrStr() { return "use_exp2"; }
}];
}

Expand Down Expand Up @@ -747,6 +748,7 @@ def IREELinalgExt_OnlineAttentionOp : IREELinalgExt_PureOp<"online_attention",
// Attributes to set on QK and PV matmul after decomposition.
static StringRef getQKAttrStr() { return "qk_attrs"; }
static StringRef getPVAttrStr() { return "pv_attrs"; }
static StringRef getUseExp2AttrStr() { return "use_exp2"; }
}];
}
//===----------------------------------------------------------------------===//
Expand Down
2 changes: 0 additions & 2 deletions compiler/src/iree/compiler/Dialect/LinalgExt/IR/TilingInterfaceImpl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1213,8 +1213,6 @@ LogicalResult PackOp::generateScalarImplementation(OpBuilder &builder,
LogicalResult UnPackOp::generateScalarImplementation(OpBuilder &builder,
Location loc,
ValueRange ivs) {
assert(ivs.size() == getOutputRank() &&
"number of ivs must match the rank of the output tensor");
OpBuilder::InsertionGuard g(builder);
ReifiedRankedShapedTypeDims outputShape;
if (failed(reifyResultShapes(builder, outputShape))) {
Expand Down
Loading