Workaround issue 1802 (#1800)

Workaround for issue 1802

Author: dhernandez0

mlir/lib/Dialect/Rock/Transforms/GridwiseGemmToBlockwise.cpp

@@ -1760,6 +1760,13 @@ struct GridwiseAttentionAccelRewritePattern
     if (doBypassLDSForQ) {
       ldsLayoutCfgNG0.doSwapThreadIterSubDims = false;
     }
+#ifndef ROCK_DEBUG_ATTENTION_REMOVE_SOFTMAX
+    // TODO: Workaround for issue
+    // https://github.com/ROCm/rocMLIR-internal/issues/1802 If sumRowBuffer and
+    // expMaxDiffRowBuffer are filled with doSwapThreadIterSubDims=true, it does
+    // not match with the second GEMM N dimension. Find a good solution to this.
+    ldsLayoutCfgNG0.doSwapThreadIterSubDims = false;
+#endif
     FailureOr<VectorDimInfo> maybeVectorDimInfoK =
         getVectorDim(rewriter, loc, inK, elemTypeK, blockSize, gemm0KPerBlock,
                      gemm0MPerBlock, gemm0kpack);
@@ -1834,7 +1841,7 @@ struct GridwiseAttentionAccelRewritePattern
     Value accRegBufferGemm0 =
         createBufferForAccelGemmOut(loc, accelParamsGemm0, rewriter);
     // Currently, there is a working assumption that this kernel is meant
-    // support fp32/fp16 This should be guranteed by op verifiers.
+    // support fp32/fp16/bf16. This should be guranteed by op verifiers.
     Type gemmOutElemType = elemTypeQxK;
     Type softmaxInElemType = elemTypeQxK;
     if (elemTypeQ == rewriter.getI8Type()) {
@@ -1985,12 +1992,13 @@ struct GridwiseAttentionAccelRewritePattern
         if (failed(statusLoadQ)) {
           return failure();
         }
+        rewriter.create<LDSBarrierOp>(loc);
 
         TypedValue<MemRefType> ldsTileBufferQ = viewBufferAs(
             rewriter, ldsByteBufferQ, vectorTypeOrSelf(elemTypeQ, gemm0kpack));
         loadGemmOperandsFromLDSToRegs(
             rewriter, loc, ldsTileBufferQ, preAccelRegBuffersQ, "n", blockSize,
-            gemm0InMPerThread, *accelEmitterPtrGemm0.get(),
+            gemm0InNPerThread, *accelEmitterPtrGemm0.get(),
             ldsLayoutCfgNG0.doRotateWithK);
         rewriter.create<GpuDeallocOp>(loc, ldsByteBufferQ);
       }

mlir/test/e2e/PrAttentionBF16.toml

@@ -3,15 +3,25 @@ prefix = "rocmlir-gen"
 suffix = "--operation attention -t bf16 --arch %arch -pv %random_data %rocmlir_gen_flags -relDiff_threshold 0.3 -absDiff_threshold 0.3 -RMS_threshold 0.15 | rocmlir-driver -c | mlir-runner -O2 --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext,%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_float16_utils%shlibext --entry-point-result=void | FileCheck %s --check-prefix="
 
 [[axis]]
-name = "operation"
-values = ["attention"]
-prefix = "--operation "
+name = "transQ"
+values = ["true", "false"]
+prefix = "--transQ="
 
 [[axis]]
 name = "transK"
 values = ["true", "false"]
 prefix = "--transK="
 
+[[axis]]
+name = "transV"
+values = ["true", "false"]
+prefix = "--transV="
+
+[[axis]]
+name = "transO"
+values = ["true", "false"]
+prefix = "--transO="
+
 ## attention variant
 [[suite]]
 name = "pr_attention_bf16"

mlir/test/e2e/PrAttentionF16.toml

@@ -3,15 +3,25 @@ prefix = "rocmlir-gen"
 suffix = "--operation attention -t f16 --arch %arch -pv %random_data %rocmlir_gen_flags -relDiff_threshold 0.02 -absDiff_threshold 0.1 -RMS_threshold 0.015 | rocmlir-driver -c | mlir-runner -O2 --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext,%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_float16_utils%shlibext --entry-point-result=void | FileCheck %s --check-prefix="
 
 [[axis]]
-name = "operation"
-values = ["attention"]
-prefix = "--operation "
+name = "transQ"
+values = ["true", "false"]
+prefix = "--transQ="
 
 [[axis]]
 name = "transK"
 values = ["true", "false"]
 prefix = "--transK="
 
+[[axis]]
+name = "transV"
+values = ["true", "false"]
+prefix = "--transV="
+
+[[axis]]
+name = "transO"
+values = ["true", "false"]
+prefix = "--transO="
+
 ## attention variant
 [[suite]]
 name = "pr_attention_f16"

mlir/test/e2e/PrAttentionF32.toml

@@ -3,15 +3,25 @@ prefix = "rocmlir-gen"
 suffix = "--operation attention -t f32 --arch %arch -pv %random_data %rocmlir_gen_flags -relDiff_threshold 0.00005 | rocmlir-driver -c | mlir-runner -O2 --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext,%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_float16_utils%shlibext --entry-point-result=void | FileCheck %s --check-prefix="
 
 [[axis]]
-name = "operation"
-values = ["attention"]
-prefix = "--operation "
+name = "transQ"
+values = ["true", "false"]
+prefix = "--transQ="
 
 [[axis]]
 name = "transK"
 values = ["true", "false"]
 prefix = "--transK="
 
+[[axis]]
+name = "transV"
+values = ["true", "false"]
+prefix = "--transV="
+
+[[axis]]
+name = "transO"
+values = ["true", "false"]
+prefix = "--transO="
+
 ## attention variant
 [[suite]]
 name = "pr_attention_f32"

mlir/test/e2e/PrAttentionI8.toml

@@ -2,11 +2,6 @@ directory = "PrAttentionI8"
 prefix = "rocmlir-gen"
 suffix = "--operation attention -t i8 --arch %arch -pv %random_data %rocmlir_gen_flags -RMS_threshold 0.01 | rocmlir-driver -c | mlir-runner -O2 --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext,%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_float16_utils%shlibext --entry-point-result=void | FileCheck %s --check-prefix="
 
-[[axis]]
-name = "operation"
-values = ["attention"]
-prefix = "--operation "
-
 [[axis]]
 name = "transK"
 values = ["true", "false"]