add E2E loweing for migraphx.attention

Author: umangyadav

mlir/include/mlir-c/Dialect/MIGraphX.h

@@ -83,8 +83,8 @@ mlirMIGraphXAddBackendPipeline(MlirPassManager pm,
 /// \p queries, \p keys, \p values are the required Q, K, V operands.
 /// \p preSoftmaxElemWiseInputs is an array of \p numPreSoftmaxInputs additional
 ///    operands for element-wise fusion before softmax (can be NULL if 0).
-/// \p resultType is the tensor type of the attention result (required).
-/// \p lseType is the tensor type of the optional log-sum-exp output; pass
+/// \p resultType is the MIXRShaped type of the attention result (required).
+/// \p lseType is the MIXRShaped type of the optional log-sum-exp output; pass
 ///    a null type (via mlirTypeIsNull) to omit.
 /// \p softmaxType is the optional element type for softmax computation; pass
 ///    a null type to omit.

mlir/include/mlir/Conversion/MIGraphXAttentionToRock/MIGraphXAttentionToRock.h

@@ -0,0 +1,24 @@
+//===-- MIGraphXAttentionToRock.h - Lower migraphx.attention to rock --*- C++
+//-*-===//
+//
+// Part of the rocMLIR Project, under the Apache License v2.0 with LLVM
+// Exceptions. See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// Copyright (c) 2025 Advanced Micro Devices
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_MIGRAPHXATTENTIONTOROCK_H
+#define MLIR_CONVERSION_MIGRAPHXATTENTIONTOROCK_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+#define GEN_PASS_DECL_MIGRAPHXATTENTIONTOROCKPASS
+#include "mlir/Conversion/RocMLIRPasses.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_MIGRAPHXATTENTIONTOROCK_H

mlir/include/mlir/Conversion/RocMLIRPasses.h

@@ -12,6 +12,7 @@
 #include "mlir/Conversion/EmulateFp8ExtTrunc/EmulateFp8ExtTrunc.h"
 #include "mlir/Conversion/FixTosaCastRounding/FixTosaCastRounding.h"
 #include "mlir/Conversion/LinalgToRock/LinalgToRock.h"
+#include "mlir/Conversion/MIGraphXAttentionToRock/MIGraphXAttentionToRock.h"
 #include "mlir/Conversion/MIGraphXToLinalg/MIGraphXToLinalg.h"
 #include "mlir/Conversion/MIGraphXToTosa/MIGraphXToTosa.h"
 #include "mlir/Conversion/Passes.h"

mlir/include/mlir/Conversion/RocMLIRPasses.td

@@ -193,4 +193,25 @@ def LinalgToRockPass : Pass<"linalg-to-rock", "::mlir::func::FuncOp"> {
   let dependentDialects = ["rock::RockDialect",
                            "bufferization::BufferizationDialect"];
 }
+//===----------------------------------------------------------------------===//
+// MIGraphXAttentionToRock
+//===----------------------------------------------------------------------===//
+def MIGraphXAttentionToRockPass
+    : Pass<"migraphx-attention-to-rock"> {
+  let summary = "Lower migraphx.attention to rock.attention";
+  let description = [{
+    Pass that converts migraphx.attention operations directly to
+    rock.attention operations for GPU compilation.
+  }];
+  let dependentDialects = [
+    "arith::ArithDialect",
+    "linalg::LinalgDialect",
+    "rock::RockDialect",
+    "bufferization::BufferizationDialect",
+    "memref::MemRefDialect",
+    "tensor::TensorDialect",
+    "migraphx::MIGraphXDialect",
+  ];
+}
+
 #endif // ROCMLIR_CONVERSION_PASSES

mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td

@@ -598,10 +598,10 @@ def MIGraphX_AttentionOp
       Arguments<(ins MIXRShapedOf<AttentionQKTypes>:$queries,
           MIXRShapedOf<AttentionQKTypes>:$keys,
           MIXRShapedOf<AttentionVTypes>:$values,
-          Variadic<AnyTypeOf<[AnyMIXRShaped, AnyRankedTensor]>>:$preSoftmaxElemWiseInputs,
+          Variadic<AnyMIXRShaped>:$preSoftmaxElemWiseInputs,
           OptionalAttr<TypeAttr>:$softmaxType)>,
-      Results<(outs AnyRankedTensor:$result,
-          Optional<AnyRankedTensor>:$lse)> {
+      Results<(outs AnyMIXRShaped:$result,
+          Optional<AnyMIXRShaped>:$lse)> {
   let summary = "Attention operation of transformer models";
   let description = [{
     Performs the operation result = SOFTMAX(preSoftmaxBody(queries * keys^T, preSoftmaxElemWiseInputs)) * values.

mlir/lib/Conversion/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_subdirectory(EmulateFp8ExtTrunc)
 add_subdirectory(FixTosaCastRounding)
+add_subdirectory(MIGraphXAttentionToRock)
 add_subdirectory(MIGraphXToTosa)
 add_subdirectory(RocmlirCustomTosaDecompose)
 add_subdirectory(RocmlirCustomTosaToLinalg)

mlir/lib/Conversion/MIGraphXAttentionToRock/CMakeLists.txt

@@ -0,0 +1,20 @@
+add_rocmlir_conversion_library(MLIRMIGraphXAttentionToRock
+  MIGraphXAttentionToRock.cpp
+
+  DEPENDS
+  RocMLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRArithDialect
+  MLIRIR
+  MLIRFuncDialect
+  MLIRLinalgDialect
+  MLIRMIGraphXDialect
+  MLIRMemRefDialect
+  MLIRRockOps
+  MLIRPass
+  MLIRTensorDialect
+  MLIRBufferizationDialect
+  MLIRTransformUtils
+  MLIRSupport
+  )