[Encoding] Introduce matmul_k encoding. (#20484)

The attribute is specialized for tracking the matmul reduction
dimensions only. The encoded dimensions are the reduction dimensions
that will be used by matmuls.

The below message will be removed before we land the PR.

Please check #20485, if you want to
get the full picture about how the new encoding is set up.

---------

Signed-off-by: hanhanW <[email protected]>

Author: hanhanW

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td

@@ -138,6 +138,55 @@ def EncodingAttr :
   let genVerifyDecl = 1;
 }
 
+//===---------------------------------------------------------------------===//
+// encoding.matmul_k
+//===---------------------------------------------------------------------===//
+
+def MatmulKAttr : IREEEncoding_Attr<"MatmulK"> {
+  let mnemonic = "matmul_k";
+  let summary = [{An attribute that tracks reduction dimensions for matmul}];
+  let description = [{
+    The attribute is specialized for tracking the matmul reduction dimensions
+    only. The encoded dimensions are the reduction dimensions that will be
+    further used in a matmul.
+
+    Example: linalg.matmul ins(%lhs, %rhs) outs(%acc), where the indexing maps
+             are:
+
+      lhs: (M, N, K) -> (M, K)
+      rhs: (M, N, K) -> (K, N)
+      acc: (M, N, K) -> (M, N)
+
+    The encoding for the lhs is iree_encoding.matmul_k<k_dims = [1]>.
+    The encoding for the rhs is iree_encoding.matmul_k<k_dims = [0]>.
+    The encoding for the acc is iree_encoding.matmul_k<k_dims = []>.
+
+    Example: linalg.matmul_transpose_b ins(%lhs, %rhs) outs(%acc), where the
+             indexing maps are:
+
+      lhs: (M, N, K) -> (M, K)
+      rhs: (M, N, K) -> (N, K)
+      acc: (M, N, K) -> (M, N)
+
+    The encoding for the lhs is iree_encoding.matmul_k<k_dims = [1]>.
+    The encoding for the rhs is iree_encoding.matmul_k<k_dims = [1]>.
+    The encoding for the acc is iree_encoding.matmul_k<k_dims = []>.
+
+    There is no difference between static cases and dynamic cases because the
+    dimension sizes are not encoded.
+
+    Any encoding propagation that transforms reduction dimensions could result
+    in undefined behavior, because it does not encode the transformations. The
+    information is missing in this context.
+  }];
+
+  let assemblyFormat = "`<` struct(params) `>`";
+
+  let parameters = (ins
+    "DenseI32ArrayAttr":$k_dims
+  );
+}
+
 //===---------------------------------------------------------------------===//
 // encoding.pad_encoding_layout
 //===---------------------------------------------------------------------===//

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingDialect.cpp

@@ -30,14 +30,14 @@ 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.
+  // 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, TestingEncodingAttr, UnknownEncodingAttr>(
-            attr)) {
+    if (llvm::isa<EncodingAttr, MatmulKAttr, TestingEncodingAttr,
+                  UnknownEncodingAttr>(attr)) {
       os << "encoding";
       return AliasResult::OverridableAlias;
     }

compiler/src/iree/compiler/Dialect/Encoding/IR/test/roundtrip.mlir

@@ -243,3 +243,14 @@ func.func @testing_encoding_with_layouts(%arg0: tensor<?x?xf32, #encoding>) -> t
 // CHECK:      func.func @testing_encoding_with_layouts(
 // CHECK-SAME:   %[[ARG0:.+]]: tensor<?x?xf32, #[[ENCODING]]>
 // CHECK         return %[[ARG0]]
+
+// -----
+
+#encoding = #iree_encoding.matmul_k<k_dims = [1]>
+func.func @matmul_k_encoding(%arg0: tensor<?x?xf32, #encoding>) -> tensor<?x?xf32, #encoding> {
+  return %arg0 : tensor<?x?xf32, #encoding>
+}
+// CHECK:     #[[ENCODING:.+]] = #iree_encoding.matmul_k<k_dims = [1]>
+// CHECK:      func.func @matmul_k_encoding(
+// CHECK-SAME:   %[[ARG0:.+]]: tensor<?x?xf32, #[[ENCODING]]>
+// CHECK         return %[[ARG0]]