[Codegen] PV and QK matmul's must have same acc layout (iree-org#21729)

Fixes issue iree-org#21602 where vector distribute failed due to an unresolvable
layout change in attention.

Check that the 2 matmuls have the same accumulator layout. 

With this change, the reproducer in iree-org#21602 compiles down to a .vmfb. I
have not checked numerics or looked at any performance benchmarks.

---------

Signed-off-by: James Newling <[email protected]>

Author: newling

compiler/src/iree/compiler/Codegen/Common/GPU/GPUHeuristics.cpp

@@ -9,6 +9,7 @@
 
 #include <cstdint>
 
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/Support/DebugLog.h"
@@ -666,6 +667,13 @@ struct ChainedMMAIntrinsics {
   bool canReuseAOutputForB;
 };
 
+static bool matchLayout(IREE::GPU::MMASingleSubgroupLayout layoutA,
+                        IREE::GPU::MMASingleSubgroupLayout layoutB) {
+  return (layoutA.element == layoutB.element) &&
+         (layoutA.thread == layoutB.thread) &&
+         (layoutA.tstrides == layoutB.tstrides);
+};
+
 FailureOr<std::pair<GPUMMASchedule, GPUMMASchedule>> deduceAttentionSchedule(
     const GPUMatmulShapeType &qkMatmul, const GPUMatmulShapeType &pvMatmul,
     ArrayRef<GPUIntrinsicType> intrinsics,
@@ -677,28 +685,33 @@ FailureOr<std::pair<GPUMMASchedule, GPUMMASchedule>> deduceAttentionSchedule(
          qkMatmul.nSizes.size() == 1 && qkMatmul.kSizes.size() == 1 &&
          "unimplemented: multi M/N/K attention schedule");
 
-  std::vector<ChainedMMAIntrinsics> intrinsicPairs;
+  SmallVector<uint64_t> qkViableIntrinsicIndices;
+  SmallVector<uint64_t> pvViableIntrinsicIndices;
+  for (const auto &[index, intrinsic] : llvm::enumerate(intrinsics)) {
+    if (!failed(canTargetIntrinsic(qkMatmul, intrinsic, subgroupSize,
+                                   canUpcastAcc, mustBeAligned))) {
+      qkViableIntrinsicIndices.push_back(index);
+    }
+    if (!failed(canTargetIntrinsic(pvMatmul, intrinsic, subgroupSize,
+                                   canUpcastAcc, mustBeAligned))) {
+      pvViableIntrinsicIndices.push_back(index);
+    }
+  }
 
-  for (const GPUIntrinsicType &intrinsicA : intrinsics) {
-    for (const GPUIntrinsicType &intrinsicB : intrinsics) {
-      if (failed(canTargetIntrinsic(qkMatmul, intrinsicA, subgroupSize,
-                                    canUpcastAcc, mustBeAligned))) {
+  std::vector<ChainedMMAIntrinsics> intrinsicPairs;
+  for (unsigned qkIndex : qkViableIntrinsicIndices) {
+    for (unsigned pvIndex : pvViableIntrinsicIndices) {
+      const GPUIntrinsicType &intrinsicA = intrinsics[qkIndex];
+      const GPUIntrinsicType &intrinsicB = intrinsics[pvIndex];
+      if (!matchLayout(getSingleSubgroupLayout(intrinsicA.mmaKind,
+                                               IREE::GPU::MMAFragment::Acc),
+                       getSingleSubgroupLayout(intrinsicB.mmaKind,
+                                               IREE::GPU::MMAFragment::Acc))) {
         continue;
       }
 
-      if (failed(canTargetIntrinsic(pvMatmul, intrinsicB, subgroupSize,
-                                    canUpcastAcc, mustBeAligned))) {
-        continue;
-      }
       // Check if we can reuse the output of intrinsicA for lhs/rhs of
       // intrinsicB.
-      auto matchLayout =
-          [](IREE::GPU::MMASingleSubgroupLayout layoutA,
-             IREE::GPU::MMASingleSubgroupLayout layoutB) -> bool {
-        return (layoutA.element == layoutB.element) &&
-               (layoutA.thread == layoutB.thread) &&
-               (layoutA.tstrides == layoutB.tstrides);
-      };
       bool canReuseAOutForBLhs =
           matchLayout(getSingleSubgroupLayout(intrinsicA.mmaKind,
                                               IREE::GPU::MMAFragment::Acc),

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h

@@ -17,8 +17,6 @@
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinTypes.h"
 
 namespace mlir::iree_compiler::IREE::GPU {
 

compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_gfx950.mlir

@@ -282,3 +282,33 @@ func.func @attention_large_head_dim_shared_mem() {
 // CHECK-SAME: subgroup_n_count = 1
 // CHECK-SAME: reduction =  [0, 0, 64, 0]
 // CHECK-SAME: workgroup =  [64, 0, 0, 64]
+
+// -----
+
+// The fix introduced for bug https://github.com/iree-org/iree/issues/21602 was
+// to constrain the MMA layout to be the same for the 2 matmuls inside
+// attention. Before this fix, the PV matmul used MFMA_F32_16x16x128_F8E4M3FN
+// and the QK matmul used MFMA_F32_32x32x64_F8E4M3FN. Vector distribution failed
+// to distribute these layouts to threads.
+
+//       CHECK: #iree_codegen.translation_info<pipeline = LLVMGPUVectorDistribute workgroup_size = [256, 1, 1] subgroup_size = 64, {}>
+// CHECK-LABEL: func.func @attention_check_mma_accs_compatable
+
+#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3)>
+#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d2)>
+#map3 = affine_map<(d0, d1, d2, d3, d4) -> ()>
+#map4 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
+func.func @attention_check_mma_accs_compatable(%arg0: f32, %arg1: tensor<960x4096x64xf8E4M3FN>, %arg2: tensor<960x4096x64xf8E4M3FN>, %arg3: tensor<960x4096x64xf8E4M3FN>, %arg4: tensor<960x4096x64xf32>, %arg5: !iree_tensor_ext.dispatch.tensor<writeonly:tensor<960x4096x64xf32>>) {
+  %0 = iree_linalg_ext.attention {indexing_maps = [#map, #map1, #map2, #map3, #map4]} ins(%arg1, %arg2, %arg3, %arg0 : tensor<960x4096x64xf8E4M3FN>, tensor<960x4096x64xf8E4M3FN>, tensor<960x4096x64xf8E4M3FN>, f32) outs(%arg4 : tensor<960x4096x64xf32>) {
+  ^bb0(%arg6: f32):
+    iree_linalg_ext.yield %arg6 : f32
+  } -> tensor<960x4096x64xf32>
+  iree_tensor_ext.dispatch.tensor.store %0, %arg5, offsets = [0, 0, 0], sizes = [960, 4096, 64], strides = [1, 1, 1] : tensor<960x4096x64xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<960x4096x64xf32>>
+  return
+}
+//      CHECK: decomposition_config =
+// CHECK-SAME: attention_pv_matmul
+// CHECK-SAME:   #iree_gpu.mma_layout<MFMA_F32_32x32x64_F8E4M3FN>
+// CHECK-SAME: attention_qk_matmul
+// CHECK-SAME:   #iree_gpu.mma_layout<MFMA_F32_32x32x64_F8E4M3FN>