[Codegen][CPU] Materialize data-tiled encodings to inner_tiled.

compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingPatterns.cpp

@@ -662,6 +662,10 @@ class MaterializeLinalgOp
     }
     Operation *newOp =
         layoutAttr.lowerOp(rewriter, op, convertedResTypes, operands);
+    if (!newOp) {
+      return op->emitOpError(
+          "unsupported operation for encoding materialization");
+    }
     rewriter.replaceOp(op, newOp->getResults());
     return success();
   }

compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_x86_64.mlir

@@ -148,6 +148,98 @@ func.func @pack_gemm_fill_dynamic(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf3
 
 // -----
 
+// With `enable_inner_tiled` and +avx512f (TileMxNxK 16x16x1), matmul lowers to
+// iree_codegen.inner_tiled instead of linalg.mmt4d.
+
+#map_it = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map_it1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map_it2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#encoding_lhs_it = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_it, #map_it1, #map_it2], iteration_sizes = [?, ?, ?]>
+#encoding_rhs_it = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_it, #map_it1, #map_it2], iteration_sizes = [?, ?, ?]>
+#encoding_result_it = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_it, #map_it1, #map_it2], iteration_sizes = [?, ?, ?]>
+func.func @pack_gemm_fill_dynamic_inner_tiled_avx512(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %m: index, %n: index, %k: index) -> tensor<?x?xf32> attributes {
+   hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f", enable_inner_tiled = true, iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %cst = arith.constant 0.0 : f32
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %0 = iree_encoding.set_encoding %arg0 encoding_dims{%m, %n, %k} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_lhs_it>
+  %1 = iree_encoding.set_encoding %arg1 encoding_dims{%m, %n, %k} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_rhs_it>
+  %2 = tensor.empty(%d0, %d1) : tensor<?x?xf32, #encoding_result_it>
+  %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?xf32, #encoding_result_it>)
+      -> tensor<?x?xf32, #encoding_result_it>
+  %4 = linalg.matmul ins(%0, %1 : tensor<?x?xf32, #encoding_lhs_it>, tensor<?x?xf32, #encoding_rhs_it>)
+      outs(%3 : tensor<?x?xf32, #encoding_result_it>) -> tensor<?x?xf32, #encoding_result_it>
+  %5 = iree_encoding.unset_encoding %4 encoding_dims{%m, %n, %k} : tensor<?x?xf32, #encoding_result_it> -> tensor<?x?xf32>{%d0, %d1}
+  return %5 : tensor<?x?xf32>
+}
+//   CHECK-DAG: #[[$MAP_N:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>
+// CHECK-LABEL: func @pack_gemm_fill_dynamic_inner_tiled_avx512(
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//   CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
+//   CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
+//   CHECK-DAG:   %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
+//   CHECK-DAG:   %[[D1:.+]] = tensor.dim %[[ARG1]], %[[C1]]
+//   CHECK-DAG:   %[[OUT_D1:.+]] = affine.apply #[[$MAP_N]]()[%[[D1]]]
+//   CHECK-DAG:   %[[PACK_LHS:.+]] = linalg.pack {{.*}}%[[ARG0]]
+//       CHECK:   %[[PACK_RHS:.+]] = linalg.pack
+//  CHECK-SAME:     %[[ARG1]]
+//   CHECK-DAG:   %[[EMPTY:.+]] = tensor.empty(%[[D0]], %[[OUT_D1]]) : tensor<?x?x1x16xf32>
+//       CHECK:   %[[FILL:.+]] = linalg.fill
+//  CHECK-SAME:       outs(%[[EMPTY]] :
+//       CHECK:   %[[INNER:.+]] = iree_codegen.inner_tiled ins(%[[PACK_LHS]], %[[PACK_RHS]]) outs(%[[FILL]])
+//  CHECK-SAME:       kind = #iree_cpu.data_tiled_mma_layout<intrinsic = MMA_X86_AVX512_1x16x1_F32_F32>, semantics = #iree_cpu.mma_semantics<>
+//       CHECK:   %[[UNPACK:.+]] = linalg.unpack %[[INNER]]
+//       CHECK:   return %[[UNPACK]]
+
+// -----
+
+#map_it_se = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map_it_se1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map_it_se2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#encoding_it_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_it_se, #map_it_se1, #map_it_se2], iteration_sizes = [127, 255, ?]>
+#encoding_it_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_it_se, #map_it_se1, #map_it_se2], iteration_sizes = [127, 255, ?]>
+#encoding_it_res = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_it_se, #map_it_se1, #map_it_se2], iteration_sizes = [127, 255, ?]>
+func.func @set_encoding_matmul_LHS_inner_tiled_avx512(%arg0: tensor<127x255xf32>, %k: index) -> tensor<127x255xf32, #encoding_it_lhs> attributes {
+   hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f", enable_inner_tiled = true, iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+  %0 = iree_encoding.set_encoding %arg0 encoding_dims{%k} : tensor<127x255xf32> -> tensor<127x255xf32, #encoding_it_lhs>
+  return %0 : tensor<127x255xf32, #encoding_it_lhs>
+}
+func.func @set_encoding_matmul_RHS_inner_tiled_avx512(%arg0: tensor<127x255xf32>, %k: index) -> tensor<127x255xf32, #encoding_it_rhs> attributes {
+   hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f", enable_inner_tiled = true, iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+  %0 = iree_encoding.set_encoding %arg0 encoding_dims{%k} : tensor<127x255xf32> -> tensor<127x255xf32, #encoding_it_rhs>
+  return %0 : tensor<127x255xf32, #encoding_it_rhs>
+}
+func.func @unset_encoding_matmul_RESULT_inner_tiled_avx512(%arg0: tensor<127x255xf32, #encoding_it_res>, %k: index) -> tensor<127x255xf32> attributes {
+   hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", cpu_features = "+avx512f", enable_inner_tiled = true, iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+  %0 = iree_encoding.unset_encoding %arg0 encoding_dims{%k} : tensor<127x255xf32, #encoding_it_res> -> tensor<127x255xf32>
+  return %0 : tensor<127x255xf32>
+}
+// CHECK-LABEL: func @set_encoding_matmul_LHS_inner_tiled_avx512(
+//  CHECK-SAME:   %[[INPUT:[a-zA-Z0-9]+]]: tensor<127x255xf32>
+//       CHECK:   %[[EMPTY:.+]] = tensor.empty() : tensor<127x255x1x1xf32>
+//       CHECK:   %[[PACK:.+]] = linalg.pack %[[INPUT]] outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [1, 1] into %[[EMPTY]] : tensor<127x255xf32> -> tensor<127x255x1x1xf32>
+//       CHECK:   return %[[PACK]] : tensor<127x255x1x1xf32>
+// CHECK-LABEL: func @set_encoding_matmul_RHS_inner_tiled_avx512(
+//  CHECK-SAME:   %[[INPUT_R:[a-zA-Z0-9]+]]: tensor<127x255xf32>
+//   CHECK-DAG:   %[[CST_R:.+]] = arith.constant 0.0
+//       CHECK:   %[[EMPTY_R:.+]] = tensor.empty() : tensor<16x127x16x1xf32>
+//       CHECK:   %[[PACK_R:.+]] = linalg.pack %[[INPUT_R]] padding_value(%[[CST_R]] : f32) outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 1] into %[[EMPTY_R]] : tensor<127x255xf32> -> tensor<16x127x16x1xf32>
+//       CHECK:   return %[[PACK_R]] : tensor<16x127x16x1xf32>
+// CHECK-LABEL: func @unset_encoding_matmul_RESULT_inner_tiled_avx512(
+//  CHECK-SAME:   %[[PACKED:[a-zA-Z0-9]+]]: tensor<127x16x1x16xf32>
+//       CHECK:   %[[EMPTY_U:.+]] = tensor.empty() : tensor<127x255xf32>
+//       CHECK:   %[[UNPACK:.+]] = linalg.unpack %[[PACKED]] outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [1, 16] into %[[EMPTY_U]] : tensor<127x16x1x16xf32> -> tensor<127x255xf32>
+//       CHECK:   return %[[UNPACK]]
+
+// -----
+
 // It tests with bindings and checks that the reshape ops are folded into bindings.
 
 #executable_target_xyz = #hal.executable.target<"llvm-cpu", "xyz", {target_triple = "x86_64-xyz-xyz", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>

compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp

@@ -332,6 +332,8 @@ getRowMajorTilesMNKShape(MMAIntrinsic intrinsic) {
   switch (intrinsic) {
   case MMAIntrinsic::None:
     return Tuple{0, 0, 0};
+  case MMAIntrinsic::MMA_X86_AVX2_FMA_1x8x1_F32_F32:
+    return Tuple{1, 8, 1};
   case MMAIntrinsic::MMA_X86_AVX512_1x8x1_F64_F64:
     return Tuple{1, 8, 1};
   case MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F32:
@@ -370,7 +372,7 @@ static Codegen::TileSwizzle fixupSwizzle(Codegen::TileSwizzle swizzle) {
 }
 
 Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma,
-                                         int operandIdx) {
+                                         bool transposed, int operandIdx) {
   using TileSwizzle = Codegen::TileSwizzle;
   using Dim = TileSwizzle::Dim;
 
@@ -394,6 +396,13 @@ Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma,
     return TileSwizzle();
   }
   auto [mSize, nSize, kSize] = *maybeMnkTuple;
+  // In the transposed orientation, the intrinsic's hardware (M, N) roles are
+  // logically swapped: what the matmul code sees as M is driven by the
+  // intrinsic's N-dim, and vice versa. Swap the sizes up front so that the
+  // per-operand expansion code below can stay oblivious to orientation.
+  if (transposed) {
+    std::swap(mSize, nSize);
+  }
   TileSwizzle swizzle;
   swizzle.expandShape().resize(2);
   auto expandIfNonUnit = [](TileSwizzle &swizzle, int dim, int size) {
@@ -402,6 +411,13 @@ Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma,
     }
   };
 
+  // For every operand, expandShape[0] is the outer physical dim and
+  // expandShape[1] is the inner physical dim, with identity permutation. For
+  // the ACC in particular, `transposed_intrinsic` flips the logical (M, N) to
+  // physical (N, M), which we encode by swapping which logical dim fills each
+  // expandShape group rather than by a non-identity permutation. That way all
+  // three operand swizzles can be read on equal footing, just like LHS (M, K)
+  // and RHS (N, K).
   if (operandIdx == 0) {
     constexpr int M = 0, K = 1;
     expandIfNonUnit(swizzle, K, kSize);
@@ -411,17 +427,19 @@ Codegen::TileSwizzle getIntrinsicSwizzle(IREE::CPU::MMAIntrinsic mma,
     expandIfNonUnit(swizzle, K, kSize);
     expandIfNonUnit(swizzle, N, nSize);
   } else {
-    constexpr int M = 0, N = 1;
-    expandIfNonUnit(swizzle, N, nSize);
-    expandIfNonUnit(swizzle, M, mSize);
+    int64_t accOuter = transposed ? nSize : mSize;
+    int64_t accInner = transposed ? mSize : nSize;
+    expandIfNonUnit(swizzle, 1, accInner);
+    expandIfNonUnit(swizzle, 0, accOuter);
   }
   return fixupSwizzle(std::move(swizzle));
 }
 
 Codegen::TileSwizzle getSwizzle(IREE::CPU::DataTiledMMAAttr mma,
                                 int operandIdx) {
   using TileSwizzle = Codegen::TileSwizzle;
-  TileSwizzle swizzle = getIntrinsicSwizzle(mma.getIntrinsic(), operandIdx);
+  TileSwizzle swizzle = getIntrinsicSwizzle(
+      mma.getIntrinsic(), mma.getTransposedIntrinsic(), operandIdx);
   TileSwizzle::Dim intrinsicsM =
       TileSwizzle::Dim::crossIntrinsic(mma.getIntrinsicsM());
   TileSwizzle::Dim intrinsicsN =
@@ -430,8 +448,9 @@ Codegen::TileSwizzle getSwizzle(IREE::CPU::DataTiledMMAAttr mma,
       TileSwizzle::Dim::crossIntrinsic(mma.getIntrinsicsK());
   // Each swizzle is built as (outer physical dim, inner physical dim) in
   // expandShape[0], expandShape[1]. LHS is (M, K), RHS is (N, K), ACC is
-  // (M, N). The expansion below injects the intrinsics_* cross-intrinsic
-  // factors into whichever group represents each logical dim.
+  // (M, N) normally and (N, M) when `transposed_intrinsic` is set. The
+  // expansion below injects the intrinsics_* cross-intrinsic factors into
+  // whichever group represents each logical dim.
   if (operandIdx == 0) {
     constexpr int M = 0, K = 1;
     if (intrinsicsK.size() > 1) {
@@ -449,12 +468,14 @@ Codegen::TileSwizzle getSwizzle(IREE::CPU::DataTiledMMAAttr mma,
       Codegen::expand(swizzle, N, intrinsicsN);
     }
   } else {
-    constexpr int M = 0, N = 1;
-    if (intrinsicsN.size() > 1) {
-      Codegen::expand(swizzle, N, intrinsicsN);
+    bool transposed = mma.getTransposedIntrinsic();
+    TileSwizzle::Dim accOuterIntr = transposed ? intrinsicsN : intrinsicsM;
+    TileSwizzle::Dim accInnerIntr = transposed ? intrinsicsM : intrinsicsN;
+    if (accInnerIntr.size() > 1) {
+      Codegen::expand(swizzle, /*srcIdx=*/1, accInnerIntr);
     }
-    if (intrinsicsM.size() > 1) {
-      Codegen::expand(swizzle, M, intrinsicsM);
+    if (accOuterIntr.size() > 1) {
+      Codegen::expand(swizzle, /*srcIdx=*/0, accOuterIntr);
     }
   }
   return swizzle;
@@ -472,6 +493,8 @@ static std::tuple<Type, Type, Type> getABCElementTypes(MLIRContext *context,
   switch (intrinsic) {
   case MMAIntrinsic::None:
     return {Type(), Type(), Type()};
+  case MMAIntrinsic::MMA_X86_AVX2_FMA_1x8x1_F32_F32:
+    return {f32, f32, f32};
   case MMAIntrinsic::MMA_X86_AVX512_1x8x1_F64_F64:
     return {f64, f64, f64};
   case MMAIntrinsic::MMA_X86_AVX512_1x16x1_F32_F32:
@@ -536,7 +559,10 @@ void DataTiledMMAAttr::getUndistributedTileTypes(
   // inner physical dim) in expandShape[0], expandShape[1]. This mirrors GPU's
   // DataTiledMMA, where the tile types encode the layout directly and no
   // separate `permutations` attribute is needed on `inner_tiled`. LHS is
-  // (M, K), RHS is (N, K), ACC is (M, N).
+  // (M, K), RHS is (N, K), and ACC is (M, N) for non-transposed intrinsics
+  // and (N, M) for transposed ones; that transposition is baked into the ACC
+  // swizzle itself (see getIntrinsicSwizzle), so we can query all three
+  // operands uniformly here.
   auto tileType = [&](Codegen::TileSwizzle swizzle, Type elemType) {
     auto [outer, outerScalable] =
         getVectorAxisSizeAndScalability(swizzle.expandShape()[0]);

compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td

@@ -83,6 +83,12 @@ def IREECPU_DataTiledMMAAttr
     thread-distribution, no concept of subgroups and no interleaving of
     intrinsics' layout.
 
+    The `transposed_intrinsic` boolean toggles between the intrinsic's natural
+    orientation and an M↔N-swapped orientation (same hardware instruction with
+    LHS/RHS roles exchanged, accumulator laid out column-major). Orthogonal to
+    the intrinsic itself: any intrinsic may be used transposed, including
+    square ones where the effect is only the accumulator layout change.
+
     Some GPU-specific methods in IREECodegen_InnerTileDescAttrInterface are left
     here but are unused.
   }];
@@ -99,7 +105,12 @@ def IREECPU_DataTiledMMAAttr
               "Intrinsic count along the N dimension.">:$intrinsics_n,
           DefaultValuedParameter<
               "int64_t", "1",
-              "Intrinsic count along the K dimension.">:$intrinsics_k);
+              "Intrinsic count along the K dimension.">:$intrinsics_k,
+          DefaultValuedParameter<
+              "bool", "false",
+              "If true, the intrinsic is used in an M↔N-swapped orientation "
+              "(LHS/RHS roles exchanged, accumulator "
+              "column-major).">:$transposed_intrinsic);
 }
 
 //===----------------------------------------------------------------------===//

compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUEnums.td

@@ -96,6 +96,18 @@ class IREECPU_I32EnumAttr<string name, string summary,
 // * D enumerates intrinsics that share the same 0xABC* bits.
 //
 def IREECPU_MMA_None : I32EnumAttrCase<"None", 0>;
+
+// X86 AVX/AVX2 + FMA3 (256-bit YMM registers). FMA3 is a separate ISA extension
+// from AVX2 itself; these intrinsics lower to `vfmadd...`, so the name includes
+// `AVX2_FMA_` and `getMmaIntrinsicRequiredFeatures` requires both `+avx2` and
+// `+fma`.
+def MMA_X86_AVX2_FMA_1x8x1_F32_F32
+    : I32EnumAttrCase<"MMA_X86_AVX2_FMA_1x8x1_F32_F32", 0x1210>;
+
+// X86 AVX-512 (ZMM registers). Only the "natural" (canonical M<=N) orientation
+// of each intrinsic is listed here; the M/N-swapped ("transposed") variant is
+// expressed by the `transposed_intrinsic` boolean on DataTiledMMAAttr rather
+// than as a separate enum value.
 def MMA_X86_AVX512_1x8x1_F64_F64
     : I32EnumAttrCase<"MMA_X86_AVX512_1x8x1_F64_F64", 0x1300>;
 def MMA_X86_AVX512_1x16x1_F32_F32
@@ -125,6 +137,9 @@ def IREECPU_MMAIntrinsic
           "MMAIntrinsic", "Descriptor for different MMA intrinsics",
           [IREECPU_MMA_None,
 
+           // X86 AVX/AVX2 + FMA
+           MMA_X86_AVX2_FMA_1x8x1_F32_F32,
+
            // X86 AVX-512
            MMA_X86_AVX512_1x8x1_F64_F64, MMA_X86_AVX512_1x16x1_F32_F32,
            MMA_X86_AVX512_1x16x1_F32_F16_CASTF32,

compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h

@@ -64,7 +64,11 @@ SmallVector<int> getTilingLevelsAsInts();
 StringRef getTilingLevelName(TilingLevel level);
 
 // Returns the TileSwizzle for the given intrinsic and operand index.
-Codegen::TileSwizzle getIntrinsicSwizzle(MMAIntrinsic mma, int operandIdx);
+// If `transposed` is true, the intrinsic is used in an M↔N-swapped
+// orientation: the physical tile layouts reflect LHS/RHS roles being
+// exchanged, and the accumulator is laid out column-major.
+Codegen::TileSwizzle getIntrinsicSwizzle(MMAIntrinsic mma, bool transposed,
+                                         int operandIdx);
 
 // Returns the TileSwizzle for the given MMA attr and operand index.
 Codegen::TileSwizzle getSwizzle(DataTiledMMAAttr mma, int operandIdx);