Skip to content

Commit 0cc753c

Browse files
ziereisziereis
authored
[TorchToTMTenor] enable broadcasting query- and key_seq_len of attention mask (#4464)
I ran into a problem where i would get a attention mask of shape [1,1,1,1] which would need to be broadcasted along the `key_seq_len` dimension which is not possible with the current implementation. This led to verification errors down the line. This change enables broadcasting the `query_seq_len` and `key_seq_len` dimension of the attention mask if required.
1 parent b249135 commit 0cc753c

2 files changed

Lines changed: 104 additions & 40 deletions

File tree

‎lib/Conversion/TorchToTMTensor/TorchToTMTensor.cpp‎

Lines changed: 41 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1871,63 +1871,64 @@ class ConvertAtenScaledDotProductAttentionOp
18711871
mask = genericOp.getResult(0);
18721872
}
18731873

1874-
// Broadcast the batch dimensions of the mask:
1874+
// Broadcast the mask to the expected attention-mask shape:
1875+
// [..., query_seq_len, key_seq_len].
18751876
if (!isa<Torch::NoneType>(mask.getType())) {
18761877
auto maskTy = cast<RankedTensorType>(mask.getType());
18771878
int64_t rank = maskTy.getRank();
1879+
// Target mask shape: [..., query_seq_len, key_seq_len].
1880+
SmallVector<int64_t> targetMaskShape(rank);
1881+
SmallVector<Value> targetMaskDynDimValues(rank);
1882+
for (int64_t i = 0; i < rank - 2; ++i) {
1883+
targetMaskShape[i] = queryTy.getDimSize(i);
1884+
if (targetMaskShape[i] == ShapedType::kDynamic)
1885+
targetMaskDynDimValues[i] =
1886+
tensor::DimOp::create(rewriter, loc, query, i);
1887+
}
1888+
1889+
targetMaskShape[rank - 2] = queryTy.getDimSize(queryTy.getRank() - 2);
1890+
if (targetMaskShape[rank - 2] == ShapedType::kDynamic)
1891+
targetMaskDynDimValues[rank - 2] =
1892+
tensor::DimOp::create(rewriter, loc, query, queryTy.getRank() - 2);
1893+
1894+
targetMaskShape[rank - 1] = keyTy.getDimSize(keyTy.getRank() - 2);
1895+
if (targetMaskShape[rank - 1] == ShapedType::kDynamic)
1896+
targetMaskDynDimValues[rank - 1] =
1897+
tensor::DimOp::create(rewriter, loc, key, keyTy.getRank() - 2);
1898+
18781899
bool needsBroadcast = false;
1879-
for (int i = 0, s = rank - 2; i < s; ++i) {
1880-
needsBroadcast |= maskTy.getDimSize(i) != queryTy.getDimSize(i);
1900+
for (int64_t i = 0; i < rank; ++i) {
1901+
needsBroadcast |= maskTy.getDimSize(i) != targetMaskShape[i];
18811902
}
18821903

18831904
if (needsBroadcast) {
1884-
SmallVector<int64_t> maskShape;
18851905
SmallVector<Value> maskDynDims;
1886-
18871906
SmallVector<AffineExpr> maskExprs;
1888-
for (int i = 0, s = rank - 2; i < s; ++i) {
1889-
maskShape.push_back(queryTy.getDimSize(i));
1890-
1891-
if (maskTy.getDimSize(i) != queryTy.getDimSize(i)) {
1892-
maskExprs.push_back(rewriter.getAffineConstantExpr(0));
1893-
} else {
1894-
maskExprs.push_back(rewriter.getAffineDimExpr(i));
1895-
}
1896-
1897-
if (queryTy.isDynamicDim(i)) {
1898-
maskDynDims.push_back(
1899-
tensor::DimOp::create(rewriter, loc, query, i));
1900-
}
1907+
for (int64_t i = 0; i < rank; ++i) {
1908+
bool broadcastDim = maskTy.getDimSize(i) != targetMaskShape[i];
1909+
maskExprs.push_back(broadcastDim ? rewriter.getAffineConstantExpr(0)
1910+
: rewriter.getAffineDimExpr(i));
1911+
if (targetMaskShape[i] == ShapedType::kDynamic)
1912+
maskDynDims.push_back(targetMaskDynDimValues[i]);
19011913
}
19021914

1903-
maskExprs.push_back(rewriter.getAffineDimExpr(rank - 2));
1904-
maskExprs.push_back(rewriter.getAffineDimExpr(rank - 1));
1905-
maskShape.push_back(maskTy.getDimSize(rank - 2));
1906-
maskShape.push_back(maskTy.getDimSize(rank - 1));
1907-
if (maskTy.isDynamicDim(rank - 2))
1908-
maskDynDims.push_back(
1909-
tensor::DimOp::create(rewriter, loc, mask, rank - 2));
1910-
if (maskTy.isDynamicDim(rank - 1))
1911-
maskDynDims.push_back(
1912-
tensor::DimOp::create(rewriter, loc, mask, rank - 1));
1913-
19141915
SmallVector<AffineMap> affineMaps = {
19151916
AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0, maskExprs,
19161917
op.getContext()),
19171918
rewriter.getMultiDimIdentityMap(rank)};
1918-
SmallVector<utils::IteratorType> findMaxIteratorTypes(
1919+
SmallVector<utils::IteratorType> iteratorTypes(
19191920
rank, utils::IteratorType::parallel);
19201921

1921-
Value emptyMask = tensor::EmptyOp::create(
1922-
rewriter, loc, maskShape, maskTy.getElementType(), maskDynDims);
1923-
Value newMask =
1924-
linalg::GenericOp::create(
1925-
rewriter, loc, emptyMask.getType(), mask,
1926-
ValueRange({emptyMask}), affineMaps, findMaxIteratorTypes,
1927-
[&](OpBuilder &b, Location loc, ValueRange args) {
1928-
linalg::YieldOp::create(b, loc, args[0]);
1929-
})
1930-
.getResult(0);
1922+
Value emptyMask =
1923+
tensor::EmptyOp::create(rewriter, loc, targetMaskShape,
1924+
maskTy.getElementType(), maskDynDims);
1925+
Value newMask = linalg::GenericOp::create(
1926+
rewriter, loc, emptyMask.getType(), mask,
1927+
ValueRange({emptyMask}), affineMaps, iteratorTypes,
1928+
[&](OpBuilder &b, Location loc, ValueRange args) {
1929+
linalg::YieldOp::create(b, loc, args[0]);
1930+
})
1931+
.getResult(0);
19311932
mask = newMask;
19321933
}
19331934
}

‎test/Conversion/TorchToTMTensor/basic.mlir‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,69 @@ func.func @sdpa_scale_dynamic_head_dim(%query: !torch.vtensor<[1,4,8,?],f32>, %k
9797

9898
// -----
9999

100+
// CHECK: #map = affine_map<(d0, d1, d2) -> (0, d1, d2)>
101+
// CHECK: #map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
102+
// CHECK-LABEL: @sdpa_bool_mask_key_seq_dynamic
103+
// CHECK: %[[MASK_IN:.*]] = torch_c.to_builtin_tensor %arg3 : !torch.vtensor<[1,1,?],i1> -> tensor<1x1x?xi1>
104+
// CHECK: %[[KEY:.*]] = torch_c.to_builtin_tensor %arg1 : !torch.vtensor<[16,?,128],f16> -> tensor<16x?x128xf16>
105+
// CHECK: %[[C1:.*]] = arith.constant 1 : index
106+
// CHECK: %[[KEY_SEQ:.*]] = tensor.dim %[[KEY]], %[[C1]] : tensor<16x?x128xf16>
107+
// CHECK: %[[EMPTY_MASK:.*]] = tensor.empty(%[[KEY_SEQ]]) : tensor<16x1x?xi1>
108+
// CHECK: %[[BCAST_MASK:.*]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[MASK_IN]] : tensor<1x1x?xi1>) outs(%[[EMPTY_MASK]] : tensor<16x1x?xi1>)
109+
// CHECK: tm_tensor.attention ins(%{{.*}}, %{{.*}}, %{{.*}}, %[[BCAST_MASK]] : tensor<16x1x128xf16>, tensor<16x?x128xf16>, tensor<16x?x128xf16>, tensor<16x1x?xi1>)
110+
func.func @sdpa_bool_mask_key_seq_dynamic(%query: !torch.vtensor<[16,1,128],f16>, %key: !torch.vtensor<[16,?,128],f16>, %value: !torch.vtensor<[16,?,128],f16>, %mask: !torch.vtensor<[1,1,?],i1>) -> !torch.vtensor<[16,1,128],f16> {
111+
%float0 = torch.constant.float 0.000000e+00
112+
%false = torch.constant.bool false
113+
%none = torch.constant.none
114+
%0 = torch.aten.scaled_dot_product_attention %query, %key, %value, %mask, %float0, %false, %none, %false : !torch.vtensor<[16,1,128],f16>, !torch.vtensor<[16,?,128],f16>, !torch.vtensor<[16,?,128],f16>, !torch.vtensor<[1,1,?],i1>, !torch.float, !torch.bool, !torch.none, !torch.bool -> !torch.vtensor<[16,1,128],f16>
115+
return %0 : !torch.vtensor<[16,1,128],f16>
116+
}
117+
118+
// -----
119+
120+
// CHECK: #map = affine_map<(d0, d1, d2) -> (0, d1, d2)>
121+
// CHECK: #map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
122+
// CHECK-LABEL: @sdpa_bool_mask_both_seq_dynamic
123+
// CHECK: %[[MASK_IN:.*]] = torch_c.to_builtin_tensor %arg3 : !torch.vtensor<[1,?,?],i1> -> tensor<1x?x?xi1>
124+
// CHECK: %[[KEY:.*]] = torch_c.to_builtin_tensor %arg1 : !torch.vtensor<[16,?,128],f16> -> tensor<16x?x128xf16>
125+
// CHECK: %[[QUERY:.*]] = torch_c.to_builtin_tensor %arg0 : !torch.vtensor<[16,?,128],f16> -> tensor<16x?x128xf16>
126+
// CHECK: %[[C1_A:.*]] = arith.constant 1 : index
127+
// CHECK: %[[QSEQ:.*]] = tensor.dim %[[QUERY]], %[[C1_A]] : tensor<16x?x128xf16>
128+
// CHECK: %[[C1_B:.*]] = arith.constant 1 : index
129+
// CHECK: %[[KSEQ:.*]] = tensor.dim %[[KEY]], %[[C1_B]] : tensor<16x?x128xf16>
130+
// CHECK: %[[EMPTY_MASK:.*]] = tensor.empty(%[[QSEQ]], %[[KSEQ]]) : tensor<16x?x?xi1>
131+
// CHECK: %[[BCAST_MASK:.*]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[MASK_IN]] : tensor<1x?x?xi1>) outs(%[[EMPTY_MASK]] : tensor<16x?x?xi1>)
132+
// CHECK: tm_tensor.attention ins(%{{.*}}, %{{.*}}, %{{.*}}, %[[BCAST_MASK]] : tensor<16x?x128xf16>, tensor<16x?x128xf16>, tensor<16x?x128xf16>, tensor<16x?x?xi1>)
133+
func.func @sdpa_bool_mask_both_seq_dynamic(%query: !torch.vtensor<[16,?,128],f16>, %key: !torch.vtensor<[16,?,128],f16>, %value: !torch.vtensor<[16,?,128],f16>, %mask: !torch.vtensor<[1,?,?],i1>) -> !torch.vtensor<[16,?,128],f16> {
134+
%float0 = torch.constant.float 0.000000e+00
135+
%false = torch.constant.bool false
136+
%none = torch.constant.none
137+
%0 = torch.aten.scaled_dot_product_attention %query, %key, %value, %mask, %float0, %false, %none, %false : !torch.vtensor<[16,?,128],f16>, !torch.vtensor<[16,?,128],f16>, !torch.vtensor<[16,?,128],f16>, !torch.vtensor<[1,?,?],i1>, !torch.float, !torch.bool, !torch.none, !torch.bool -> !torch.vtensor<[16,?,128],f16>
138+
return %0 : !torch.vtensor<[16,?,128],f16>
139+
}
140+
141+
// -----
142+
143+
// CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, 0, d2, 0)>
144+
// CHECK-LABEL: @sdpa_bool_mask_4d_static_ones
145+
// CHECK: %[[MASK_IN:.*]] = torch_c.to_builtin_tensor %arg3 : !torch.vtensor<[1,1,1,1],i1> -> tensor<1x1x1x1xi1>
146+
// CHECK: %[[KEY:.*]] = torch_c.to_builtin_tensor %arg1 : !torch.vtensor<[1,16,?,128],f16> -> tensor<1x16x?x128xf16>
147+
// CHECK: %[[C2:.*]] = arith.constant 2 : index
148+
// CHECK: %[[KSEQ:.*]] = tensor.dim %[[KEY]], %[[C2]] : tensor<1x16x?x128xf16>
149+
// CHECK: %[[EMPTY_MASK:.*]] = tensor.empty(%[[KSEQ]]) : tensor<1x16x1x?xi1>
150+
// CHECK: %[[BCAST_MASK:.*]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%[[MASK_IN]] : tensor<1x1x1x1xi1>) outs(%[[EMPTY_MASK]] : tensor<1x16x1x?xi1>)
151+
// CHECK: %[[COLLAPSED_MASK:.*]] = tensor.collapse_shape %[[BCAST_MASK]] {{.*}} : tensor<1x16x1x?xi1> into tensor<16x1x?xi1>
152+
// CHECK: tm_tensor.attention ins(%{{.*}}, %{{.*}}, %{{.*}}, %[[COLLAPSED_MASK]] : tensor<16x1x128xf16>, tensor<16x?x128xf16>, tensor<16x?x128xf16>, tensor<16x1x?xi1>)
153+
func.func @sdpa_bool_mask_4d_static_ones(%query: !torch.vtensor<[1,16,1,128],f16>, %key: !torch.vtensor<[1,16,?,128],f16>, %value: !torch.vtensor<[1,16,?,128],f16>, %mask: !torch.vtensor<[1,1,1,1],i1>) -> !torch.vtensor<[1,16,1,128],f16> {
154+
%float0 = torch.constant.float 0.000000e+00
155+
%false = torch.constant.bool false
156+
%none = torch.constant.none
157+
%0 = torch.aten.scaled_dot_product_attention %query, %key, %value, %mask, %float0, %false, %none, %false : !torch.vtensor<[1,16,1,128],f16>, !torch.vtensor<[1,16,?,128],f16>, !torch.vtensor<[1,16,?,128],f16>, !torch.vtensor<[1,1,1,1],i1>, !torch.float, !torch.bool, !torch.none, !torch.bool -> !torch.vtensor<[1,16,1,128],f16>
158+
return %0 : !torch.vtensor<[1,16,1,128],f16>
159+
}
160+
161+
// -----
162+
100163
// CHECK-LABEL: @scatter_src_i64_index
101164
// CHECK: tm_tensor.scatter {dimension_map = array<i64: 0, 1, 2>} unique_indices(false) ins(%{{.*}}, %{{.*}} : tensor<?xf32>, tensor<?x3xi64>) outs(%{{.*}} : tensor<10x8x6xf32>) {
102165
// CHECK: ^bb0(%arg3: f32, %arg4: f32):

0 commit comments

Comments
 (0)