[PATCH] D154363: [WIP][mlir] Add an interface to decompose complex ops

[lorenzo chelini via Phabricator via llvm-commits]

chelini added inline comments.


================
Comment at: mlir/include/mlir/Interfaces/AggregatedOpInterface.td:1
+//===- TilingInterface.td - Interface for tiling operations *- tablegen -*-===//
+//
----------------
`AggregatedOpInterface.td`


================
Comment at: mlir/include/mlir/Interfaces/AggregatedOpInterface.td:2
+//===- TilingInterface.td - Interface for tiling operations *- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
----------------
Why this interface is not part of Linalg? 


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2317
+  auto identityMap =
+      AffineMap::getMultiDimIdentityMap(inputRank, builder.getContext());
+  SmallVector<AffineExpr, 2> affineExprs;
----------------
nit: I would create a variable for the context to avoid `builder.getContext()` multiple times.


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2324
+  auto reductionMap =
+      AffineMap::get(inputRank, 0, affineExprs, builder.getContext());
+  SmallVector<AffineMap> indexingMaps{identityMap, reductionMap};
----------------
nit: /*symbols=*/0


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2332
+template <typename T>
+static Value reduce(Value input, Value output, int64_t dim, Location loc,
+                    OpBuilder &builder) {
----------------
nit: I think we want to have the builder as first arg. 


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2340
+  assert(indexingMaps.size() == 2 &&
+         "We should two maps: 1 for the input, 1 for the output");
+  assert(indexingMaps[0].isIdentity() && "input map should be identity");
----------------
"we should have two maps"


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2362
+      inputRank, dim, builder, /*allParallel=*/true);
+  assert(indexingMaps.size() == 2 && "We should one map for each input");
+  assert(indexingMaps[0].isIdentity() &&
----------------
nit: We should have..


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2363
+  assert(indexingMaps.size() == 2 && "We should one map for each input");
+  assert(indexingMaps[0].isIdentity() &&
+         "All parallel should give us identity map for input");
----------------
I would expect to have identity map for both input operands. Is this not the case?


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2382
+/// }
+static Value computeSoftmax(Value numerator, Value denominator, Value output,
+                            int64_t dim, Location loc, OpBuilder &builder) {
----------------
nit: I think naming is a bit misleading here. Simply, `buildDivOp`?


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2444
+  // Step 2: Subtract max from input and exponentiate.
+  Value numerator = subtractAndExp(input, max, outputNd, reductionDim, loc, b);
+
----------------
nit: I would rename: `subtractAndExp` to `buildSubAndExpOp`.


Repository:
  rG LLVM Github Monorepo

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

https://reviews.llvm.org/D154363