Skip to content

Commit bcc050b

Browse files
jtuylsjtuyls
committed
[Encoding] Add encoding_dims flow/stream conversion
Signed-off-by: Jorn Tuyls <jorn.tuyls@gmail.com>
1 parent a7328f7 commit bcc050b

22 files changed

+507
-48
lines changed

‎compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
88

9+
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
910
#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
1011
#include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
1112
#include "iree/compiler/Dialect/Util/IR/ClosureOpUtils.h"
@@ -1657,6 +1658,22 @@ LogicalResult TensorEncodeOp::verify() {
16571658
resultType.getShape()))) {
16581659
return emitOpError("the operand shape and result shape are not compatible");
16591660
}
1661+
// Verify encoding_dims count matches what the encoding expects.
1662+
// Check both operand (for decode) and result (for encode) encodings.
1663+
for (RankedTensorType type : {operandType, resultType}) {
1664+
if (IREE::Encoding::SerializableAttr encodingAttr =
1665+
dyn_cast_or_null<IREE::Encoding::SerializableAttr>(
1666+
type.getEncoding())) {
1667+
std::optional<int64_t> expectedDims =
1668+
encodingAttr.getNumDynamicEncodingDims();
1669+
if (expectedDims.has_value() &&
1670+
static_cast<int64_t>(getEncodingDims().size()) != *expectedDims) {
1671+
return emitOpError() << "encoding expects " << *expectedDims
1672+
<< " encoding dim(s), but "
1673+
<< getEncodingDims().size() << " provided";
1674+
}
1675+
}
1676+
}
16601677
return success();
16611678
}
16621679

‎compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1306,12 +1306,17 @@ def FLOW_TensorEncodeOp : FLOW_PureOp<"tensor.encode", [
13061306
let summary = [{Performs a full tensor encode operation.}];
13071307
let description = [{
13081308
Encode the input tensor into an encoded output tensor.
1309+
1310+
The optional `encoding_dims` operand carries encoding-specific dynamic
1311+
values (e.g., M, N, K dimensions for matmul encodings). These values are
1312+
used for runtime layout selection based on problem size.
13091313
}];
13101314

13111315
let arguments = (ins
13121316
FLOW_Tensor:$operand,
13131317
FLOW_ShapeDynamicDims:$operand_dims,
1314-
FLOW_ShapeDynamicDims:$result_dims
1318+
FLOW_ShapeDynamicDims:$result_dims,
1319+
Variadic<Index>:$encoding_dims
13151320
);
13161321
let results = (outs
13171322
FLOW_Tensor:$result
@@ -1321,6 +1326,7 @@ def FLOW_TensorEncodeOp : FLOW_PureOp<"tensor.encode", [
13211326
$operand `:`
13221327
type($operand) (`{` $operand_dims^ `}`)? `->`
13231328
type($result) (`{` $result_dims^ `}`)?
1329+
(`encoding_dims` `{` $encoding_dims^ `}`)?
13241330
attr-dict-with-keyword
13251331
}];
13261332

‎compiler/src/iree/compiler/Dialect/Flow/IR/test/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ iree_lit_test_suite(
2121
"dispatch_folding.mlir",
2222
"dispatch_ops.mlir",
2323
"executable_ops.mlir",
24+
"invalid.mlir",
2425
"resolve_dim_ops.mlir",
2526
"tensor_folding.mlir",
2627
"tensor_ops.mlir",

‎compiler/src/iree/compiler/Dialect/Flow/IR/test/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ iree_lit_test_suite(
1818
"dispatch_folding.mlir"
1919
"dispatch_ops.mlir"
2020
"executable_ops.mlir"
21+
"invalid.mlir"
2122
"resolve_dim_ops.mlir"
2223
"tensor_folding.mlir"
2324
"tensor_ops.mlir"

‎compiler/src/iree/compiler/Dialect/Flow/IR/test/invalid.mlir‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
// RUN: iree-opt --split-input-file --verify-diagnostics %s
2+
3+
// Test: encoding expects 1 encoding dim but 0 provided (encode direction).
4+
#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
5+
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
6+
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
7+
#encoding = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map0, #map1, #map2], iteration_sizes = [?, 128, 64]>
8+
util.func public @encode_missing_encoding_dims(%arg0 : tensor<?x64xf32>, %m : index) -> tensor<?x64xf32, #encoding> {
9+
// expected-error @+1 {{encoding expects 1 encoding dim(s), but 0 provided}}
10+
%0 = flow.tensor.encode %arg0 : tensor<?x64xf32>{%m} -> tensor<?x64xf32, #encoding>{%m}
11+
util.return %0 : tensor<?x64xf32, #encoding>
12+
}
13+
14+
// -----
15+
16+
// Test: encoding expects 1 encoding dim but 2 provided (encode direction).
17+
#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
18+
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
19+
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
20+
#encoding = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map0, #map1, #map2], iteration_sizes = [?, 128, 64]>
21+
util.func public @encode_too_many_encoding_dims(%arg0 : tensor<?x64xf32>, %m : index) -> tensor<?x64xf32, #encoding> {
22+
// expected-error @+1 {{encoding expects 1 encoding dim(s), but 2 provided}}
23+
%0 = flow.tensor.encode %arg0 : tensor<?x64xf32>{%m} -> tensor<?x64xf32, #encoding>{%m} encoding_dims{%m, %m}
24+
util.return %0 : tensor<?x64xf32, #encoding>
25+
}
26+
27+
// -----
28+
29+
// Test: encoding expects 1 encoding dim but 0 provided (decode direction).
30+
#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
31+
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
32+
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
33+
#encoding = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map0, #map1, #map2], iteration_sizes = [?, 128, 64]>
34+
util.func public @decode_missing_encoding_dims(%arg0 : tensor<?x64xf32, #encoding>, %m : index) -> tensor<?x64xf32> {
35+
// expected-error @+1 {{encoding expects 1 encoding dim(s), but 0 provided}}
36+
%0 = flow.tensor.encode %arg0 : tensor<?x64xf32, #encoding>{%m} -> tensor<?x64xf32>{%m}
37+
util.return %0 : tensor<?x64xf32>
38+
}

‎compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_ops.mlir‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,46 @@ util.func public @tensorEncodeChangeEncoding(%arg0 : tensor<?x4xf32, #encoding>,
227227

228228
// -----
229229

230+
// Test encoding with encoding_dims for dynamic iteration sizes.
231+
// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
232+
// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
233+
// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
234+
// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iteration_sizes = [?, 128, 64]>
235+
// CHECK-LABEL: @tensorEncodeWithEncodingDims
236+
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
237+
// CHECK-SAME: %[[M:[a-zA-Z0-9]+]]
238+
#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
239+
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
240+
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
241+
#encoding_with_dims = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map0, #map1, #map2], iteration_sizes = [?, 128, 64]>
242+
util.func public @tensorEncodeWithEncodingDims(%arg0 : tensor<?x64xf32>, %m : index) -> tensor<?x64xf32, #encoding_with_dims> {
243+
// CHECK: %[[RES:.+]] = flow.tensor.encode %[[ARG0]] : tensor<?x64xf32>{%[[M]]} -> tensor<?x64xf32, #[[$ENCODING]]>{%[[M]]} encoding_dims{%[[M]]}
244+
%0 = flow.tensor.encode %arg0 : tensor<?x64xf32>{%m} -> tensor<?x64xf32, #encoding_with_dims>{%m} encoding_dims{%m}
245+
util.return %0 : tensor<?x64xf32, #encoding_with_dims>
246+
}
247+
248+
// -----
249+
250+
// Test decoding (unset_encoding direction) - operand has encoding, result doesn't.
251+
// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
252+
// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
253+
// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
254+
// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iteration_sizes = [?, 128, 64]>
255+
// CHECK-LABEL: @tensorDecodeWithEncodingDims
256+
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
257+
// CHECK-SAME: %[[M:[a-zA-Z0-9]+]]
258+
#map0_dec = affine_map<(d0, d1, d2) -> (d0, d2)>
259+
#map1_dec = affine_map<(d0, d1, d2) -> (d2, d1)>
260+
#map2_dec = affine_map<(d0, d1, d2) -> (d0, d1)>
261+
#encoding_dec = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map0_dec, #map1_dec, #map2_dec], iteration_sizes = [?, 128, 64]>
262+
util.func public @tensorDecodeWithEncodingDims(%arg0 : tensor<?x64xf32, #encoding_dec>, %m : index) -> tensor<?x64xf32> {
263+
// CHECK: %[[RES:.+]] = flow.tensor.encode %[[ARG0]] : tensor<?x64xf32, #[[$ENCODING]]>{%[[M]]} -> tensor<?x64xf32>{%[[M]]} encoding_dims{%[[M]]}
264+
%0 = flow.tensor.encode %arg0 : tensor<?x64xf32, #encoding_dec>{%m} -> tensor<?x64xf32>{%m} encoding_dims{%m}
265+
util.return %0 : tensor<?x64xf32>
266+
}
267+
268+
// -----
269+
230270
// CHECK-LABEL: @tensorSlice
231271
util.func public @tensorSlice(%arg0 : tensor<4x4xf32>, %arg1 : index, %arg2 : index) -> tensor<2x2xf32> {
232272
// CHECK-NEXT: %0 = flow.tensor.slice %arg0[%arg1, %arg2 for %arg2, %arg1] : tensor<4x4xf32> -> tensor<2x2xf32>

‎compiler/src/iree/compiler/Dialect/Flow/Transforms/InjectDispatchTracing.cpp‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
#include <utility>
88

9+
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
910
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
1011
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
1112
#include "mlir/IR/Builders.h"
@@ -44,8 +45,9 @@ static SmallVector<TensorValue> filterTensorValues(ValueRange &&range,
4445
/// Sets all `Value`s of the `TensorValue`s in `tensorValues` to the row major
4546
/// layout by inserting flow.tensor.encode ops before any Value that has an
4647
/// encoding. Populates `decodedIndices` with the indices of `tensorValues` that
47-
/// were decoded.
48-
static SmallVector<Value>
48+
/// were decoded. Returns failure if an encoding with dynamic encoding
49+
/// dimensions is encountered (not yet supported).
50+
static FailureOr<SmallVector<Value>>
4951
getInRowMajorLayout(OpBuilder &builder, SmallVector<TensorValue> tensorValues,
5052
SmallVector<int64_t> &decodedIndices) {
5153
SmallVector<Value> rowMajorTensors;
@@ -55,11 +57,27 @@ getInRowMajorLayout(OpBuilder &builder, SmallVector<TensorValue> tensorValues,
5557
rowMajorTensors.push_back(v.value);
5658
continue;
5759
}
60+
// Fail for encodings with dynamic encoding dimensions. We can't retrieve
61+
// encoding_dims through dispatch boundaries, so tracing with dynamic
62+
// layout specialization is not yet supported.
63+
if (auto encodingAttr = dyn_cast<IREE::Encoding::SerializableAttr>(
64+
rankedTensorType.getEncoding())) {
65+
if (std::optional<int64_t> numDynamic =
66+
encodingAttr.getNumDynamicEncodingDims()) {
67+
if (*numDynamic > 0) {
68+
emitError(v.value.getLoc())
69+
<< "tracing encoded tensors with dynamic encoding dimensions is "
70+
"not yet supported";
71+
return failure();
72+
}
73+
}
74+
}
5875
OpBuilder::InsertionGuard g(builder);
5976
builder.setInsertionPointAfterValue(v.value);
6077
Value rowMajorTensor = IREE::Flow::TensorEncodeOp::create(
6178
builder, v.value.getLoc(), rankedTensorType.dropEncoding(), v.value,
62-
/*operand_dims=*/v.dynamicDims, /*result_dims=*/v.dynamicDims);
79+
/*operand_dims=*/v.dynamicDims, /*result_dims=*/v.dynamicDims,
80+
/*encoding_dims=*/ValueRange{});
6381
rowMajorTensors.push_back(rowMajorTensor);
6482
decodedIndices.push_back(idx);
6583
}
@@ -91,36 +109,42 @@ struct InjectDispatchTracingPass
91109
SmallVector<TensorValue> inputTensorValues = filterTensorValues(
92110
dispatchOp.getArguments(), dispatchOp.getArgumentDims());
93111
SmallVector<int64_t> decodedInputIndices;
94-
SmallVector<Value> decodedInputValues =
112+
FailureOr<SmallVector<Value>> decodedInputValues =
95113
getInRowMajorLayout(builder, inputTensorValues, decodedInputIndices);
114+
if (failed(decodedInputValues)) {
115+
return signalPassFailure();
116+
}
96117
std::string inputsLabelStr = appendDecodedValuesToLabel(
97118
entryPointName + " inputs", decodedInputIndices);
98119
StringAttr inputsLabel = builder.getStringAttr(inputsLabelStr);
99120
IREE::Flow::TensorTraceOp::create(builder, dispatchOp.getLoc(),
100-
inputsLabel, decodedInputValues);
121+
inputsLabel, *decodedInputValues);
101122

102123
// Output tensors:
103124
SmallVector<TensorValue> resultTensorValues = filterTensorValues(
104125
dispatchOp.getResults(), dispatchOp.getResultDims());
105126
SmallVector<int64_t> decodedOutputIndices;
106-
SmallVector<Value> decodedResultValues = getInRowMajorLayout(
127+
FailureOr<SmallVector<Value>> decodedResultValues = getInRowMajorLayout(
107128
builder, resultTensorValues, decodedOutputIndices);
129+
if (failed(decodedResultValues)) {
130+
return signalPassFailure();
131+
}
108132
std::string outputsLabelStr = appendDecodedValuesToLabel(
109133
entryPointName + " outputs", decodedOutputIndices);
110134
StringAttr outputsLabel = builder.getStringAttr(outputsLabelStr);
111135

112136
// Set insertion point to the last decoded value before creating the
113137
// trace op.
114-
Operation *lastResult = decodedResultValues.front().getDefiningOp();
138+
Operation *lastResult = decodedResultValues->front().getDefiningOp();
115139
DominanceInfo domInfo(funcOp);
116-
for (Value v : decodedResultValues) {
140+
for (Value v : *decodedResultValues) {
117141
if (domInfo.dominates(lastResult, v.getDefiningOp())) {
118142
lastResult = v.getDefiningOp();
119143
}
120144
}
121145
builder.setInsertionPointAfter(lastResult);
122146
IREE::Flow::TensorTraceOp::create(builder, dispatchOp.getLoc(),
123-
outputsLabel, decodedResultValues);
147+
outputsLabel, *decodedResultValues);
124148
}
125149
}
126150
};

‎compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,8 @@ struct ConvertTensorEncodeOp
257257
rewriter, op.getLoc(), unknownType, operand.resource,
258258
op.getOperand().getType(), op.getOperandDims(), operand.resourceSize,
259259
op.getResult().getType(), flattenValues(adaptor.getResultDims()),
260-
resultSize, executionAffinityAttr);
260+
resultSize, flattenValues(adaptor.getEncodingDims()),
261+
executionAffinityAttr);
261262
rewriter.replaceOpWithMultiple(op, {{encodeOp, resultSize}});
262263
return success();
263264
}

‎compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,30 @@ util.func public @tensorEncodeChangeEncoding(%arg0 : tensor<?x4xf32, #encoding>,
204204
// CHECK: util.return %[[RESULT]], %[[SIZE]] : !stream.resource<*>, index
205205
util.return %0 : tensor<?x4xf32, #encoding1>
206206
}
207+
208+
// -----
209+
210+
// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.testing<>
211+
// CHECK-LABEL: @tensorEncodeWithEncodingDims
212+
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
213+
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
214+
// CHECK-SAME: %[[D0:[a-zA-Z0-9]+]]
215+
// CHECK-SAME: %[[D1:[a-zA-Z0-9]+]]
216+
// CHECK-SAME: %[[M:[a-zA-Z0-9]+]]
217+
// CHECK-SAME: %[[N:[a-zA-Z0-9]+]]
218+
// CHECK-SAME: %[[K:[a-zA-Z0-9]+]]
219+
#encoding = #iree_encoding.testing<>
220+
util.func public @tensorEncodeWithEncodingDims(%input: tensor<?x?xf32>, %d0: index, %d1: index, %m: index, %n: index, %k: index) -> tensor<?x?xf32, #encoding> {
221+
// CHECK: %[[SIZE:.+]] = stream.tensor.sizeof tensor<?x?xf32, #[[$ENCODING]]>{%[[D0]], %[[D1]]} : index
222+
// CHECK: %[[RESULT:.+]] = stream.tensor.encode %[[ARG0]]
223+
// CHECK-SAME: : tensor<?x?xf32>{%[[D0]], %[[D1]]} in !stream.resource<*>{%[[ARG1]]}
224+
// CHECK-SAME: -> tensor<?x?xf32, #[[$ENCODING]]>{%[[D0]], %[[D1]]} in !stream.resource<*>{%[[SIZE]]}
225+
// CHECK-SAME: encoding_dims{%[[M]], %[[N]], %[[K]]}
226+
%0 = flow.tensor.encode %input : tensor<?x?xf32>{%d0, %d1} -> tensor<?x?xf32, #encoding>{%d0, %d1} encoding_dims {%m, %n, %k}
227+
// CHECK: util.return %[[RESULT]], %[[SIZE]] : !stream.resource<*>, index
228+
util.return %0 : tensor<?x?xf32, #encoding>
229+
}
230+
207231
// -----
208232

209233
// CHECK-LABEL: @tensorAlloca

‎compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2281,6 +2281,23 @@ LogicalResult TensorEncodeOp::verify() {
22812281
failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
22822282
return failure();
22832283
}
2284+
// Verify encoding_dims count matches what the encoding expects.
2285+
// Check both source (for decode) and result (for encode) encodings.
2286+
for (Type type : {op.getSourceEncoding(), op.getResultEncoding()}) {
2287+
RankedTensorType tensorType = cast<RankedTensorType>(type);
2288+
if (IREE::Encoding::SerializableAttr encodingAttr =
2289+
dyn_cast_or_null<IREE::Encoding::SerializableAttr>(
2290+
tensorType.getEncoding())) {
2291+
std::optional<int64_t> expectedDims =
2292+
encodingAttr.getNumDynamicEncodingDims();
2293+
if (expectedDims.has_value() &&
2294+
static_cast<int64_t>(op.getEncodingDims().size()) != *expectedDims) {
2295+
return op.emitOpError() << "encoding expects " << *expectedDims
2296+
<< " encoding dim(s), but "
2297+
<< op.getEncodingDims().size() << " provided";
2298+
}
2299+
}
2300+
}
22842301
return success();
22852302
}
22862303

0 commit comments

Comments
 (0)