[Encoding] Implement OpAsmDialectInterface for IREE Encoding dialect. (#19146)
The revision implements the dialect interface which hoists the encodings
out of the tensor type. It makes IR dumps much easier in terms of
readability because the same encodings are collapsed into a single one
and the alias is shorter.
The revision refines the lit checks for
Encoding/IR/test/roundtrip.mlir which makes the style consistent; it
adapts some tests to hoist the encodings, which matches the IR dumps.
The `IREEEncodingInlinerInterface` is renamed to
`EncodingInlinerInterface` in this patch for consistency. The `IREE`
prefix is redundant because it is already in an anonymous namespace
under the IREE namespace.
Fixes https://github.com/iree-org/iree/issues/19024
Signed-off-by: hanhanW <hanhan0912@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/test/vmvx_materialize_encoding.mlir b/compiler/src/iree/compiler/Codegen/Common/CPU/test/vmvx_materialize_encoding.mlir
index f93d138..85dd416 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CPU/test/vmvx_materialize_encoding.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/CPU/test/vmvx_materialize_encoding.mlir
@@ -43,24 +43,29 @@
-> !flow.dispatch.tensor<readwrite:tensor<?x?xi32, #encoding_result>>{%M, %N}
return
}
-
// CHECK-DAG: #[[MAP_CEILDIV:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
+// 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: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: func @matmul_lowering_i8i8i32_vmvx_ukernel()
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[M:.+]] = hal.interface.constant.load layout({{.+}}) ordinal(0)
// CHECK-DAG: %[[N:.+]] = hal.interface.constant.load layout({{.+}}) ordinal(1)
// CHECK-DAG: %[[K:.+]] = hal.interface.constant.load layout({{.+}}) ordinal(2)
-// CHECK: %[[LHS_TILE_SIZES:.+]]:2 = iree_codegen.query_tile_sizes tensor<?x?xi8, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>> -> index, index
+// CHECK: %[[LHS_TILE_SIZES:.+]]:2 = iree_codegen.query_tile_sizes tensor<?x?xi8, #[[LHS_ENCODING]]> -> index, index
// CHECK-DAG: %[[LHS_OUTER_SIZE0:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[M]], %[[LHS_TILE_SIZES]]#0]
// CHECK-DAG: %[[LHS_OUTER_SIZE1:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[K]], %[[LHS_TILE_SIZES]]#1]
// CHECK: %[[LHS_BINDING:.+]] = hal.interface.binding.subspan layout({{.+}}) binding(0)
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%[[LHS_OUTER_SIZE0]], %[[LHS_OUTER_SIZE1]], %[[LHS_TILE_SIZES]]#0, %[[LHS_TILE_SIZES]]#1}
-// CHECK: %[[RHS_TILE_SIZES:.+]]:2 = iree_codegen.query_tile_sizes tensor<?x?xi8, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>> -> index, index
+// CHECK: %[[RHS_TILE_SIZES:.+]]:2 = iree_codegen.query_tile_sizes tensor<?x?xi8, #[[RHS_ENCODING]]> -> index, index
// CHECK-DAG: %[[RHS_OUTER_SIZE0:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[N]], %[[RHS_TILE_SIZES]]#0]
// CHECK-DAG: %[[RHS_OUTER_SIZE1:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[K]], %[[RHS_TILE_SIZES]]#1]
// CHECK: %[[RHS_BINDING:.+]] = hal.interface.binding.subspan layout({{.+}}) binding(1)
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%[[RHS_OUTER_SIZE0]], %[[RHS_OUTER_SIZE1]], %[[RHS_TILE_SIZES]]#0, %[[RHS_TILE_SIZES]]#1}
-// CHECK: %[[RESULT_TILE_SIZES:.+]]:2 = iree_codegen.query_tile_sizes tensor<?x?xi32, #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>> -> index, index
+// CHECK: %[[RESULT_TILE_SIZES:.+]]:2 = iree_codegen.query_tile_sizes tensor<?x?xi32, #[[OUT_ENCODING]]> -> index, index
// CHECK-DAG: %[[RESULT_OUTER_SIZE0:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[M]], %[[RESULT_TILE_SIZES]]#0]
// CHECK-DAG: %[[RESULT_OUTER_SIZE1:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[N]], %[[RESULT_TILE_SIZES]]#1]
// CHECK: %[[OUTS_BINDING:.+]] = hal.interface.binding.subspan layout({{.+}}) binding(2)
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingDialect.cpp b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingDialect.cpp
index 46925b5..f48e73a 100644
--- a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingDialect.cpp
+++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingDialect.cpp
@@ -23,9 +23,27 @@
#undef GET_ATTRDEF_CLASSES
namespace mlir::iree_compiler::IREE::Encoding {
+namespace {
+
+// Used for custom printing support.
+struct EncodingOpAsmInterface : public OpAsmDialectInterface {
+ using OpAsmDialectInterface::OpAsmDialectInterface;
+ /// Hooks for getting an alias identifier alias for a given symbol, that is
+ /// not necessarily a part of this dialect. The identifier is used in place
+ /// of the symbol when printing textual IR. These aliases must not contain
+ /// `.` or end with a numeric digit([0-9]+). Returns success if an alias was
+ /// provided, failure otherwise.
+ AliasResult getAlias(Attribute attr, raw_ostream &os) const override {
+ if (llvm::isa<EncodingAttr>(attr)) {
+ os << "encoding";
+ return AliasResult::OverridableAlias;
+ }
+ return AliasResult::NoAlias;
+ }
+};
// Used to control inlining behavior.
-struct IREEEncodingInlinerInterface : public DialectInlinerInterface {
+struct EncodingInlinerInterface : public DialectInlinerInterface {
using DialectInlinerInterface::DialectInlinerInterface;
bool isLegalToInline(Operation *call, Operation *callable,
@@ -45,9 +63,10 @@
return true;
}
};
+} // namespace
void IREEEncodingDialect::initialize() {
- addInterfaces<IREEEncodingInlinerInterface>();
+ addInterfaces<EncodingOpAsmInterface, EncodingInlinerInterface>();
addAttributes<
#define GET_ATTRDEF_LIST
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/test/roundtrip.mlir b/compiler/src/iree/compiler/Dialect/Encoding/IR/test/roundtrip.mlir
index 2032f39..775f2fa 100644
--- a/compiler/src/iree/compiler/Dialect/Encoding/IR/test/roundtrip.mlir
+++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/test/roundtrip.mlir
@@ -1,108 +1,141 @@
// RUN: iree-opt --split-input-file %s | FileCheck %s
-// CHECK: @set_encoding_ops(%[[ARG0:.+]]: tensor<?x?xf32>)
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
func.func @set_encoding_ops(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>> {
- // CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
- %0 = iree_encoding.set_encoding %arg0 : tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
- return %0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
+ %0 = iree_encoding.set_encoding %arg0 : tensor<?x?xf32> -> tensor<?x?xf32, #encoding>
+ return %0 : tensor<?x?xf32, #encoding>
}
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK: func.func @set_encoding_ops
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?xf32> -> tensor<?x?xf32, #[[ENCODING]]>
// -----
-// CHECK: @set_encoding_ops_mixed_dynamic_static(%[[ARG0:.+]]: tensor<?x10xf32>)
-func.func @set_encoding_ops_mixed_dynamic_static(%arg0: tensor<?x10xf32>) -> tensor<20x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>> {
- // CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x10xf32> -> tensor<20x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
- %0 = iree_encoding.set_encoding %arg0 : tensor<?x10xf32> -> tensor<20x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
- return %0 : tensor<20x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
+#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+func.func @set_encoding_ops_mixed_dynamic_static(%arg0: tensor<?x10xf32>) -> tensor<20x?xf32, #encoding> {
+ %0 = iree_encoding.set_encoding %arg0 : tensor<?x10xf32> -> tensor<20x?xf32, #encoding>
+ return %0 : tensor<20x?xf32, #encoding>
}
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK: func.func @set_encoding_ops_mixed_dynamic_static
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x10xf32> -> tensor<20x?xf32, #[[ENCODING]]>
// -----
-// CHECK: @set_encoding_with_batch_matmul_user(%[[ARG0:.+]]: tensor<?x?x?xf32>)
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding1 = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [i8, i8, i32]>
+#encoding2 = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f16, f16, f32]>
+#encoding3 = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f16, f16, f16]>
+#encoding4 = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [bf16, bf16, f32]>
+#encoding5 = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [bf16, bf16, bf16]>
func.func @set_encoding_with_batch_matmul_user(%arg0: tensor<?x?x?xf32>) {
- // CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
- iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
- // CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [i8, i8, i32]>>
- iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [i8, i8, i32]>>
- // CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f16, f16, f32]>>
- iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f16, f16, f32]>>
- // CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f16, f16, f16]>>
- iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f16, f16, f16]>>
- // CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [bf16, bf16, f32]>>
- iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [bf16, bf16, f32]>>
- // CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [bf16, bf16, bf16]>>
- iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [bf16, bf16, bf16]>>
+ iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #encoding>
+ iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #encoding1>
+ iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #encoding2>
+ iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #encoding3>
+ iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #encoding4>
+ iree_encoding.set_encoding %arg0 : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #encoding5>
return
}
+// CHECK-DAG: #[[ENCODING0:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK-DAG: #[[ENCODING1:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [i8, i8, i32]>
+// CHECK-DAG: #[[ENCODING2:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f16, f16, f32]>
+// CHECK-DAG: #[[ENCODING3:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f16, f16, f16]>
+// CHECK-DAG: #[[ENCODING4:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [bf16, bf16, f32]>
+// CHECK-DAG: #[[ENCODING5:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [bf16, bf16, bf16]>
+// CHECK: func.func @set_encoding_with_batch_matmul_user
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #[[ENCODING0]]>
+// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #[[ENCODING1]]>
+// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #[[ENCODING2]]>
+// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #[[ENCODING3]]>
+// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #[[ENCODING4]]>
+// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #[[ENCODING5]]>
// -----
-// CHECK: @unset_encoding_fully_static(%[[ARG0:.+]]: tensor<3x5xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>>)
-func.func @unset_encoding_fully_static(%arg0: tensor<3x5xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>>) -> tensor<3x5xf32> {
- // CHECK: iree_encoding.unset_encoding %[[ARG0]] : tensor<3x5xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<3x5xf32>
- %0 = iree_encoding.unset_encoding %arg0 : tensor<3x5xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<3x5xf32>
+#encoding = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>
+func.func @unset_encoding_fully_static(%arg0: tensor<3x5xf32, #encoding>) -> tensor<3x5xf32> {
+ %0 = iree_encoding.unset_encoding %arg0 : tensor<3x5xf32, #encoding> -> tensor<3x5xf32>
return %0 : tensor<3x5xf32>
}
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK: func.func @unset_encoding_fully_static
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<3x5xf32, #[[ENCODING]]
+// CHECK: iree_encoding.unset_encoding %[[ARG0]] : tensor<3x5xf32, #[[ENCODING]]> -> tensor<3x5xf32>
+
// -----
-// CHECK: @unset_encoding_fully_dynamic(%[[ARG0:.+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
-// CHECK-SAME: %[[D0:[a-zA-Z0-9]+]]
-// CHECK-SAME: %[[D1:[a-zA-Z0-9]+]]
-func.func @unset_encoding_fully_dynamic(%arg0: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>>,
- %d0 : index, %d1 : index) -> tensor<?x?xf32> {
- // CHECK: iree_encoding.unset_encoding %[[ARG0]] : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<?x?xf32>{%[[D0]], %[[D1]]}
- %0 = iree_encoding.unset_encoding %arg0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<?x?xf32>{%d0, %d1}
+#encoding = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+func.func @unset_encoding_fully_dynamic(%arg0: tensor<?x?xf32, #encoding>, %d0 : index, %d1 : index) -> tensor<?x?xf32> {
+ %0 = iree_encoding.unset_encoding %arg0 : tensor<?x?xf32, #encoding> -> tensor<?x?xf32>{%d0, %d1}
return %0 : tensor<?x?xf32>
}
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK: func.func @unset_encoding_fully_dynamic
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32, #[[ENCODING]]
+// CHECK-SAME: %[[D0:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[D1:[a-zA-Z0-9]+]]
+// CHECK: iree_encoding.unset_encoding %[[ARG0]] : tensor<?x?xf32, #[[ENCODING]]> -> tensor<?x?xf32>{%[[D0]], %[[D1]]}
// -----
-// CHECK: @unset_encoding_ops_mixed_dynamic_static(%[[ARG0:.+]]: tensor<10x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
-// CHECK-SAME: %[[D0:[a-zA-Z0-9]+]]
-func.func @unset_encoding_ops_mixed_dynamic_static(%arg0: tensor<10x?xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>>,
- %d0 : index) -> tensor<?x20xf32> {
- // CHECK: iree_encoding.unset_encoding %[[ARG0]] : tensor<10x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<?x20xf32>{%[[D0]]}
- %0 = iree_encoding.unset_encoding %arg0 : tensor<10x?xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<?x20xf32>{%d0}
+#encoding = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+func.func @unset_encoding_ops_mixed_dynamic_static(%arg0: tensor<10x?xf32, #encoding>, %d0 : index) -> tensor<?x20xf32> {
+ %0 = iree_encoding.unset_encoding %arg0 : tensor<10x?xf32, #encoding> -> tensor<?x20xf32>{%d0}
return %0 : tensor<?x20xf32>
}
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK: func.func @unset_encoding_ops_mixed_dynamic_static
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<10x?xf32, #[[ENCODING]]>
+// CHECK-SAME: %[[D0:[a-zA-Z0-9]+]]
+// CHECK: iree_encoding.unset_encoding %[[ARG0]] : tensor<10x?xf32, #[[ENCODING]]> -> tensor<?x20xf32>{%[[D0]]}
// -----
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding1 = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding2 = #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>
func.func @encoding_tensors_with_ops(%arg0 : tensor<?x?xf32>,
%arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%M = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%N = tensor.dim %arg1, %c1 : tensor<?x?xf32>
- %0 = iree_encoding.set_encoding %arg0 : tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
- %1 = iree_encoding.set_encoding %arg1 : tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>>
- %2 = iree_encoding.set_encoding %arg2 : tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>
+ %0 = iree_encoding.set_encoding %arg0 : tensor<?x?xf32> -> tensor<?x?xf32, #encoding>
+ %1 = iree_encoding.set_encoding %arg1 : tensor<?x?xf32> -> tensor<?x?xf32, #encoding1>
+ %2 = iree_encoding.set_encoding %arg2 : tensor<?x?xf32> -> tensor<?x?xf32, #encoding2>
%3 = linalg.matmul
- ins(%0, %1 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>, tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>>)
- outs(%2 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>)
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>
- %4 = iree_encoding.unset_encoding %3 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<?x?xf32>{%M, %N}
+ ins(%0, %1 : tensor<?x?xf32, #encoding>, tensor<?x?xf32, #encoding1>)
+ outs(%2 : tensor<?x?xf32, #encoding2>)
+ -> tensor<?x?xf32, #encoding2>
+ %4 = iree_encoding.unset_encoding %3 : tensor<?x?xf32, #encoding2> -> tensor<?x?xf32>{%M, %N}
return %4 : tensor<?x?xf32>
}
-// CHECK-LABEL: func.func @encoding_tensors_with_ops
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
-// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>
-// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?xf32>
-// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK-DAG: %[[M:.+]] = tensor.dim %[[ARG0]], %[[C0]]
-// CHECK-DAG: %[[N:.+]] = tensor.dim %[[ARG1]], %[[C1]]
-// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
-// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
-// CHECK: %[[OUT:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
-// CHECK: %[[GEMM:.+]] = linalg.matmul
-// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
-// CHECK-SAME: outs(%[[OUT]] :
-// CHECK: %[[RESULT:.+]] = iree_encoding.unset_encoding %[[GEMM]] {{.+}} -> tensor<?x?xf32>{%[[M]], %[[N]]}
-// CHECK: return %[[RESULT]]
+// CHECK-DAG: #[[ENCODING0:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK-DAG: #[[ENCODING1:.+]] = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK-DAG: #[[ENCODING2:.+]] = #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK: func.func @encoding_tensors_with_ops
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[M:.+]] = tensor.dim %[[ARG0]], %[[C0]]
+// CHECK-DAG: %[[N:.+]] = tensor.dim %[[ARG1]], %[[C1]]
+// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
+// CHECK-SAME: tensor<?x?xf32> -> tensor<?x?xf32, #[[ENCODING0]]>
+// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
+// CHECK-SAME: tensor<?x?xf32> -> tensor<?x?xf32, #[[ENCODING1]]>
+// CHECK: %[[OUT:.+]] = iree_encoding.set_encoding %[[ARG2]]
+// CHECK-SAME: tensor<?x?xf32> -> tensor<?x?xf32, #[[ENCODING2]]>
+// CHECK: %[[GEMM:.+]] = linalg.matmul
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]]
+// CHECK-SAME: outs(%[[OUT]]
+// CHECK: %[[RESULT:.+]] = iree_encoding.unset_encoding %[[GEMM]] : tensor<?x?xf32, #[[ENCODING2]]> -> tensor<?x?xf32>{%[[M]], %[[N]]}
+// CHECK: return %[[RESULT]]
// -----
@@ -110,14 +143,17 @@
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-func.func @set_encoding_ops_with_indexing_maps(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], bcast_map = #map3>> {
- %0 = iree_encoding.set_encoding %arg0 : tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], bcast_map = #map3>>
- return %0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], bcast_map = #map3>>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], bcast_map = #map3>
+func.func @set_encoding_ops_with_indexing_maps(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32, #encoding> {
+ %0 = iree_encoding.set_encoding %arg0 : tensor<?x?xf32> -> tensor<?x?xf32, #encoding>
+ return %0 : tensor<?x?xf32, #encoding>
}
-// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32],
+// CHECK-SAME: user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], bcast_map = #[[MAP3]]
// CHECK: func.func @set_encoding_ops_with_indexing_maps(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]:
-// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP2]]], bcast_map = #[[MAP3]]>>
+// CHECK: iree_encoding.set_encoding %[[ARG0]] : tensor<?x?xf32> -> tensor<?x?xf32, #[[ENCODING]]>
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_on_tensors.mlir b/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_on_tensors.mlir
index 1522bce..2b20e30 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_on_tensors.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_on_tensors.mlir
@@ -1822,13 +1822,14 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @set_encoding_op(%arg0 : tensor<?x?xf32>)
-> tensor<?x?xf32, #encoding> {
%0 = iree_encoding.set_encoding %arg0
: tensor<?x?xf32> -> tensor<?x?xf32, #encoding>
util.return %0 : tensor<?x?xf32, #encoding>
}
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
// CHECK: util.func public @set_encoding_op
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
@@ -1839,15 +1840,15 @@
// CHECK-NEXT: %[[INARG:.+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32>>
// CHECK-SAME: %[[INDEXARG0:[a-zA-Z0-9]+]]: index
// CHECK-SAME: %[[INDEXARG1:[a-zA-Z0-9]+]]: index
-// CHECK-SAME: %[[OUTARG:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
+// CHECK-SAME: %[[OUTARG:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #[[ENCODING]]>>
// CHECK-DAG: %[[W0:.+]] = flow.dispatch.workload.ordinal %[[INDEXARG0]], 0
// CHECK-DAG: %[[W1:.+]] = flow.dispatch.workload.ordinal %[[INDEXARG1]], 1
// CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load %[[INARG]]
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?xf32>>{%[[W0]], %[[W1]]}
-// CHECK: %[[ENCODING:.+]] = iree_encoding.set_encoding %[[LOAD]]
-// CHECK: flow.dispatch.tensor.store %[[ENCODING]], %[[OUTARG]]
+// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[LOAD]]
+// CHECK: flow.dispatch.tensor.store %[[SET_ENCODING]], %[[OUTARG]]
// CHECK-SAME: sizes = [%[[W0]], %[[W1]]]
-// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>{%[[W0]], %[[W1]]}
+// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #[[ENCODING]]>>{%[[W0]], %[[W1]]}
// CHECK: flow.return
// CHECK: count(%[[WL0:[a-zA-Z0-9]+]]: index, %[[WL1:[a-zA-Z0-9]+]]: index)
// CHECK: %[[X:[a-zA-Z0-9]+]], %[[Y:[a-zA-Z0-9]+]], %[[Z:.+]] = flow.dispatch.workgroup_count_from_slice %[[WL0]], %[[WL1]]
@@ -1856,15 +1857,16 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @unset_encoding_op(%arg0 : tensor<?x?xf32, #encoding>, %d0 : index, %d1 : index)
-> tensor<?x?xf32> {
%0 = iree_encoding.unset_encoding %arg0
: tensor<?x?xf32, #encoding> -> tensor<?x?xf32>{%d0, %d1}
util.return %0 : tensor<?x?xf32>
}
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
// CHECK: util.func public @unset_encoding_op
-// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
+// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32, #[[ENCODING]]>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
@@ -1872,7 +1874,7 @@
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
// CHECK: %[[DISPATCH:.+]] = flow.dispatch.workgroups[%[[D0]], %[[D1]], %[[ARG1]], %[[ARG2]]](%[[ARG0]], %[[D0]], %[[D1]], %[[ARG1]], %[[ARG2]])
-// CHECK-NEXT: %[[INARG:.+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
+// CHECK-NEXT: %[[INARG:.+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #[[ENCODING]]>>
// CHECK-SAME: %[[ARG1_INDEX:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[ARG2_INDEX:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[INARG_INDEX0:[a-zA-Z0-9]+]]
@@ -1884,9 +1886,9 @@
// CHECK-DAG: %[[D3_W:.+]] = flow.dispatch.workload.ordinal %[[INARG_INDEX1]], 3
// CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load %[[INARG]]
// CHECK-SAME: sizes = [%[[D0_W]], %[[D1_W]]]
-// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>{%[[D0_W]], %[[D1_W]]}
-// CHECK: %[[ENCODING:.+]] = iree_encoding.unset_encoding %[[LOAD]]{{.+}}{%[[D2_W]], %[[D3_W]]
-// CHECK: flow.dispatch.tensor.store %[[ENCODING]], %[[OUTARG]]
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #[[ENCODING]]>>{%[[D0_W]], %[[D1_W]]}
+// CHECK: %[[UNSET_ENCODING:.+]] = iree_encoding.unset_encoding %[[LOAD]]{{.+}}{%[[D2_W]], %[[D3_W]]
+// CHECK: flow.dispatch.tensor.store %[[UNSET_ENCODING]], %[[OUTARG]]
// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<?x?xf32>>{%[[D2_W]], %[[D3_W]]}
// CHECK: flow.return
// CHECK: count(%[[WL0:[a-zA-Z0-9]+]]: index, %[[WL1:[a-zA-Z0-9]+]]: index, %[[WL2:[a-zA-Z0-9]+]]: index, %[[WL3:[a-zA-Z0-9]+]]: index)
@@ -1897,7 +1899,7 @@
// -----
#map = affine_map<()[s0] -> (-s0 + (s0 ceildiv 16) * 16)>
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @pad_and_set_encoding_op(%arg0 : tensor<?x?xf32>)
-> tensor<?x?xf32, #encoding> {
%c0 = arith.constant 0 : index
@@ -1917,6 +1919,7 @@
}
// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> ((s0 ceildiv 16) * 16)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (-s0 + (s0 ceildiv 16) * 16)>
+// CHECK-DAG: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
// CHECK: util.func public @pad_and_set_encoding
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
@@ -1931,7 +1934,7 @@
// CHECK-SAME: %[[INARG:.+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32>>
// CHECK-SAME: %[[PADDED_D0:[a-zA-Z0-9]+]]: index
// CHECK-SAME: %[[PADDED_D1:[a-zA-Z0-9]+]]: index
-// CHECK-SAME: %[[OUTARG:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
+// CHECK-SAME: %[[OUTARG:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #[[ENCODING]]>>
// CHECK-DAG: %[[D1_W:.+]] = flow.dispatch.workload.ordinal %[[INDEXARG1]], 0
// CHECK-DAG: %[[D0_W:.+]] = flow.dispatch.workload.ordinal %[[INDEXARG0]], 1
// CHECK-DAG: %[[PADDED_D0_W:.+]] = flow.dispatch.workload.ordinal %[[PADDED_D0]], 2
@@ -1944,7 +1947,7 @@
// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[PADDED]]
// CHECK: flow.dispatch.tensor.store %[[SET_ENCODING]], %[[OUTARG]]
// CHECK-SAME: sizes = [%[[PADDED_D0_W]], %[[PADDED_D1_W]]]
-// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>{%[[PADDED_D0_W]], %[[PADDED_D1_W]]}
+// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #[[ENCODING]]>>{%[[PADDED_D0_W]], %[[PADDED_D1_W]]}
// CHECK: flow.return
// CHECK: count(%[[WL0:[a-zA-Z0-9]+]]: index, %[[WL1:[a-zA-Z0-9]+]]: index,
// CHECK-SAME: %[[WL2:[a-zA-Z0-9]+]]: index, %[[WL3:[a-zA-Z0-9]+]]: index)
@@ -1954,7 +1957,7 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @unset_encoding_and_slice(
%arg0: tensor<?x?xf32, #encoding>,
%arg1 : index, %arg2 : index) -> tensor<?x?xf32> {
@@ -1964,8 +1967,9 @@
: tensor<?x?xf32> to tensor<?x?xf32>
util.return %1 : tensor<?x?xf32>
}
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
// CHECK: util.func public @unset_encoding_and_slice
-// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
+// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32, #[[ENCODING]]>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
@@ -1973,7 +1977,7 @@
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
// CHECK: %[[DISPATCH:.+]] = flow.dispatch.workgroups[%[[D0]], %[[D1]], %[[ARG1]], %[[ARG2]]](%[[ARG0]], %[[D0]], %[[D1]], %[[ARG1]], %[[ARG2]])
-// CHECK-NEXT: %[[INARG:.+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
+// CHECK-NEXT: %[[INARG:.+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #[[ENCODING]]>>
// CHECK-SAME: %[[INDEXARG0:[a-zA-Z0-9]+]]: index
// CHECK-SAME: %[[INDEXARG1:[a-zA-Z0-9]+]]: index
// CHECK-SAME: %[[INDEXARG2:[a-zA-Z0-9]+]]: index
@@ -1985,9 +1989,9 @@
// CHECK-DAG: %[[ARG1_W:.+]] = flow.dispatch.workload.ordinal %[[INDEXARG3]], 3
// CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load %[[INARG]]
// CHECK-SAME: sizes = [%[[D0_W]], %[[D1_W]]]
-// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>{%[[D0_W]], %[[D1_W]]}
-// CHECK: %[[ENCODING:.+]] = iree_encoding.unset_encoding %[[LOAD]]
-// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ENCODING]][0, 0] [%[[ARG0_W]], %[[ARG1_W]]]
+// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #[[ENCODING]]>>{%[[D0_W]], %[[D1_W]]}
+// CHECK: %[[UNSET_ENCODING:.+]] = iree_encoding.unset_encoding %[[LOAD]]
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[UNSET_ENCODING]][0, 0] [%[[ARG0_W]], %[[ARG1_W]]]
// CHECK: flow.dispatch.tensor.store %[[SLICE]], %[[OUTARG]]
// CHECK-SAME: sizes = [%[[ARG0_W]], %[[ARG1_W]]]
// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<?x?xf32>>{%[[ARG0_W]], %[[ARG1_W]]}
@@ -1997,7 +2001,7 @@
#map = affine_map<(d0, d1) -> (d1)>
#map1 = affine_map<(d0, d1) -> (d0, d1)>
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @root_on_unset_encoding(
%arg0: tensor<784x96xf32, #encoding>,
%arg1: tensor<96xf32>) -> tensor<784x96xf32> {
@@ -2017,21 +2021,22 @@
} -> tensor<784x96xf32>
util.return %4 : tensor<784x96xf32>
}
-// CHECL: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d1)>
-// CHECK: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d1)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
// CHECK: util.func public @root_on_unset_encoding
-// CHECK-SAME: %[[ARG0:.+]]: tensor<784x96xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
+// CHECK-SAME: %[[ARG0:.+]]: tensor<784x96xf32, #[[ENCODING]]>
// CHECK-SAME: %[[ARG1:.+]]: tensor<96xf32>
// CHECK: %[[DISPATCH:.+]] = flow.dispatch.workgroups(%[[ARG0]], %[[ARG1]])
-// CHECK-NEXT: %[[INARG0:.+]]: !flow.dispatch.tensor<readonly:tensor<784x96xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
+// CHECK-NEXT: %[[INARG0:.+]]: !flow.dispatch.tensor<readonly:tensor<784x96xf32, #[[ENCODING]]>>
// CHECK-SAME: %[[INARG1:.+]]: !flow.dispatch.tensor<readonly:tensor<96xf32>>
// CHECK-SAME: %[[OUTARG:.+]]: !flow.dispatch.tensor<writeonly:tensor<784x96xf32>>
// CHECK: %[[LOAD0:.+]] = flow.dispatch.tensor.load %[[INARG0]], offsets = [0, 0], sizes = [784, 96], strides = [1, 1]
-// CHECK-SAME: : !flow.dispatch.tensor<readonly:tensor<784x96xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
+// CHECK-SAME: : !flow.dispatch.tensor<readonly:tensor<784x96xf32, #[[ENCODING]]>>
// CHECK: %[[LOAD1:.+]] = flow.dispatch.tensor.load %[[INARG1]], offsets = [0], sizes = [96], strides = [1]
// CHECK-SAME: : !flow.dispatch.tensor<readonly:tensor<96xf32>>
// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<784x96xf32>
-// CHECK: %[[UNSET:.+]] = iree_encoding.unset_encoding %[[LOAD0]] : tensor<784x96xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<784x96xf32>
+// CHECK: %[[UNSET:.+]] = iree_encoding.unset_encoding %[[LOAD0]] : tensor<784x96xf32, #[[ENCODING]]> -> tensor<784x96xf32>
// CHECK: %[[GENERIC0:.+]] = linalg.generic {indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP1]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel"]} ins(%[[LOAD1]], %[[UNSET]]
// CHECK-SAME: : tensor<96xf32>, tensor<784x96xf32>) outs(%[[OUT]] : tensor<784x96xf32>)
@@ -2043,28 +2048,31 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#lhs_encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+#rhs_encoding = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+#out_encoding = #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @gemm_encoded(
- %arg0 : tensor<?x?xf32, #encoding>,
- %arg1 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>>,
- %arg2 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>)
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>> {
+ %arg0 : tensor<?x?xf32, #lhs_encoding>,
+ %arg1 : tensor<?x?xf32, #rhs_encoding>,
+ %arg2 : tensor<?x?xf32, #out_encoding>)
+ -> tensor<?x?xf32, #out_encoding> {
%0 = linalg.matmul
- ins(%arg0, %arg1
- : tensor<?x?xf32, #encoding>,
- tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>>)
- outs(%arg2 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>)
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>
- util.return %0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>
+ ins(%arg0, %arg1 : tensor<?x?xf32, #lhs_encoding>, tensor<?x?xf32, #rhs_encoding>)
+ outs(%arg2 : tensor<?x?xf32, #out_encoding>)
+ -> tensor<?x?xf32, #out_encoding>
+ util.return %0 : tensor<?x?xf32, #out_encoding>
}
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>
// CHECK: util.func public @gemm_encoded
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
-// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
-// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32, #[[LHS_ENCODING]]>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32, #[[RHS_ENCODING]]>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?xf32, #[[OUT_ENCODING]]>
// CHECK: %[[DISPATCH:.+]] = flow.dispatch.workgroups
-// CHECK-NEXT: %[[LHS_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
-// CHECK-SAME: %[[RHS_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
-// CHECK-SAME: %[[INIT_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
+// CHECK-NEXT: %[[LHS_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #[[LHS_ENCODING]]>>
+// CHECK-SAME: %[[RHS_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #[[RHS_ENCODING]]>>
+// CHECK-SAME: %[[INIT_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readwrite:tensor<?x?xf32, #[[OUT_ENCODING]]>>
// CHECK-DAG: %[[LHS:.+]] = flow.dispatch.tensor.load %[[LHS_IN]]
// CHECK-DAG: %[[RHS:.+]] = flow.dispatch.tensor.load %[[RHS_IN]]
// CHECK-DAG: %[[INIT:.+]] = flow.dispatch.tensor.load %[[INIT_IN]]
@@ -2075,34 +2083,39 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#lhs_encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+#rhs_encoding = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+#out_encoding = #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @gemm_fill_encoded(
- %arg0 : tensor<?x?xf32, #encoding>,
- %arg1 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>>)
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>> {
+ %arg0 : tensor<?x?xf32, #lhs_encoding>,
+ %arg1 : tensor<?x?xf32, #rhs_encoding>)
+ -> tensor<?x?xf32, #out_encoding> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%cst = arith.constant 0.0 : f32
- %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32, #encoding>
- %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>>
- %empty = tensor.empty(%d0, %d1) : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>
- %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>)
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>
+ %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32, #lhs_encoding>
+ %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32, #rhs_encoding>
+ %empty = tensor.empty(%d0, %d1) : tensor<?x?xf32, #out_encoding>
+ %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<?x?xf32, #out_encoding>)
+ -> tensor<?x?xf32, #out_encoding>
%0 = linalg.matmul
ins(%arg0, %arg1
- : tensor<?x?xf32, #encoding>,
- tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32]>>)
- outs(%fill : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>)
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>
- util.return %0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32]>>
+ : tensor<?x?xf32, #lhs_encoding>,
+ tensor<?x?xf32, #rhs_encoding>)
+ outs(%fill : tensor<?x?xf32, #out_encoding>)
+ -> tensor<?x?xf32, #out_encoding>
+ util.return %0 : tensor<?x?xf32, #out_encoding>
}
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>
// CHECK: util.func public @gemm_fill_encoded
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
-// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32, #[[LHS_ENCODING]]>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32, #[[RHS_ENCODING]]>
// CHECK: %[[DISPATCH:.+]] = flow.dispatch.workgroups
-// CHECK-NEXT: %[[LHS_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
-// CHECK-SAME: %[[RHS_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
-// CHECK-SAME: %[[RESULT:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32]>>>
+// CHECK-NEXT: %[[LHS_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #[[LHS_ENCODING]]>>
+// CHECK-SAME: %[[RHS_IN:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<readonly:tensor<?x?xf32, #[[RHS_ENCODING]]>>
+// CHECK-SAME: %[[RESULT:[a-zA-Z0-9]+]]: !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #[[OUT_ENCODING]]>>
// CHECK-DAG: %[[LHS:.+]] = flow.dispatch.tensor.load %[[LHS_IN]]
// CHECK-DAG: %[[RHS:.+]] = flow.dispatch.tensor.load %[[RHS_IN]]
// CHECK: %[[EMPTY:.+]] = tensor.empty
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/form_dispatch_regions.mlir b/compiler/src/iree/compiler/DispatchCreation/test/form_dispatch_regions.mlir
index 1f71258..b29f43e 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/form_dispatch_regions.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/form_dispatch_regions.mlir
@@ -137,7 +137,7 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @set_encoding_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
%arg2 : index, %arg3 : index) -> tensor<?x?xf32, #encoding> {
%cst = arith.constant 0.0 : f32
@@ -188,7 +188,7 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @set_encoding_pad_fusion(%arg0 : tensor<?x?xf32>,
%arg1 : index, %arg2 : index) -> tensor<?x?xf32, #encoding> {
%cst = arith.constant 0.0 : f32
@@ -210,7 +210,7 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @set_encoding_pad_elementwise_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
%arg2 : index, %arg3 : index) -> tensor<?x?xf32, #encoding> {
%cst = arith.constant 0.0 : f32
@@ -266,7 +266,7 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @unset_encoding_elementwise_fusion(
%arg0: tensor<?x?xf32, #encoding>,
%arg1: tensor<?xf32>,
@@ -290,8 +290,9 @@
} -> tensor<?x?xf32>
util.return %2 : tensor<?x?xf32>
}
+// CHECK: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
// CHECK-LABEL: util.func public @unset_encoding_elementwise_fusion(
-// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
+// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32, #[[$ENCODING]]>
// CHECK-SAME: %[[ARG1:.+]]: tensor<?xf32>
// CHECK: %[[RESULT:.+]] = flow.dispatch.region
// CHECK: %[[UNSET_ENCODING:.+]] = iree_encoding.unset_encoding %[[ARG0]]
@@ -302,7 +303,7 @@
// -----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @unset_encoding_slice_elementwise_fusion(
%arg0: tensor<?x?xf32, #encoding>,
%arg1: tensor<?xf32>, %arg2 : index, %arg3 : index) -> tensor<?x?xf32> {
@@ -327,8 +328,9 @@
} -> tensor<?x?xf32>
util.return %5 : tensor<?x?xf32>
}
+// CHECK: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>
// CHECK-LABEL: util.func public @unset_encoding_slice_elementwise_fusion(
-// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32]>>
+// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32, #[[$ENCODING]]>
// CHECK-SAME: %[[ARG1:.+]]: tensor<?xf32>
// CHECK: %[[RESULT0:.+]] = flow.dispatch.region
// CHECK: %[[UNSET_ENCODING:.+]] = iree_encoding.unset_encoding %[[ARG0]]
@@ -339,7 +341,7 @@
// -----
-util.func public @unpack_encoding_elementwise_fusion(
+util.func public @unpack_elementwise_fusion(
%arg0: tensor<?x?x?x?xf32>,
%arg1: tensor<?xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index
@@ -368,7 +370,7 @@
} -> tensor<?x?xf32>
util.return %1 : tensor<?x?xf32>
}
-// CHECK-LABEL: util.func public @unpack_encoding_elementwise_fusion(
+// CHECK-LABEL: util.func public @unpack_elementwise_fusion(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<?xf32>)
// CHECK: %[[RESULT:.+]] = flow.dispatch.region
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/fuse_encoding_ops_into_dispatch_regions.mlir b/compiler/src/iree/compiler/DispatchCreation/test/fuse_encoding_ops_into_dispatch_regions.mlir
index 7496b79..e6fbec2 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/fuse_encoding_ops_into_dispatch_regions.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/fuse_encoding_ops_into_dispatch_regions.mlir
@@ -25,14 +25,17 @@
util.return %2 : tensor<2x11008x128xf32, #encoding>
}
}
-
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK-LABEL: @parallel_fusion
-// CHECK: %[[DISPATCH0:.+]] = flow.dispatch.region -> (tensor<2x11008x128xf32, #iree_encoding.encoding
+// CHECK: %[[DISPATCH0:.+]] = flow.dispatch.region -> (tensor<2x11008x128xf32, #[[$ENCODING]]>)
// CHECK: %[[ADD:.+]] = linalg.generic
// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding
// CHECK: flow.return %[[SET_ENCODING]] :
// CHECK: }
-// CHECK: util.return %[[DISPATCH0]] : tensor<2x11008x128xf32, #iree_encoding.encoding
+// CHECK: util.return %[[DISPATCH0]] : tensor<2x11008x128xf32, #[[$ENCODING]]>
// -----
@@ -40,14 +43,13 @@
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
-#map4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
#encoding = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map1, #map2, #map3], round_dims_to = array<i64: 32, 32, 32>>
module {
util.func public @reduction_fusion(%arg0: tensor<2x11008x128x16xf32>) -> tensor<2x11008x128xf32, #encoding> {
%0 = tensor.empty() : tensor<2x11008x128xf32>
%1 = flow.dispatch.region -> (tensor<2x11008x128xf32>) {
%5 = linalg.generic {
- indexing_maps = [#map, #map4],
+ indexing_maps = [#map, #map3],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
ins(%arg0 : tensor<2x11008x128x16xf32>)
outs(%0 : tensor<2x11008x128xf32>) {
@@ -62,13 +64,17 @@
}
}
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK-LABEL: @reduction_fusion
-// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region -> (tensor<2x11008x128xf32, #iree_encoding.encoding
+// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region -> (tensor<2x11008x128xf32, #[[$ENCODING]]
// CHECK: %[[REDUCTION:.+]] = linalg.generic
// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[REDUCTION]]
// CHECK: flow.return %[[SET_ENCODING]] :
// CHECK: }
-// CHECK: util.return %[[DISPATCH]] : tensor<2x11008x128xf32, #iree_encoding.encoding
+// CHECK: util.return %[[DISPATCH]] : tensor<2x11008x128xf32, #[[$ENCODING]]>
// -----
@@ -97,13 +103,17 @@
}
}
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK-LABEL: @transpose_fusion
-// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region -> (tensor<2x11008x128xf32, #iree_encoding.encoding
+// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region -> (tensor<2x11008x128xf32, #[[$ENCODING]]>
// CHECK: %[[TRANSPOSE:.+]] = linalg.generic
// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[TRANSPOSE]]
// CHECK: flow.return %[[SET_ENCODING]]
// CHECK: }
-// CHECK: util.return %[[DISPATCH]] : tensor<2x11008x128xf32, #iree_encoding.encoding
+// CHECK: util.return %[[DISPATCH]] : tensor<2x11008x128xf32, #[[$ENCODING]]>
// -----
@@ -133,12 +143,16 @@
}
}
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK-LABEL: @fusion_dynamic
// CHECK-SAME: {{.+}}: tensor<?x?x?xf32>, %[[D0:.+]]: index, %[[D1:.+]]: index, %[[D2:.+]]: index)
-// CHECK: %[[DISPATCH0:.+]] = flow.dispatch.region -> (tensor<?x?x?xf32, #iree_encoding.encoding
+// CHECK: %[[DISPATCH0:.+]] = flow.dispatch.region -> (tensor<?x?x?xf32, #[[$ENCODING]]>
// CHECK-SAME: {%[[D0]], %[[D1]], %[[D2]]}
// CHECK: %[[ADD:.+]] = linalg.generic
// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding
// CHECK: flow.return %[[SET_ENCODING]] :
// CHECK: }
-// CHECK: util.return %[[DISPATCH0]] : tensor<?x?x?xf32, #iree_encoding.encoding
+// CHECK: util.return %[[DISPATCH0]] : tensor<?x?x?xf32, #[[$ENCODING]]>
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir b/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
index a6c3012..5bd5642 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
@@ -31,13 +31,19 @@
}
}
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[$RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[$OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK-LABEL: @hoist_matmul_encodings
// CHECK-SAME: (%[[ARG0:.+]]: tensor<2x128x64xf32>, %[[ARG1:.+]]: tensor<2x11008x128xf32>)
-// CHECK-DAG: %[[SET_ENCODING0:.+]] = iree_encoding.set_encoding %[[ARG0]] : tensor<2x128x64xf32> -> tensor<2x128x64xf32, #iree_encoding.encoding
-// CHECK-DAG: %[[SET_ENCODING1:.+]] = iree_encoding.set_encoding %[[ARG1]] : tensor<2x11008x128xf32> -> tensor<2x11008x128xf32, #iree_encoding.encoding
+// CHECK-DAG: %[[SET_ENCODING0:.+]] = iree_encoding.set_encoding %[[ARG0]] : tensor<2x128x64xf32> -> tensor<2x128x64xf32, #[[$LHS_ENCODING]]>
+// CHECK-DAG: %[[SET_ENCODING1:.+]] = iree_encoding.set_encoding %[[ARG1]] : tensor<2x11008x128xf32> -> tensor<2x11008x128xf32, #[[$RHS_ENCODING]]>
// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region -> (tensor<2x11008x64xf32>) {
// CHECK: %[[MATMUL:.+]] = linalg.generic {{.*}} ins(%[[SET_ENCODING0]], %[[SET_ENCODING1]]
-// CHECK: %[[UNSET_ENCODING1:.+]] = iree_encoding.unset_encoding %[[MATMUL]] : tensor<2x11008x64xf32, #iree_encoding.encoding
+// CHECK: %[[UNSET_ENCODING1:.+]] = iree_encoding.unset_encoding %[[MATMUL]] : tensor<2x11008x64xf32, #[[$OUT_ENCODING]]>
// CHECK: flow.return %[[UNSET_ENCODING1]] : tensor<2x11008x64xf32>
// CHECK: }
// CHECK: util.return %[[DISPATCH]] : tensor<2x11008x64xf32>
@@ -46,7 +52,10 @@
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
-#encoding = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>], round_dims_to = array<i64: 32, 32, 32>>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#encoding = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map2, #map3, #map4], round_dims_to = array<i64: 32, 32, 32>>
util.func public @bubble_through_dequant(
%arg0: tensor<2x11008x128xi8>, %arg1: tensor<2x11008xf32>, %arg2: tensor<2x11008xf32>) -> tensor<2x11008x128xf32, #encoding> {
%6 = flow.dispatch.region -> (tensor<2x11008x128xf32, #encoding>) {
@@ -69,19 +78,22 @@
util.return %6 : tensor<2x11008x128xf32, #encoding>
}
-// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
-// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
-// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
-// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
-// CHECK-DAG: #[[$MAP4:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[$ENCODING_IBMAP:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP2]]], bcast_map = #[[MAP3]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[$ENCODING_BMAP:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP2]]], bcast_map = #[[MAP4]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK-LABEL: @bubble_through_dequant
// CHECK-SAME: %[[ARG0:.+]]: tensor<2x11008x128xi8>,
// CHECK-SAME: %[[ARG1:.+]]: tensor<2x11008xf32>, %[[ARG2:.+]]: tensor<2x11008xf32>
-// CHECK-DAG: %[[SET_ENCODING0:.+]] = iree_encoding.set_encoding %[[ARG0]] : {{.*}} bcast_map = #[[$MAP4]]
-// CHECK-DAG: %[[SET_ENCODING1:.+]] = iree_encoding.set_encoding %[[ARG1]] : {{.*}} bcast_map = #[[$MAP3]]
-// CHECK-DAG: %[[SET_ENCODING2:.+]] = iree_encoding.set_encoding %[[ARG2]] : {{.*}} bcast_map = #[[$MAP3]]
+// CHECK-DAG: %[[SET_ENCODING0:.+]] = iree_encoding.set_encoding %[[ARG0]] : tensor<2x11008x128xi8> -> tensor<2x11008x128xi8, #[[$ENCODING_IBMAP]]>
+// CHECK-DAG: %[[SET_ENCODING1:.+]] = iree_encoding.set_encoding %[[ARG1]] : tensor<2x11008xf32> -> tensor<2x11008xf32, #[[$ENCODING_BMAP]]>
+// CHECK-DAG: %[[SET_ENCODING2:.+]] = iree_encoding.set_encoding %[[ARG2]] : tensor<2x11008xf32> -> tensor<2x11008xf32, #[[$ENCODING_BMAP]]>
// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region
-// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x11008x128xf32, #iree_encoding.encoding
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x11008x128xf32, #[[$ENCODING]]>
// CHECK: %[[DEQUANT:.+]] = linalg.generic {{.*}} ins(%[[SET_ENCODING0]], %[[SET_ENCODING1]], %[[SET_ENCODING2]] : {{.*}} outs(%[[INIT]] :
// CHECK: flow.return %[[DEQUANT]]
// CHECK: }
@@ -110,16 +122,18 @@
util.return %6 : tensor<2x11008x128xf32, #encoding>
}
-// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
-// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
-// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
-// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
-// CHECK-DAG: #[[$MAP4:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[$ENCODING_BMAP:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP2]]], bcast_map = #[[MAP3]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK-LABEL: @bubble_through_broadcast
// CHECK-SAME: %[[ARG0:.+]]: tensor<11008x128xf32>
-// CHECK-DAG: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[ARG0]] : {{.*}} bcast_map = #[[$MAP3]]
+// CHECK-DAG: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[ARG0]] : tensor<11008x128xf32> -> tensor<11008x128xf32, #[[$ENCODING_BMAP]]>
// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region
-// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x11008x128xf32, #iree_encoding.encoding
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x11008x128xf32, #[[$ENCODING]]>
// CHECK: %[[BROADCAST:.+]] = linalg.generic {{.*}} ins(%[[SET_ENCODING]] : {{.*}} outs(%[[INIT]] :
// CHECK: flow.return %[[BROADCAST]]
// CHECK: }
@@ -180,7 +194,11 @@
util.return %0#0, %0#1 : tensor<?x?x?xf32>, tensor<?x?x?xf32, #encoding>
}
}
-
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK-LABEL: @hoist_dynamic
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>, %[[D0:.+]]: index, %[[D1:.+]]: index, %[[D2:.+]]: index)
// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region -> (tensor<?x?x?xf32>{%[[D0]], %[[D1]], %[[D2]]})
@@ -188,5 +206,5 @@
// CHECK: %[[ADD:.+]] = linalg.generic {{.*}} ins(%[[ARG0]], %[[ARG0]] : {{.*}} outs(%[[INIT]] :
// CHECK: flow.return %[[ADD]]
// CHECK: }
-// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[DISPATCH]]
+// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[DISPATCH]] : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #[[$ENCODING]]>
// CHECK: util.return %[[DISPATCH]], %[[SET_ENCODING]]
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/set_encoding.mlir b/compiler/src/iree/compiler/DispatchCreation/test/set_encoding.mlir
index 617aec0..1d2aa3b 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/set_encoding.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/set_encoding.mlir
@@ -9,20 +9,23 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_f32f32f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<100x250xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<250x500xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<100x250xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x250xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<250x500xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<250x500xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[MATMUL:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
-// CHECK: %[[RESULT:.+]] = iree_encoding.unset_encoding %[[MATMUL]]
+// CHECK: %[[RESULT:.+]] = iree_encoding.unset_encoding %[[MATMUL]] : tensor<100x500xf32, #[[OUT_ENCODING]]> -> tensor<100x500xf32>
// CHECK: util.return %[[RESULT]]
// -----
@@ -36,22 +39,25 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_f32f32f32_dynamic(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>, %[[ARG1:.+]]: tensor<?x?xf32>, %[[ARG2:.+]]: tensor<?x?xf32>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<?x?xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<?x?xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<?x?xf32, #[[OUT_ENCODING]]>
// CHECK: %[[MATMUL:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG2]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG2]], %[[C1]]
-// CHECK: %[[RESULT:.+]] = iree_encoding.unset_encoding %[[MATMUL]]{{.+}} -> tensor<?x?xf32>{%[[D0]], %[[D1]]}
+// CHECK: %[[RESULT:.+]] = iree_encoding.unset_encoding %[[MATMUL]] : tensor<?x?xf32, #[[OUT_ENCODING]]> -> tensor<?x?xf32>{%[[D0]], %[[D1]]}
// CHECK: util.return %[[RESULT]]
// -----
@@ -65,16 +71,19 @@
// 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: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_i8i8i32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<100x250xi8>
// CHECK-SAME: %[[ARG1:.+]]: tensor<250x500xi8>
// CHECK-SAME: %[[ARG2:.+]]: tensor<100x500xi32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<100x250xi8, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x250xi8, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<250x500xi8, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<250x500xi8, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<100x500xi32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x500xi32, #[[OUT_ENCODING]]>
// CHECK: %[[MATMUL:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -92,16 +101,19 @@
// 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: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_f16f16f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<100x250xf16>
// CHECK-SAME: %[[ARG1:.+]]: tensor<250x500xf16>
// CHECK-SAME: %[[ARG2:.+]]: tensor<100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<100x250xf16, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x250xf16, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<250x500xf16, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<250x500xf16, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[MATMUL:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -119,16 +131,19 @@
// 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: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_f16f16f16(
// CHECK-SAME: %[[ARG0:.+]]: tensor<100x250xf16>
// CHECK-SAME: %[[ARG1:.+]]: tensor<250x500xf16>
// CHECK-SAME: %[[ARG2:.+]]: tensor<100x500xf16>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<100x250xf16, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x250xf16, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<250x500xf16, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<250x500xf16, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<100x500xf16, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x500xf16, #[[OUT_ENCODING]]>
// CHECK: %[[MATMUL:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -146,16 +161,19 @@
// 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: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_bf16bf16f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<100x250xbf16>
// CHECK-SAME: %[[ARG1:.+]]: tensor<250x500xbf16>
// CHECK-SAME: %[[ARG2:.+]]: tensor<100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<100x250xbf16, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x250xbf16, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<250x500xbf16, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<250x500xbf16, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[MATMUL:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -173,16 +191,19 @@
// 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: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_bf16bf16bf16(
// CHECK-SAME: %[[ARG0:.+]]: tensor<100x250xbf16>
// CHECK-SAME: %[[ARG1:.+]]: tensor<250x500xbf16>
// CHECK-SAME: %[[ARG2:.+]]: tensor<100x500xbf16>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<100x250xbf16, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x250xbf16, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<250x500xbf16, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<250x500xbf16, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<100x500xbf16, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x500xbf16, #[[OUT_ENCODING]]>
// CHECK: %[[MATMUL:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -200,16 +221,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @batch_matmul_f32f32f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<64x100x250xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<64x250x500xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<64x100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<64x100x250xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x250xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<64x250x500xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x250x500xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<64x100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATMUL:.+]] = linalg.batch_matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -227,24 +251,27 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @batch_matmul_f32f32f32_dynamic(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?x?xf32>, %[[ARG2:.+]]: tensor<?x?x?xf32>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<?x?x?xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<?x?x?xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<?x?x?xf32, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATMUL:.+]] = linalg.batch_matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG2]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG2]], %[[C1]]
// CHECK-DAG: %[[D2:.+]] = tensor.dim %[[ARG2]], %[[C2]]
-// CHECK: %[[RESULT:.+]] = iree_encoding.unset_encoding %[[BATCH_MATMUL]]{{.+}} -> tensor<?x?x?xf32>{%[[D0]], %[[D1]], %[[D2]]}
+// CHECK: %[[RESULT:.+]] = iree_encoding.unset_encoding %[[BATCH_MATMUL]] : tensor<?x?x?xf32, #[[OUT_ENCODING]]> -> tensor<?x?x?xf32>{%[[D0]], %[[D1]], %[[D2]]}
// CHECK: util.return %[[RESULT]]
// -----
@@ -258,16 +285,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @batch_matmul_f16f16f16(
// CHECK-SAME: %[[ARG0:.+]]: tensor<64x100x250xf16>
// CHECK-SAME: %[[ARG1:.+]]: tensor<64x250x500xf16>
// CHECK-SAME: %[[ARG2:.+]]: tensor<64x100x500xf16>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<64x100x250xf16, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x250xf16, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<64x250x500xf16, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x250x500xf16, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<64x100x500xf16, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x500xf16, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATMUL:.+]] = linalg.batch_matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -285,16 +315,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @batch_matmul_f16f16f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<64x100x250xf16>
// CHECK-SAME: %[[ARG1:.+]]: tensor<64x250x500xf16>
// CHECK-SAME: %[[ARG2:.+]]: tensor<64x100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<64x100x250xf16, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x250xf16, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<64x250x500xf16, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x250x500xf16, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<64x100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATMUL:.+]] = linalg.batch_matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -312,16 +345,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @batch_matmul_bf16bf16bf16(
// CHECK-SAME: %[[ARG0:.+]]: tensor<64x100x250xbf16>
// CHECK-SAME: %[[ARG1:.+]]: tensor<64x250x500xbf16>
// CHECK-SAME: %[[ARG2:.+]]: tensor<64x100x500xbf16>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<64x100x250xbf16, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x250xbf16, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<64x250x500xbf16, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x250x500xbf16, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<64x100x500xbf16, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [bf16, bf16, bf16], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x500xbf16, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATMUL:.+]] = linalg.batch_matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -339,16 +375,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @batch_matmul_bf16bf16f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<64x100x250xbf16>
// CHECK-SAME: %[[ARG1:.+]]: tensor<64x250x500xbf16>
// CHECK-SAME: %[[ARG2:.+]]: tensor<64x100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<64x100x250xbf16, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x250xbf16, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<64x250x500xbf16, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x250x500xbf16, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<64x100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [bf16, bf16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATMUL:.+]] = linalg.batch_matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -366,16 +405,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @batch_matmul_i8i8i32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<64x100x250xi8>
// CHECK-SAME: %[[ARG1:.+]]: tensor<64x250x500xi8>
// CHECK-SAME: %[[ARG2:.+]]: tensor<64x100x500xi32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<64x100x250xi8, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x250xi8, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<64x250x500xi8, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x250x500xi8, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<64x100x500xi32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<64x100x500xi32, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATMUL:.+]] = linalg.batch_matmul
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -393,16 +435,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d1, d0)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>
// CHECK: util.func public @vecmat_f32f32f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<250xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<250x500xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<250xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>>
+// CHECK-SAME: tensor<250xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<250x500xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>>
+// CHECK-SAME: tensor<250x500xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>>
+// CHECK-SAME: tensor<500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[VECMAT:.+]] = linalg.vecmat
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -420,16 +465,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d1)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>
// CHECK: util.func public @matvec_f32f32f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<100x250xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<250xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<100xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<100x250xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>>
+// CHECK-SAME: tensor<100x250xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<250xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>>
+// CHECK-SAME: tensor<250xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<100xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>>
+// CHECK-SAME: tensor<100xf32, #[[OUT_ENCODING]]>
// CHECK: %[[MATVEC:.+]] = linalg.matvec
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -447,16 +495,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>
// CHECK: util.func public @batch_vecmat_f32f32f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<3x250xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<3x250x500xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<3x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<3x250xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>>
+// CHECK-SAME: tensor<3x250xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<3x250x500xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>>
+// CHECK-SAME: tensor<3x250x500xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<3x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>>
+// CHECK-SAME: tensor<3x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[VECMAT:.+]] = linalg.batch_vecmat
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -474,16 +525,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>
+// CHECK-DAG: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>
+// CHECK-DAG: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>
// CHECK: util.func public @batch_matvec_f32f32f32_dynamic(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?xf32>, %[[ARG2:.+]]: tensor<?x?xf32>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<?x?x?xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>>
+// CHECK-SAME: tensor<?x?x?xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>>
+// CHECK-SAME: tensor<?x?xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 1, 32>>>
+// CHECK-SAME: tensor<?x?xf32, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATVEC:.+]] = linalg.batch_matvec
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -494,25 +548,26 @@
// -----
-util.func public @fold_fill_with_set_encoding(%arg0 : index, %arg1 : index)
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>> {
+#encoding = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>
+util.func public @fold_fill_with_set_encoding(%arg0 : index, %arg1 : index) -> tensor<?x?xf32, #encoding> {
%cst = arith.constant 0.0 : f32
%0 = tensor.empty(%arg0, %arg1) : tensor<?x?xf32>
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
%2 = iree_encoding.set_encoding %1 : tensor<?x?xf32>
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>>
- util.return %2 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>>
+ -> tensor<?x?xf32, #encoding>
+ util.return %2 : tensor<?x?xf32, #encoding>
}
+// CHECK-DAG: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @fold_fill_with_set_encoding(
-// CHECK: %[[EMPTY:.+]] = tensor.empty(%{{.+}}, %{{.+}}) : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK: %[[EMPTY:.+]] = tensor.empty(%{{.+}}, %{{.+}}) : tensor<?x?xf32, #[[ENCODING]]>
// CHECK: %[[FILL:.+]] = linalg.fill
-// CHECK-SAME: outs(%[[EMPTY]] : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0 : i64, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<?x?xf32, #[[ENCODING]]>)
// CHECK: util.return %[[FILL]]
// -----
-util.func public @fold_fill_with_tensor_pad(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index)
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>> {
+#encoding = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>
+util.func public @fold_fill_with_tensor_pad(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index) -> tensor<?x?xf32, #encoding> {
%cst = arith.constant 0.0 : f32
%0 = tensor.empty(%arg0, %arg1) : tensor<?x?xf32>
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
@@ -520,13 +575,13 @@
^bb0(%b0: index, %b1 : index):
tensor.yield %cst : f32
} : tensor<?x?xf32> to tensor<?x?xf32>
- %3 = iree_encoding.set_encoding %2 : tensor<?x?xf32>
- -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>>
- util.return %3 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>>
+ %3 = iree_encoding.set_encoding %2 : tensor<?x?xf32> -> tensor<?x?xf32, #encoding>
+ util.return %3 : tensor<?x?xf32, #encoding>
}
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @fold_fill_with_tensor_pad(
// CHECK: %[[EMPTY:.+]] = tensor.empty(
-// CHECK-SAME: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2 : i64, op_type = matmul, element_types = [f32, f32, f32], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<?x?xf32, #[[ENCODING]]>
// CHECK: %[[FILL:.+]] = linalg.fill
// CHECK-SAME: outs(%[[EMPTY]] :
// CHECK: util.return %[[FILL]]
@@ -580,14 +635,34 @@
outs(%arg2 : tensor<64x100x500xf32>) -> tensor<64x100x500xf32>
util.return %1 : tensor<64x100x500xf32>
}
+// CHECK: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK: #[[LHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK: #[[RHS_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK: #[[OUT_ENCODING:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK: util.func public @batch_matmul_truncf_f16f16f32
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]
+// CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<64x250x500xf16>
+// CHECK-DAG: %[[TRUNCF0:.+]] = arith.truncf %[[ARG0]] : tensor<64x100x250xf32> to tensor<64x100x250xf16>
+// CHECK: %[[TRUNCF1:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[ARG1]] : tensor<64x250x500xf32>)
+// CHECK-SAME: outs(%[[INIT]] : tensor<64x250x500xf16>)
+// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[TRUNCF0]]
+// CHECK-SAME: tensor<64x100x250xf16, #[[LHS_ENCODING]]>
+// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[TRUNCF1]]
+// CHECK-SAME: tensor<64x250x500xf16, #[[RHS_ENCODING]]>
+// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
+// CHECK-SAME: tensor<64x100x500xf32, #[[OUT_ENCODING]]>
+// CHECK: %[[MATMUL:.+]] = linalg.batch_matmul
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: %[[RESULT:.+]] = iree_encoding.unset_encoding %[[MATMUL]] : tensor<64x100x500xf32, #[[OUT_ENCODING]]> -> tensor<64x100x500xf32>
+// CHECK: util.return %[[RESULT]]
-// CHECK: util.func public @batch_matmul_truncf_f16f16f32(%[[ARG0:.+]]: tensor<64x100x250xf32>, %[[ARG1:.+]]: tensor<64x250x500xf32>
-// CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<64x250x500xf16>
-// CHECK-DAG: arith.truncf %[[ARG0]] : tensor<64x100x250xf32> to tensor<64x100x250xf16>
-// CHECK: linalg.generic
-// CHECK-SAME: ins(%[[ARG1]] : tensor<64x250x500xf32>)
-// CHECK-SAME: outs(%[[INIT]] : tensor<64x250x500xf16>)
-// CHECK: element_types = [f16, f16, f32]
// -----
@@ -603,10 +678,13 @@
// 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: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_casted_from_i1_f32f32f32
-// CHECK: set_encoding {{.+}} tensor<64x256xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
-// CHECK: set_encoding {{.+}} tensor<256x128xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
-// CHECK: set_encoding {{.+}} tensor<64x128xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK: set_encoding {{.+}} tensor<64x256xf32, #[[LHS_ENCODING]]>
+// CHECK: set_encoding {{.+}} tensor<256x128xf32, #[[RHS_ENCODING]]>
+// CHECK: set_encoding {{.+}} tensor<64x128xf32, #[[OUT_ENCODING]]>
// -----
@@ -631,10 +709,13 @@
// 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: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_generic_casted_from_i1_f32f32f32
-// CHECK: set_encoding {{.+}} tensor<64x256xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
-// CHECK: set_encoding {{.+}} tensor<256x128xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
-// CHECK: set_encoding {{.+}} tensor<64x128xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK: set_encoding {{.+}} tensor<64x256xf32, #[[LHS_ENCODING]]>
+// CHECK: set_encoding {{.+}} tensor<256x128xf32, #[[RHS_ENCODING]]>
+// CHECK: set_encoding {{.+}} tensor<64x128xf32, #[[OUT_ENCODING]]>
// -----
@@ -647,13 +728,16 @@
// 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: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 2, 32, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 2, 32, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 2, 32, 32>>
// CHECK: util.func public @matmul_f32f32f32_narrow_M(
// CHECK: iree_encoding.set_encoding
-// CHECK-SAME: tensor<2x250xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 2, 32, 32>>>
+// CHECK-SAME: tensor<2x250xf32, #[[LHS_ENCODING]]>
// CHECK: iree_encoding.set_encoding
-// CHECK-SAME: tensor<250x500xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 2, 32, 32>>>
+// CHECK-SAME: tensor<250x500xf32, #[[RHS_ENCODING]]>
// CHECK: iree_encoding.set_encoding
-// CHECK-SAME: tensor<2x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 2, 32, 32>>>
+// CHECK-SAME: tensor<2x500xf32, #[[OUT_ENCODING]]>
// CHECK: linalg.matmul
// -----
@@ -667,13 +751,16 @@
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 2, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 2, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 2, 32>>
// CHECK: util.func public @batch_matmul_f32f32f32_narrow_MN(
// CHECK: iree_encoding.set_encoding
-// CHECK-SAME: tensor<64x4x250xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 2, 32>>>
+// CHECK-SAME: tensor<64x4x250xf32, #[[LHS_ENCODING]]>
// CHECK: iree_encoding.set_encoding
-// CHECK-SAME: tensor<64x250x2xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 2, 32>>>
+// CHECK-SAME: tensor<64x250x2xf32, #[[RHS_ENCODING]]>
// CHECK: iree_encoding.set_encoding
-// CHECK-SAME: tensor<64x4x2xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]], round_dims_to = array<i64: 32, 2, 32>>>
+// CHECK-SAME: tensor<64x4x2xf32, #[[OUT_ENCODING]]>
// CHECK: linalg.batch_matmul
// -----
@@ -688,16 +775,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_transpose_a_f32f32f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<250x100xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<250x500xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<250x100xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<250x100xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<250x500xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<250x500xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[MATMUL:.+]] = linalg.matmul_transpose_a
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -715,16 +805,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @matmul_transpose_b_f32f32f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<100x250xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<500x250xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<100x250xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x250xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<500x250xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<500x250xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[MATMUL:.+]] = linalg.matmul_transpose_b
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -742,16 +835,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @batch_matmul_transpose_a_f32f32f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<2x250x100xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<2x250x500xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<2x100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<2x250x100xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<2x250x100xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<2x250x500xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<2x250x500xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<2x100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<2x100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATMUL:.+]] = linalg.batch_matmul_transpose_a
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -769,16 +865,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @batch_matmul_transpose_b_f32f32f32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<2x100x250xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<2x500x250xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<2x100x500xf32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<2x100x250xf32, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<2x100x250xf32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<2x500x250xf32, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<2x500x250xf32, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<2x100x500xf32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 32, 32, 32>>>
+// CHECK-SAME: tensor<2x100x500xf32, #[[OUT_ENCODING]]>
// CHECK: %[[BATCH_MATMUL:.+]] = linalg.batch_matmul_transpose_b
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[OUTS]] :
@@ -803,16 +902,19 @@
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [i16, ui4, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [i16, ui4, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [i16, ui4, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>
// CHECK: util.func public @generic_batch_vecmat_transposed_i16u4i32(
// CHECK-SAME: %[[ARG0:.+]]: tensor<32x128xi16>
// CHECK-SAME: %[[ARG1:.+]]: tensor<4096x32x128xi4>
// CHECK-SAME: %[[ARG2:.+]]: tensor<4096x32xi32>
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[ARG0]]
-// CHECK-SAME: tensor<32x128xi16, #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [i16, ui4, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>>
+// CHECK-SAME: tensor<32x128xi16, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %[[ARG1]]
-// CHECK-SAME: tensor<4096x32x128xi4, #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [i16, ui4, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>>
+// CHECK-SAME: tensor<4096x32x128xi4, #[[RHS_ENCODING]]>
// CHECK: %[[OUTS:.+]] = iree_encoding.set_encoding %[[ARG2]]
-// CHECK-SAME: tensor<4096x32xi32> -> tensor<4096x32xi32, #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [i16, ui4, i32], user_indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], round_dims_to = array<i64: 1, 32, 32>>>
+// CHECK-SAME: tensor<4096x32xi32> -> tensor<4096x32xi32, #[[OUT_ENCODING]]>
// CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
@@ -956,6 +1058,9 @@
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[LHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [i32, i32, i32], user_indexing_maps = [#[[MAP2]], #[[MAP3]], #[[MAP4]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[RHS_encoding:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [i32, i32, i32], user_indexing_maps = [#[[MAP2]], #[[MAP3]], #[[MAP4]]], round_dims_to = array<i64: 32, 32, 32>>
+// CHECK-DAG: #[[OUT_encoding:.+]] = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [i32, i32, i32], user_indexing_maps = [#[[MAP2]], #[[MAP3]], #[[MAP4]]], round_dims_to = array<i64: 32, 32, 32>>
// CHECK: util.func public @broadcasting_dequant_op(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
@@ -967,21 +1072,10 @@
// CHECK: %{{.+}} = flow.dispatch.region
// CHECK: %[[BCAST:.+]] = linalg.generic
// CHECK: %[[LHS:.+]] = iree_encoding.set_encoding %[[BCAST]] : tensor<?x?x?xi32>
-// CHECK-SAME: -> tensor<?x?x?xi32, #iree_encoding.encoding
-// CHECK-SAME: operand_index = 0 : index
-// CHECK-SAME: element_types = [i32, i32, i32]
-// CHECK-SAME: user_indexing_maps = [#[[MAP2]], #[[MAP3]], #[[MAP4]]]
-// CHECK-SAME: round_dims_to = array<i64: 32, 32, 32>
+// CHECK-SAME: -> tensor<?x?x?xi32, #[[LHS_ENCODING]]>
// CHECK: %[[RHS:.+]] = iree_encoding.set_encoding %{{.+}} : tensor<?x?x?xi32>
-// CHECK-SAME: operand_index = 1 : index
-// CHECK-SAME: element_types = [i32, i32, i32]
-// CHECK-SAME: user_indexing_maps = [#[[MAP2]], #[[MAP3]], #[[MAP4]]]
-// CHECK-SAME: round_dims_to = array<i64: 32, 32, 32>
-// CHECK: %[[INIT:.+]] = tensor.empty({{.+}}) : tensor<?x?x?xi32, #iree_encoding.encoding
-// CHECK-SAME: operand_index = 2 : index
-// CHECK-SAME: element_types = [i32, i32, i32]
-// CHECK-SAME: user_indexing_maps = [#[[MAP2]], #[[MAP3]], #[[MAP4]]]
-// CHECK-SAME: round_dims_to = array<i64: 32, 32, 32>
+// CHECK-SAME: -> tensor<?x?x?xi32, #[[RHS_ENCODING]]>
+// CHECK: %[[INIT:.+]] = tensor.empty({{.+}}) : tensor<?x?x?xi32, #[[OUT_ENCODING]]>
// CHECK: %[[FILL:.+]] = linalg.fill ins({{.+}}) outs(%[[INIT]]
// CHECK: %[[GEMM:.+]] = linalg.batch_matmul_transpose_b
// CHECK-SAME: ins(%[[LHS]], %[[RHS]]
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/tensor_pad_to_tensor_insert_slice.mlir b/compiler/src/iree/compiler/DispatchCreation/test/tensor_pad_to_tensor_insert_slice.mlir
index d990f8b..ec6efa7 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/tensor_pad_to_tensor_insert_slice.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/tensor_pad_to_tensor_insert_slice.mlir
@@ -79,7 +79,7 @@
// ----
-#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>
+#encoding = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32]>
util.func public @dispatch_dispatch_0_generic_512x1024_f32(
%arg0: !flow.dispatch.tensor<readonly:tensor<512x1024xf32>>,
%arg1: index, %arg2: index, %arg3: index, %arg4: index,
@@ -99,8 +99,9 @@
flow.dispatch.tensor.store %11, %2, offsets = [0, 0], sizes = [%0, %1], strides = [1, 1] : tensor<?x?xf32, #encoding> -> !flow.dispatch.tensor<writeonly:tensor<?x?xf32, #encoding>>{%0, %1}
util.return
}
-
-// CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load
-// CHECK: %[[PAD:.+]] = tensor.pad %[[LOAD]] low
-// CHECK: %[[ENCODE:.+]] = iree_encoding.set_encoding %[[PAD]] : tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding
-// CHECK: flow.dispatch.tensor.store %[[ENCODE]],
+// CHECK: #[[ENCODING:.+]] = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32]>
+// CHECK: util.func public @dispatch_dispatch_0_generic_512x1024_f32
+// CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load
+// CHECK: %[[PAD:.+]] = tensor.pad %[[LOAD]] low
+// CHECK: %[[ENCODE:.+]] = iree_encoding.set_encoding %[[PAD]] : tensor<?x?xf32> -> tensor<?x?xf32, #[[ENCODING]]>
+// CHECK: flow.dispatch.tensor.store %[[ENCODE]],
