[rocjitsu] Pair mixed-format WMMA operands by K coordinate

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/mma_exec.h

@@ -367,8 +367,11 @@ inline InputLoc wmma_f8f6f4_mixed_subbyte_input_loc(uint32_t dim, uint32_t K, ui
                                                     uint32_t k, uint32_t data_bits) {
   if (dim == 16 && K == 128 && data_bits < 8) {
     const uint32_t lane = i + 16u * ((k >> 5) & 1u);
-    const uint32_t slot = 32u * ((k >> 6) & 1u) + 16u * ((k >> 2) & 1u) + 8u * ((k >> 4) & 1u) +
-                          4u * ((k >> 3) & 1u) + 2u * ((k >> 1) & 1u) + (k & 1u);
+    // Use the same logical K order as the 8-bit operand in wmma_input_loc:
+    // contiguous groups of 16 there pair with groups of 32 here. The old
+    // sub-byte permutation assumed an 8-bit layout split every four elements,
+    // so it paired different K values when only one operand was sub-byte.
+    const uint32_t slot = 32u * (k >> 6) + (k & 31u);
     return wmma_packed_input_loc(lane, slot, data_bits);
   }
   return wmma_f8f6f4_ab_input_loc(dim, K, i, k, data_bits);
@@ -399,7 +402,7 @@ inline InputLoc wmma_b_input_loc(uint32_t N, uint32_t K, uint32_t col, uint32_t
   return wmma_f8f6f4_input_loc(N, K, col, k, b_bits, b_bits < 8 && a_bits == 8);
 }
 
-/// Compute the CDNA5 block-scaled WMMA input layout. Unlike the ordinary
+/// Compute the CDNA5 block-scaled WMMA input layout. Unlike the same-format ordinary
 /// F8F6F4 WMMA layout, each operand format stores contiguous K ranges in each
 /// lane group: 16 values for 8-bit inputs and 32 values for 4/6-bit inputs.
 inline InputLoc wmma_block_scaled_input_loc(uint32_t dim, uint32_t K, uint32_t index, uint32_t k,

emulation/rocjitsu/tests/cdna5_architecture_test.cpp

@@ -1161,6 +1161,65 @@ TEST(Gfx1250ExecutionTest, WmmaScaleExecutesAllMatrixFormatPairs) {
   }
 }
 
+TEST(Gfx1250ExecutionTest, WmmaMixedFormatsPairTheSameKElements) {
+  Gfx1250Sim sim;
+  auto *cu = sim.cu();
+  auto *wf = sim.dispatch_scratch_wf();
+  ASSERT_NE(wf, nullptr);
+  wf->set_exec(0xffffffffu);
+  const uint32_t base = wf->vgpr_alloc().base;
+  auto decoder = Decoder::create(ROCJITSU_CODE_ARCH_CDNA5);
+  ASSERT_NE(decoder, nullptr);
+  ForceScalarGuard scalar_guard;
+
+  for (bool force_scalar : {true, false}) {
+    util::set_force_scalar_for_testing(force_scalar);
+    for (uint32_t afmt = 0; afmt <= 4; ++afmt) {
+      for (uint32_t bfmt = 0; bfmt <= 4; ++bfmt) {
+        if ((afmt < 2) == (bfmt < 2))
+          continue;
+        SCOPED_TRACE(::testing::Message()
+                     << "scalar=" << force_scalar << " A=" << afmt << " B=" << bfmt);
+        for (uint32_t reg = 0; reg < 72; ++reg)
+          for (uint32_t lane = 0; lane < 32; ++lane)
+            cu->write_vgpr(base + reg, lane, 0);
+        // A single nonzero product in each row/column isolates K pairing.
+        // These independent physical locations also cover 6-bit word crossings.
+        for (uint32_t row = 0; row < 16; ++row) {
+          const uint32_t k = 5u + 7u * row;
+          write_wmma_packed(*cu, base,
+                            manual_block_scaled_wmma_loc(16, row, k, wmma_format_bits(afmt)),
+                            encode_wmma_one(afmt));
+          write_wmma_packed(*cu, base + 32,
+                            manual_block_scaled_wmma_loc(16, row, k, wmma_format_bits(bfmt)),
+                            encode_wmma_one(bfmt));
+        }
+        for (unsigned variant : {0u, 1u, 2u}) {
+          SCOPED_TRACE(::testing::Message() << "scale variant=" << variant);
+          for (uint32_t reg = 64; reg < 72; ++reg)
+            for (uint32_t lane = 0; lane < 32; ++lane)
+              cu->write_vgpr(base + reg, lane, 0);
+          const auto words =
+              build_scaled_wmma_execution_words(16, variant == 2, afmt, bfmt, 128, 128, 0, 0);
+          // Exercise ordinary, scale32, and scale16 with neutral inline scales.
+          std::unique_ptr<Instruction> inst(
+              decode_valid(*decoder, words.data() + (variant ? 0 : 2)));
+          ASSERT_NE(inst, nullptr);
+          ASSERT_TRUE(cu->execute_instruction(inst.get(), *wf).succeeded());
+          for (uint32_t row = 0; row < 16; ++row)
+            for (uint32_t col = 0; col < 16; ++col) {
+              const uint32_t lane = col + 16u * (row / 8u);
+              const float expected = row == col ? 1.0f : 0.0f;
+              EXPECT_EQ(cu->read_vgpr(base + 64 + row % 8u, lane),
+                        std::bit_cast<uint32_t>(expected))
+                  << "row=" << row << " col=" << col;
+            }
+        }
+      }
+    }
+  }
+}
+
 TEST(Gfx1250ExecutionTest, WmmaRegularScaleInlineZeroMatchesNeutralScalarSources) {
   Gfx1250Sim sim;
   auto *cu = sim.cu();

emulation/rocjitsu/tests/shared_infra_test.cpp

@@ -950,8 +950,8 @@ TEST(MfmaExecTest, WmmaF8f6f4K128InputLocUsesPairAwareSubbyteLayouts) {
   auto mixed6 = amdgpu::wmma_f8f6f4_input_loc(16, 128, /*i=*/3, /*k=*/4, 6,
                                               /*mixed_subbyte=*/true);
   EXPECT_EQ(mixed6.lane, 3u);
-  EXPECT_EQ(mixed6.vgpr_offset, 3u);
-  EXPECT_EQ(mixed6.bit_offset, 0u);
+  EXPECT_EQ(mixed6.vgpr_offset, 0u);
+  EXPECT_EQ(mixed6.bit_offset, 24u);
 
   auto mixed4 = amdgpu::wmma_f8f6f4_input_loc(16, 128, /*i=*/3, /*k=*/31, 4,
                                               /*mixed_subbyte=*/true);
@@ -982,7 +982,7 @@ TEST(MfmaExecTest, WmmaF8f6f4K128InputLocUsesPairAwareSubbyteLayouts) {
             3u);
 }
 
-TEST(MfmaExecTest, WmmaF8f6f4K128InputLocMatchesManualLayoutsExhaustively) {
+TEST(MfmaExecTest, WmmaF8f6f4K128InputLocMatchesOperandLayoutsExhaustively) {
   for (uint32_t data_bits : {4u, 6u, 8u}) {
     for (uint32_t index = 0; index < 16; ++index) {
       for (uint32_t k = 0; k < 128; ++k) {
@@ -1017,9 +1017,9 @@ TEST(MfmaExecTest, WmmaF8f6f4K128InputLocMatchesManualLayoutsExhaustively) {
         SCOPED_TRACE(::testing::Message()
                      << "mixed data_bits=" << data_bits << " index=" << index << " k=" << k);
         const uint32_t expected_lane = index + 16u * ((k >> 5u) & 1u);
-        const uint32_t expected_slot = 32u * ((k >> 6u) & 1u) + 16u * ((k >> 2u) & 1u) +
-                                       8u * ((k >> 4u) & 1u) + 4u * ((k >> 3u) & 1u) +
-                                       2u * ((k >> 1u) & 1u) + (k & 1u);
+        // Mixed 8-bit/sub-byte operands pair contiguous groups of 16/32 K
+        // values. The decoded one-hot execution test checks the pairing.
+        const uint32_t expected_slot = 32u * (k / 64u) + (k % 32u);
         const uint32_t expected_bit = expected_slot * data_bits;
         const auto actual =
             amdgpu::wmma_f8f6f4_input_loc(16, 128, index, k, data_bits, /*mixed_subbyte=*/true);