Attention LSE migraphx integration (#1887)

* Attention LSE (log-sum-exp): migraphx integration

This PR introduces migraphx integration for the optional LSE output (attention kernels).

Author: dhernandez0

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

@@ -561,6 +561,19 @@ def MIGraphX_ReduceSumOp :
   }];
 }
 
+def MIGraphX_ReduceMaxOp
+    : MIGraphX_Op<"reduce_max">,
+      Arguments<(ins AnyMIXRShaped:$input, I64ArrayAttr:$axes)>,
+      Results<(outs AnyMIXRShaped:$output)> {
+  let summary = "Get the max of the values in given axis";
+  let description = [{
+    The `migraphx.reduce_max` op.
+  }];
+  let assemblyFormat = [{
+    $input attr-dict `:` type($input) `->` type($output)
+  }];
+}
+
 //--------- Execution layer Ops
 def MIGraphX_CodeObjOp :
     MIGraphX_Op<"code_object">,

mlir/include/mlir/Dialect/Rock/IR/RockOps.td

@@ -209,7 +209,8 @@ def Rock_AttentionOp
     : Rock_Op<
           "attention", [DeclareOpInterfaceMethods<RockGemmGemmWrapperInterface>,
                         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-                        RockFusionRoot, AttrSizedOperandSegments]>,
+                        RockFusionRoot, AttrSizedOperandSegments,
+                        AttrSizedResultSegments]>,
       Arguments<(ins TensorOrMemRefOf<[F32, F16, BF16, I8]>:$queries,
           TensorOrMemRefOf<[F32, F16, BF16, I8]>:$keys,
           TensorOrMemRefOf<[F32, F16, BF16]>:$values,
@@ -223,7 +224,8 @@ def Rock_AttentionOp
           OptionalAttr<RockTuningParamAttrInterface>:$params0,
           OptionalAttr<RockTuningParamAttrInterface>:$params1,
           I32Attr:$firstGemmIdx)>,
-      Results<(outs Optional<TensorOf<[F32, F16, BF16]>>:$result)> {
+      Results<(outs Optional<TensorOf<[F32, F16, BF16]>>:$result,
+          Optional<TensorOf<[F32, F16, BF16]>>:$lseOut)> {
   let summary = "Attention operation of transformer models";
   let description = [{
     Performs the operation out = SOFTMAX(preSoftmaxBody(queries * keys, preSoftmaxElemWiseInputs)) * values.
@@ -261,7 +263,12 @@ def Rock_AttentionOp
         (`lse` `=` $lse^ `:` type($lse) `\n`)?
         (`qk` `=` `elementwise` (`otherIns` `(` $preSoftmaxElemWiseInputs^ `:` type($preSoftmaxElemWiseInputs) `)`)? $preSoftmaxBody^ `\n`)?
         (`tr` $oTransposed^)? $out `=` `softmax` `(` `qk` `)` `*` (`tr` $vTransposed^)? $values `:` type($values) `->` type($out) `\n`
-    `}` attr-dict (`->` type($result)^)?
+    `}` attr-dict (`->` type($result)^)? (`,` type($lseOut)^)?
+  }];
+
+  // Return operand of LSE if it exists.
+  let extraClassDeclaration = [{
+    ::mlir::OpOperand* getOutLseArgument() { return getLse() ? &(*this)->getOpOperand(getNumOperands() - 1) : nullptr; }
   }];
 }
 

mlir/lib/Conversion/MIGraphXToTosa/MIGraphXToTosa.cpp

@@ -738,16 +738,20 @@ struct ReduceMeanConverter final
                   ConversionPatternRewriter &rewriter) const final;
 };
 
-struct ReduceSumConverter final
-    : public OpConversionPattern<migraphx::ReduceSumOp> {
-  using OpConversionPattern<migraphx::ReduceSumOp>::OpConversionPattern;
+namespace {
+template <typename MIGraphXOp, typename TosaOp>
+struct ReduceConverter final : public OpConversionPattern<MIGraphXOp> {
+  using OpConversionPattern<MIGraphXOp>::OpConversionPattern;
+  using OpAdaptor = typename OpConversionPattern<MIGraphXOp>::OpAdaptor;
 
   LogicalResult
-  matchAndRewrite(migraphx::ReduceSumOp op, OpAdaptor adaptor,
+  matchAndRewrite(MIGraphXOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final;
 };
 } // namespace
 
+} // namespace
+
 tosa::ConstOp ReduceMeanConverter::createNumElementsTosaConst(
     Location loc, TypedValue<RankedTensorType> inputTensor,
     IntegerAttr axisAttr, ConversionPatternRewriter &rewriter) const {
@@ -805,9 +809,10 @@ LogicalResult ReduceMeanConverter::matchAndRewrite(
   return success();
 }
 
-LogicalResult
-ReduceSumConverter::matchAndRewrite(migraphx::ReduceSumOp op, OpAdaptor adaptor,
-                                    ConversionPatternRewriter &rewriter) const {
+template <typename MIGraphXOp, typename TosaOp>
+LogicalResult ReduceConverter<MIGraphXOp, TosaOp>::matchAndRewrite(
+    MIGraphXOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
   Location loc = op.getLoc();
   ArrayRef<Attribute> axes = op.getAxes().getValue();
   if (axes.size() != 1) {
@@ -817,9 +822,9 @@ ReduceSumConverter::matchAndRewrite(migraphx::ReduceSumOp op, OpAdaptor adaptor,
       rewriter.getI32IntegerAttr(cast<IntegerAttr>(axes[0]).getInt());
   auto input = cast<TypedValue<RankedTensorType>>(adaptor.getInput());
   Type elementType = input.getType().getElementType();
-  auto tosaReduceSum = createOpAndInfer<tosa::ReduceSumOp>(
-      rewriter, loc, elementType, input, axis);
-  rewriter.replaceOp(op, tosaReduceSum);
+  auto tosaReduce =
+      createOpAndInfer<TosaOp>(rewriter, loc, elementType, input, axis);
+  rewriter.replaceOp(op, tosaReduce);
   return success();
 }
 
@@ -1493,7 +1498,9 @@ void migraphx::populateMIGraphXToTosaConversionPatterns(
                DotConverter<DotOp>, DotConverter<QuantDotOp>,
                BroadcastConverter, MultiBroadcastConverter, TransposeConverter,
                ReshapeConverter, SliceConverter, ReduceMeanConverter,
-               ReduceSumConverter, TrivialConverter<AddOp, tosa::AddOp>,
+               ReduceConverter<ReduceSumOp, tosa::ReduceSumOp>,
+               ReduceConverter<ReduceMaxOp, tosa::ReduceMaxOp>,
+               TrivialConverter<AddOp, tosa::AddOp>,
                TrivialConverter<SubOp, tosa::SubOp>,
                TrivialConverter<PowOp, tosa::PowOp>, DivConverter, MulConverter,
                TrivialConverter<AbsOp, tosa::AbsOp>,