[Encoding] Add encoding_dims flow/stream conversion

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

@@ -6,6 +6,7 @@
 
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
 #include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
 #include "iree/compiler/Dialect/Util/IR/ClosureOpUtils.h"
@@ -1657,6 +1658,22 @@ LogicalResult TensorEncodeOp::verify() {
                                          resultType.getShape()))) {
     return emitOpError("the operand shape and result shape are not compatible");
   }
+  // Verify encoding_dims count matches what the encoding expects.
+  // Check both operand (for decode) and result (for encode) encodings.
+  for (RankedTensorType type : {operandType, resultType}) {
+    if (IREE::Encoding::SerializableAttr encodingAttr =
+            dyn_cast_or_null<IREE::Encoding::SerializableAttr>(
+                type.getEncoding())) {
+      std::optional<int64_t> expectedDims =
+          encodingAttr.getNumDynamicEncodingDims();
+      if (expectedDims.has_value() &&
+          static_cast<int64_t>(getEncodingDims().size()) != *expectedDims) {
+        return emitOpError() << "encoding expects " << *expectedDims
+                             << " encoding dim(s), but "
+                             << getEncodingDims().size() << " provided";
+      }
+    }
+  }
   return success();
 }
 

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

@@ -1306,12 +1306,17 @@ def FLOW_TensorEncodeOp : FLOW_PureOp<"tensor.encode", [
   let summary = [{Performs a full tensor encode operation.}];
   let description = [{
     Encode the input tensor into an encoded output tensor.
+
+    The optional `encoding_dims` operand carries encoding-specific dynamic
+    values (e.g., M, N, K dimensions for matmul encodings). These values are
+    used for runtime layout selection based on problem size.
   }];
 
   let arguments = (ins
     FLOW_Tensor:$operand,
     FLOW_ShapeDynamicDims:$operand_dims,
-    FLOW_ShapeDynamicDims:$result_dims
+    FLOW_ShapeDynamicDims:$result_dims,
+    Variadic<Index>:$encoding_dims
   );
   let results = (outs
     FLOW_Tensor:$result
@@ -1321,6 +1326,7 @@ def FLOW_TensorEncodeOp : FLOW_PureOp<"tensor.encode", [
     $operand `:`
     type($operand) (`{` $operand_dims^ `}`)? `->`
     type($result) (`{` $result_dims^ `}`)?
+    (`encoding_dims` `{` $encoding_dims^ `}`)?
     attr-dict-with-keyword
   }];
 

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

@@ -21,6 +21,7 @@ iree_lit_test_suite(
             "dispatch_folding.mlir",
             "dispatch_ops.mlir",
             "executable_ops.mlir",
+            "invalid.mlir",
             "resolve_dim_ops.mlir",
             "tensor_folding.mlir",
             "tensor_ops.mlir",

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

@@ -18,6 +18,7 @@ iree_lit_test_suite(
     "dispatch_folding.mlir"
     "dispatch_ops.mlir"
     "executable_ops.mlir"
+    "invalid.mlir"
     "resolve_dim_ops.mlir"
     "tensor_folding.mlir"
     "tensor_ops.mlir"

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

@@ -0,0 +1,38 @@
+// RUN: iree-opt --split-input-file --verify-diagnostics %s
+
+// Test: encoding expects 1 encoding dim but 0 provided (encode direction).
+#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#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]>
+util.func public @encode_missing_encoding_dims(%arg0 : tensor<?x64xf32>, %m : index) -> tensor<?x64xf32, #encoding> {
+  // expected-error @+1 {{encoding expects 1 encoding dim(s), but 0 provided}}
+  %0 = flow.tensor.encode %arg0 : tensor<?x64xf32>{%m} -> tensor<?x64xf32, #encoding>{%m}
+  util.return %0 : tensor<?x64xf32, #encoding>
+}
+
+// -----
+
+// Test: encoding expects 1 encoding dim but 2 provided (encode direction).
+#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#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]>
+util.func public @encode_too_many_encoding_dims(%arg0 : tensor<?x64xf32>, %m : index) -> tensor<?x64xf32, #encoding> {
+  // expected-error @+1 {{encoding expects 1 encoding dim(s), but 2 provided}}
+  %0 = flow.tensor.encode %arg0 : tensor<?x64xf32>{%m} -> tensor<?x64xf32, #encoding>{%m} encoding_dims{%m, %m}
+  util.return %0 : tensor<?x64xf32, #encoding>
+}
+
+// -----
+
+// Test: encoding expects 1 encoding dim but 0 provided (decode direction).
+#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#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]>
+util.func public @decode_missing_encoding_dims(%arg0 : tensor<?x64xf32, #encoding>, %m : index) -> tensor<?x64xf32> {
+  // expected-error @+1 {{encoding expects 1 encoding dim(s), but 0 provided}}
+  %0 = flow.tensor.encode %arg0 : tensor<?x64xf32, #encoding>{%m} -> tensor<?x64xf32>{%m}
+  util.return %0 : tensor<?x64xf32>
+}

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

@@ -227,6 +227,46 @@ util.func public @tensorEncodeChangeEncoding(%arg0 : tensor<?x4xf32, #encoding>,
 
 // -----
 
+// Test encoding with encoding_dims for dynamic iteration sizes.
+// CHECK-DAG:   #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG:   #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG:   #[[$MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// 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]>
+// CHECK-LABEL: @tensorEncodeWithEncodingDims
+// CHECK-SAME:    %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[M:[a-zA-Z0-9]+]]
+#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#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]>
+util.func public @tensorEncodeWithEncodingDims(%arg0 : tensor<?x64xf32>, %m : index) -> tensor<?x64xf32, #encoding_with_dims> {
+  // CHECK: %[[RES:.+]] = flow.tensor.encode %[[ARG0]] : tensor<?x64xf32>{%[[M]]} -> tensor<?x64xf32, #[[$ENCODING]]>{%[[M]]} encoding_dims{%[[M]]}
+  %0 = flow.tensor.encode %arg0 : tensor<?x64xf32>{%m} -> tensor<?x64xf32, #encoding_with_dims>{%m} encoding_dims{%m}
+  util.return %0 : tensor<?x64xf32, #encoding_with_dims>
+}
+
+// -----
+
+// Test decoding (unset_encoding direction) - operand has encoding, result doesn't.
+// CHECK-DAG:   #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG:   #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG:   #[[$MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// 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]>
+// CHECK-LABEL: @tensorDecodeWithEncodingDims
+// CHECK-SAME:    %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[M:[a-zA-Z0-9]+]]
+#map0_dec = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1_dec = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2_dec = affine_map<(d0, d1, d2) -> (d0, d1)>
+#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]>
+util.func public @tensorDecodeWithEncodingDims(%arg0 : tensor<?x64xf32, #encoding_dec>, %m : index) -> tensor<?x64xf32> {
+  // CHECK: %[[RES:.+]] = flow.tensor.encode %[[ARG0]] : tensor<?x64xf32, #[[$ENCODING]]>{%[[M]]} -> tensor<?x64xf32>{%[[M]]} encoding_dims{%[[M]]}
+  %0 = flow.tensor.encode %arg0 : tensor<?x64xf32, #encoding_dec>{%m} -> tensor<?x64xf32>{%m} encoding_dims{%m}
+  util.return %0 : tensor<?x64xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @tensorSlice
 util.func public @tensorSlice(%arg0 : tensor<4x4xf32>, %arg1 : index, %arg2 : index) -> tensor<2x2xf32> {
   // 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

@@ -6,6 +6,7 @@
 
 #include <utility>
 
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
 #include "mlir/IR/Builders.h"
@@ -44,8 +45,9 @@ static SmallVector<TensorValue> filterTensorValues(ValueRange &&range,
 /// Sets all `Value`s of the `TensorValue`s in `tensorValues` to the row major
 /// layout by inserting flow.tensor.encode ops before any Value that has an
 /// encoding. Populates `decodedIndices` with the indices of `tensorValues` that
-/// were decoded.
-static SmallVector<Value>
+/// were decoded. Returns failure if an encoding with dynamic encoding
+/// dimensions is encountered (not yet supported).
+static FailureOr<SmallVector<Value>>
 getInRowMajorLayout(OpBuilder &builder, SmallVector<TensorValue> tensorValues,
                     SmallVector<int64_t> &decodedIndices) {
   SmallVector<Value> rowMajorTensors;
@@ -55,11 +57,27 @@ getInRowMajorLayout(OpBuilder &builder, SmallVector<TensorValue> tensorValues,
       rowMajorTensors.push_back(v.value);
       continue;
     }
+    // Fail for encodings with dynamic encoding dimensions. We can't retrieve
+    // encoding_dims through dispatch boundaries, so tracing with dynamic
+    // layout specialization is not yet supported.
+    if (auto encodingAttr = dyn_cast<IREE::Encoding::SerializableAttr>(
+            rankedTensorType.getEncoding())) {
+      if (std::optional<int64_t> numDynamic =
+              encodingAttr.getNumDynamicEncodingDims()) {
+        if (*numDynamic > 0) {
+          emitError(v.value.getLoc())
+              << "tracing encoded tensors with dynamic encoding dimensions is "
+                 "not yet supported";
+          return failure();
+        }
+      }
+    }
     OpBuilder::InsertionGuard g(builder);
     builder.setInsertionPointAfterValue(v.value);
     Value rowMajorTensor = IREE::Flow::TensorEncodeOp::create(
         builder, v.value.getLoc(), rankedTensorType.dropEncoding(), v.value,
-        /*operand_dims=*/v.dynamicDims, /*result_dims=*/v.dynamicDims);
+        /*operand_dims=*/v.dynamicDims, /*result_dims=*/v.dynamicDims,
+        /*encoding_dims=*/ValueRange{});
     rowMajorTensors.push_back(rowMajorTensor);
     decodedIndices.push_back(idx);
   }
@@ -91,36 +109,42 @@ struct InjectDispatchTracingPass
       SmallVector<TensorValue> inputTensorValues = filterTensorValues(
           dispatchOp.getArguments(), dispatchOp.getArgumentDims());
       SmallVector<int64_t> decodedInputIndices;
-      SmallVector<Value> decodedInputValues =
+      FailureOr<SmallVector<Value>> decodedInputValues =
           getInRowMajorLayout(builder, inputTensorValues, decodedInputIndices);
+      if (failed(decodedInputValues)) {
+        return signalPassFailure();
+      }
       std::string inputsLabelStr = appendDecodedValuesToLabel(
           entryPointName + " inputs", decodedInputIndices);
       StringAttr inputsLabel = builder.getStringAttr(inputsLabelStr);
       IREE::Flow::TensorTraceOp::create(builder, dispatchOp.getLoc(),
-                                        inputsLabel, decodedInputValues);
+                                        inputsLabel, *decodedInputValues);
 
       // Output tensors:
       SmallVector<TensorValue> resultTensorValues = filterTensorValues(
           dispatchOp.getResults(), dispatchOp.getResultDims());
       SmallVector<int64_t> decodedOutputIndices;
-      SmallVector<Value> decodedResultValues = getInRowMajorLayout(
+      FailureOr<SmallVector<Value>> decodedResultValues = getInRowMajorLayout(
           builder, resultTensorValues, decodedOutputIndices);
+      if (failed(decodedResultValues)) {
+        return signalPassFailure();
+      }
       std::string outputsLabelStr = appendDecodedValuesToLabel(
           entryPointName + " outputs", decodedOutputIndices);
       StringAttr outputsLabel = builder.getStringAttr(outputsLabelStr);
 
       // Set insertion point to the last decoded value before creating the
       // trace op.
-      Operation *lastResult = decodedResultValues.front().getDefiningOp();
+      Operation *lastResult = decodedResultValues->front().getDefiningOp();
       DominanceInfo domInfo(funcOp);
-      for (Value v : decodedResultValues) {
+      for (Value v : *decodedResultValues) {
         if (domInfo.dominates(lastResult, v.getDefiningOp())) {
           lastResult = v.getDefiningOp();
         }
       }
       builder.setInsertionPointAfter(lastResult);
       IREE::Flow::TensorTraceOp::create(builder, dispatchOp.getLoc(),
-                                        outputsLabel, decodedResultValues);
+                                        outputsLabel, *decodedResultValues);
     }
   }
 };

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

@@ -257,7 +257,8 @@ struct ConvertTensorEncodeOp
         rewriter, op.getLoc(), unknownType, operand.resource,
         op.getOperand().getType(), op.getOperandDims(), operand.resourceSize,
         op.getResult().getType(), flattenValues(adaptor.getResultDims()),
-        resultSize, executionAffinityAttr);
+        resultSize, flattenValues(adaptor.getEncodingDims()),
+        executionAffinityAttr);
     rewriter.replaceOpWithMultiple(op, {{encodeOp, resultSize}});
     return success();
   }

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

@@ -204,6 +204,30 @@ util.func public @tensorEncodeChangeEncoding(%arg0 : tensor<?x4xf32, #encoding>,
   // CHECK:      util.return %[[RESULT]], %[[SIZE]] : !stream.resource<*>, index
   util.return %0 : tensor<?x4xf32, #encoding1>
 }
+
+// -----
+
+// CHECK-DAG:   #[[$ENCODING:.+]] = #iree_encoding.testing<>
+// CHECK-LABEL: @tensorEncodeWithEncodingDims
+// CHECK-SAME:    %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[D0:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[D1:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[M:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[N:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[K:[a-zA-Z0-9]+]]
+#encoding = #iree_encoding.testing<>
+util.func public @tensorEncodeWithEncodingDims(%input: tensor<?x?xf32>, %d0: index, %d1: index, %m: index, %n: index, %k: index) -> tensor<?x?xf32, #encoding> {
+  // CHECK: %[[SIZE:.+]] = stream.tensor.sizeof tensor<?x?xf32, #[[$ENCODING]]>{%[[D0]], %[[D1]]} : index
+  // CHECK: %[[RESULT:.+]] = stream.tensor.encode %[[ARG0]]
+  // CHECK-SAME:   : tensor<?x?xf32>{%[[D0]], %[[D1]]} in !stream.resource<*>{%[[ARG1]]}
+  // CHECK-SAME:   -> tensor<?x?xf32, #[[$ENCODING]]>{%[[D0]], %[[D1]]} in !stream.resource<*>{%[[SIZE]]}
+  // CHECK-SAME:   encoding_dims{%[[M]], %[[N]], %[[K]]}
+  %0 = flow.tensor.encode %input : tensor<?x?xf32>{%d0, %d1} -> tensor<?x?xf32, #encoding>{%d0, %d1} encoding_dims {%m, %n, %k}
+  // CHECK: util.return %[[RESULT]], %[[SIZE]] : !stream.resource<*>, index
+  util.return %0 : tensor<?x?xf32, #encoding>
+}
+
 // -----
 
 // CHECK-LABEL: @tensorAlloca

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

@@ -2281,6 +2281,23 @@ LogicalResult TensorEncodeOp::verify() {
       failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
     return failure();
   }
+  // Verify encoding_dims count matches what the encoding expects.
+  // Check both source (for decode) and result (for encode) encodings.
+  for (Type type : {op.getSourceEncoding(), op.getResultEncoding()}) {
+    RankedTensorType tensorType = cast<RankedTensorType>(type);
+    if (IREE::Encoding::SerializableAttr encodingAttr =
+            dyn_cast_or_null<IREE::Encoding::SerializableAttr>(
+                tensorType.getEncoding())) {
+      std::optional<int64_t> expectedDims =
+          encodingAttr.getNumDynamicEncodingDims();
+      if (expectedDims.has_value() &&
+          static_cast<int64_t>(op.getEncodingDims().size()) != *expectedDims) {
+        return op.emitOpError() << "encoding expects " << *expectedDims
+                                << " encoding dim(s), but "
+                                << op.getEncodingDims().size() << " provided";
+      }
+    }
+  }
   return success();
 }
 

compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td

@@ -1359,10 +1359,14 @@ def Stream_TensorEncodeOp : Stream_PureOp<"tensor.encode", [
 ]> {
   let summary = [{Encodes the contents of a value.}];
   let description = [{
-    Elones the contents of a value at a snapshot in time. Future changes to the
+    Encodes the contents of a value at a snapshot in time. Future changes to the
     identity encoding will not effect the result. Acts as a copy-on-write
     operation. Otherwise, an executable for encoding the value is generated
     during lowering.
+
+    The optional `encoding_dims` operand carries encoding-specific dynamic
+    values (e.g., M, N, K dimensions for matmul encodings). These values are
+    used for runtime layout selection based on problem size.
   }];
 
   let arguments = (ins
@@ -1373,6 +1377,7 @@ def Stream_TensorEncodeOp : Stream_PureOp<"tensor.encode", [
     TypeAttr:$result_encoding,
     Stream_ShapeDynamicDims:$result_encoding_dims,
     Stream_Size:$result_size,
+    Variadic<Index>:$encoding_dims,
     OptionalAttr<Stream_AffinityAttr>:$affinity
   );
   let results = (outs
@@ -1389,6 +1394,7 @@ def Stream_TensorEncodeOp : Stream_PureOp<"tensor.encode", [
     $result_encoding (`` `{` $result_encoding_dims^ `}`)?
     `in`
     type($result) `` `{` $result_size `}`
+    (`encoding_dims` `{` $encoding_dims^ `}`)?
     attr-dict-with-keyword
   }];
 

compiler/src/iree/compiler/Dialect/Stream/IR/test/BUILD.bazel

@@ -30,6 +30,7 @@ iree_lit_test_suite(
             "context_ops.mlir",
             "executable_ops.mlir",
             "file_ops.mlir",
+            "invalid.mlir",
             "resource_folding.mlir",
             "resource_ops.mlir",
             "tensor_folding.mlir",

compiler/src/iree/compiler/Dialect/Stream/IR/test/CMakeLists.txt

@@ -27,6 +27,7 @@ iree_lit_test_suite(
     "context_ops.mlir"
     "executable_ops.mlir"
     "file_ops.mlir"
+    "invalid.mlir"
     "resource_folding.mlir"
     "resource_ops.mlir"
     "tensor_folding.mlir"