[mlir][Linalg] Add a softmax op

[qcolombet]

Summary

This patch adds a softmax op.
For now, nothing interesting happens, we can only do a round trip.
Later patches will add the tiling interface and the lowering of this op to a sequence of simpler ops.

This is graduating the linag_ext.softmax op from iree to LLVM.

Original implementation from @harsh.
@nicolasvasilache co-authored this patch.

[rengolin]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

100

Would it make sense to allow more than one dimension here?

If your tensor is MB x Head x TileX x TileY you want the max over TileX x TileY, no?

[chelini]

Thanks! Some minor comments.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

125

Any specific reason for having this wrapper? Can we use the method generated by tablegen directly getInput? Same for getOutput.

144

I cannot fully understand the sentence, maybe something like: Implement functions necessary for DestinationStyleOpInterface.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

2148

I don't think getting the pointer is needed here. Below we can simply write:

if (failed(...)
  emitOpError("...")

2166

Can you please comment on why we need this folder?

mlir/test/Dialect/Linalg/invalid.mlir

743

nit optional: I would align ins and outs.

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

100

let's extend on a per-need basis if you don't mind, proabably it makes sense yes

125

yes let's drop this

144

Method to implement DestinationStyleOpInterface

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

2166

this is a general folder that rewrites cast(dynamic) -> static

[qcolombet]

• Remove redundant wrapper input()/output()
• Update stall comment
• Get rid of temporary variable op, emitOpError can be called directly

[qcolombet]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

125

Ah thanks, forgot to clean that out.
It made the porting from IREE easier.

144

Good catch

[qcolombet]

• Aligned ins and outs in test
• Rebased

[rengolin]

Should be fine for now. @chelini are you happy with this?

[chelini]

In D153422#4459072, @rengolin wrote:

Should be fine for now. @chelini are you happy with this?

yes.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LINALG_OPS
 #define LINALG_OPS
 
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/DestinationStyleOpInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 // Base class for Linalg dialect ops that do not correspond to library calls.
 class Linalg_Op<string mnemonic, list<Trait> traits = []> :
     Op<Linalg_Dialect, mnemonic, traits>;
let description = [{
     }
     ```
   }];
 
   let assemblyFormat = [{ $dim attr-dict `:` type($result) }];
   let hasVerifier = 1;
 }
 
+def Linalg_SoftmaxOp : Linalg_Op<"softmax",
+    [DestinationStyleOpInterface,
+     PredOpTrait<"input and output have same element type", TCopVTEtIsSameAs<0, 1>>,
+     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
+  let summary = "Softmax operator";
+  let description = [{
+    linalg.softmax computes a numerically stable version of softmax.
+
+    For a given input tensor and a specified dimension `d`, compute:
+      1. the max `m` along that dimension `d`
+      2. f(x) = exp(x - m)
+      3. sum f(x) along dimension d to get l(x).
+      4. compute the final result f(x) / l(x).
+
+    This is an aggregate linalg operation that further reduces to a small DAG of
+    structured operations.
+  }];
+
+  let arguments = (ins AnyShaped:$input,
+                       AnyShaped:$output,
+                       I64Attr:$dimension
+  );
+
+  let results = (outs Variadic<AnyRankedTensor>:$result);
+  let hasFolder = 1;
+  let assemblyFormat = [{
+    attr-dict
+    `dimension` `(` $dimension `)`
+    `ins` `(` $input `:` type($input) `)`
+    `outs` `(` $output `:` type($output) `)`
+    (`->` type($result)^)?
+  }];
+
+  let extraClassDeclaration = [{
+    ShapedType getInputOperandType() {
+      return getInput().getType().cast<ShapedType>();
+    }
+    ShapedType getOutputOperandType() {
+      return getOutput().getType().cast<ShapedType>();
+    }
+    int64_t getInputOperandRank() {
+      return getInputOperandType().getRank();
+    }
+    int64_t getOutputOperandRank() {
+      return getOutputOperandType().getRank();
+    }
+    // Method to implement DestinationStyleOpInterface.
+    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
+      std::pair<unsigned, unsigned> outputsIndexAndLength =
+        getODSOperandIndexAndLength(1);
+      return std::make_pair<int64_t, int64_t>(
+          outputsIndexAndLength.first,
+          outputsIndexAndLength.first + outputsIndexAndLength.second);
+    }
+  }];
+  let hasVerifier = 1;
+}
+
 #endif // LINALG_OPS

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 };
 
 } // namespace
 
 // All named ops canonicalizers and folders are auto-generated in the
 // .cpp.inc.
 
 //===----------------------------------------------------------------------===//
+// SoftmaxOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult SoftmaxOp::verify() {
+  ShapedType inputType = getInputOperandType();
+  ShapedType outputType = getOutputOperandType();
+
+  ArrayRef<int64_t> inputShape = inputType.getShape();
+  ArrayRef<int64_t> outputShape = outputType.getShape();
+  if (failed(verifyCompatibleShape(inputShape, outputShape)))
+    return emitOpError("incompatible output shape");
+
+  int64_t inputRank = getInputOperandRank();
+  int64_t dimension = getDimension();
+  if ((dimension < 0) || (dimension >= inputRank))
+    return emitOpError("incorrect dimension specified");
+
+  return success();
+}
+
+// cast(dynamic) -> static.
+LogicalResult SoftmaxOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) {
+  return memref::foldMemRefCast(*this);
+}
+
+LogicalResult
+SoftmaxOp::reifyResultShapes(OpBuilder &b,
+                             ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
+  return cast<LinalgOp>(getOperation())
+      .reifyResultShapes(b, reifiedReturnShapes);
+}
+
+//===----------------------------------------------------------------------===//
 // LinalgDialect
 //===----------------------------------------------------------------------===//
 
 void LinalgDialect::getCanonicalizationPatterns(
     RewritePatternSet &results) const {
   results.add<EraseDeadLinalgOp, FoldTensorCastConsumerOp,
               InferStaticShapeOfOperands>(getContext());
 }
 
 Operation *LinalgDialect::materializeConstant(OpBuilder &builder,
                                               Attribute value, Type type,
                                               Location loc) {
   return arith::ConstantOp::materialize(builder, value, type, loc);
 }

mlir/test/Dialect/Linalg/invalid.mlir

 
 // -----
 
 func.func @missing_iterator_types() {
   // expected-error @below {{expected "iterator_types" array attribute}}
   linalg.generic {} ins() outs()
   return
 }
+
+// -----
+
+func.func @illegal_softmax_output_shape(%arg0: tensor<2x16x32xf32>) -> tensor<2x16xf32> {
+  %0 = tensor.empty() : tensor<2x16xf32>
+  // expected-error @+1 {{incompatible output shape}}
+  %1 = linalg.softmax dimension(2) ins(%arg0 : tensor<2x16x32xf32>)
+                                   outs(%0: tensor<2x16xf32>)
+    -> tensor<2x16xf32>
+  return %1 : tensor<2x16xf32>
+}

mlir/test/Dialect/Linalg/roundtrip.mlir

 }
 // CHECK-LABEL: func @reduce_arith_with_attr
 // CHECK-NEXT:    %[[REDUCED:.*]] = linalg.reduce
 // CHECK-SAME:    { arith.addf {fastmath = #arith.fastmath<fast>} }
 // CHECK-SAME:    ins
 // CHECK-SAME:    outs
 // CHECK-SAME:    dimensions = [1]
 // CHECK-NEXT:    return %[[REDUCED]] : tensor<16x64xf32>
+
+// -----
+
+func.func @softmax(%arg0: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
+  %0 = tensor.empty() : tensor<2x16x32xf32>
+  %1 = linalg.softmax dimension(2) ins(%arg0 : tensor<2x16x32xf32>) outs(%0: tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
+  return %1 : tensor<2x16x32xf32>
+}
+// CHECK:      func.func @softmax(%[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
+// CHECK:        %[[D0:.+]] = tensor.empty() : tensor<2x16x32xf32>
+// CHECK:        %[[D1:.+]] = linalg.softmax dimension(2) ins(%[[ARG0]] : tensor<2x16x32xf32>) outs(%[[D0]] :
+// CHECK-SAME:     tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
+// CHECK:        return %[[D1]] : tensor<2x16x32xf32>
+// CHECK:      }