[PATCH] D153422: [mlir][Linalg] Add a softmax op

[Nicolas Vasilache via Phabricator via llvm-commits]

nicolasvasilache accepted this revision.
nicolasvasilache added inline comments.
This revision is now accepted and ready to land.


================
Comment at: mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td:100
+    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)
----------------
rengolin wrote:
> 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?
let's extend on a per-need basis if you don't mind, proabably it makes sense yes


================
Comment at: mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td:125
+  let extraClassDeclaration = [{
+    Value input() {
+      return getInput();
----------------
chelini wrote:
> Any specific reason for having this wrapper?  Can we use the method generated by tablegen directly `getInput`? Same for `getOutput`.
yes let's drop this


================
Comment at: mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td:144
+    // Method to implement for specifying output range for
+    // DestinationStyleOpInterface
+    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
----------------
chelini wrote:
> I cannot fully understand the sentence, maybe something like: `Implement functions necessary for DestinationStyleOpInterface.`
`Method to implement DestinationStyleOpInterface`


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2166
+LogicalResult SoftmaxOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) {
+  return memref::foldMemRefCast(*this);
+}
----------------
chelini wrote:
> Can you please comment on why we need this folder?
this is a general folder that rewrites `cast(dynamic) -> static`



Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D153422/new/

https://reviews.llvm.org/D153422