Add tanh, add, and broadcast ops + Lowering to linalg #1595

sanjoy · 2022-11-16T05:55:09Z

This PR adds tanh, add and broadcast ops to TCP and implements a lowering to linalg.

This PR was carved out of a larger set of commits that includes a full Torch->TCP->Linalg e2e flow for these ops. This larger flow is here for reference.

(@navahgar is the original author of this code, I'm sending it out on his behalf.)

MaheshRavishankar

Thanks! First round of comments. Mostly small things. I dont fully follow the logic of broadcast lowering. Some comments there would be nice.

MaheshRavishankar · 2022-11-17T07:00:58Z

...s/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Conversion/Passes.h

+} // namespace mlir
+
+
+#endif // TORCH_MLIR_DIALECTS_CONVERSION_PASSES_H


Nit: Add line at EOF here and every where else.

...-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/Tcp/IR/TcpBase.td

MaheshRavishankar · 2022-11-17T07:04:23Z

...m-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/Tcp/IR/TcpOps.td

+          AllElementTypesMatch<["in", "out"]>,
+          PredOpTrait<
+            "attribute `axes` must be in increasing order",
+            CPred<"::llvm::is_sorted(getAxes(), "


I'd move this into utility functions and just call those from here instead of this in tablegen.

Upsignalling that this is still here

MaheshRavishankar · 2022-11-17T07:10:42Z

externals/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/Elementwise.cpp

+
+  // Add indexing maps for all the tensor operands and for the result.
+  for (size_t i = 0; i < tensorOperands.size(); ++i) {
+    indexingMaps.push_back(b.getMultiDimIdentityMap(resultRank));


All AffineMaps here are uniqued. You could just do

SmallVector<AffineMap> indexingMaps(tensorOperands.size() + 1, b.getMultiDimIdentityMap(resultRank));

MaheshRavishankar · 2022-11-17T07:12:47Z

externals/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/Elementwise.cpp

+}
+
+template<typename TcpOpT>
+Value createLinalgPayloadForElementwiseOp(


I am not entirely sure why this is templated. It would probably be the same if you use Operation *op instead of TcpOpT op?

Maybe also make this return FailureOr<Value>. Not a fan of returning nullptrs around.

MaheshRavishankar · 2022-11-17T07:14:15Z

externals/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/Elementwise.cpp

+                                            ConversionTarget &target) {
+  MLIRContext *context = patterns.getContext();
+
+  target.addIllegalOp<TanhOp>();


Wouldnt it be better to just add TCP dialect as illegal.

MaheshRavishankar · 2022-11-17T07:15:32Z

externals/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/Misc.cpp

+
+namespace {
+
+void getValuesFromIndexArrayAttribute(ArrayAttr attr,


You can just create a SmallVector<int64_t> within this method and return it instead of return by reference? I think most compilers optimize the return as a move.

MaheshRavishankar · 2022-11-17T07:15:55Z

externals/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/Misc.cpp

+    SmallVector<Value> resultDimSizes;
+    SmallVector<AffineExpr> exprs;
+    int64_t pos = 0;
+    for (int64_t i = 0; i < resultRank; ++i) {


Some comments here would be nice.

Didn't add explicit comments, but refactored the code to be more understandable. LMK if this still needs comments.

MaheshRavishankar · 2022-11-17T07:17:41Z

...als/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/PopulatePatterns.h

+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace tcp_to_linalg {


If you plan to move this into MLIR core maybe this should follow LLVM coding conventions, i.e. camel case instead of snake case.

sanjoy · 2022-11-19T06:38:35Z

@MaheshRavishankar Addressed comments, PTAL.

MaheshRavishankar

Thanks! Next round of comments, but mostly looks fine now.
Please make sure all files have a new line at end of file.

MaheshRavishankar · 2022-11-19T21:32:49Z

...m-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/Tcp/IR/TcpOps.td

+          AllElementTypesMatch<["in", "out"]>,
+          PredOpTrait<
+            "attribute `axes` must be in increasing order",
+            CPred<"::llvm::is_sorted(getAxes(), "


Upsignalling that this is still here

MaheshRavishankar · 2022-11-19T21:53:00Z

externals/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/Elementwise.cpp

+  // the tensorOperands, since all the operands are expected to have the same
+  // shape.
+  auto tensorOperand = tensorOperands[0];
+  for (int64_t i = 0; i < resultRank; ++i)


Please add { } around bodies that span multiple lines (even if it is a single statement).

Also you can just use createDimValues method (here : https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h#L31 ). That returns a SmallVector<OpFoldResult> that you can just pass to build method of tensor.empty().

Nice! Changed to use createDimValues.

MaheshRavishankar · 2022-11-19T21:59:40Z

externals/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/Misc.cpp

+          pos < static_cast<int64_t>(axes.size()) && axes[pos] == i;
+      if (isOutputDimBroadcasted) {
+        resultDimSizes.push_back(op->getOperands()[pos + 1]);
+        inputIndexingMapExprs.push_back(b.getAffineConstantExpr(0));


This is the design for broadcast chosen here, but just pointing out that having 0 in indexing maps of linalg ops is kind of bad form.

Given the current definition of tcp.broadcast, is there a better way to lower it to linalg?

No, there isnt. Just pointing out (minor) downside of the current design

MaheshRavishankar

Looks good to me. Dont know if anyone else needs to sign off as well.

MaheshRavishankar · 2022-11-22T07:41:23Z

externals/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/CMakeLists.txt

+  TorchMLIRTcpDialect
+)
+
+torch_mlir_dialects_target_includes(TorchMLIRTcpToLinalg)


Nit: new line at EOF

MaheshRavishankar · 2022-11-22T07:42:37Z

externals/llvm-external-projects/torch-mlir-dialects/lib/Conversion/TcpToLinalg/Misc.cpp

+          pos < static_cast<int64_t>(axes.size()) && axes[pos] == i;
+      if (isOutputDimBroadcasted) {
+        resultDimSizes.push_back(op->getOperands()[pos + 1]);
+        inputIndexingMapExprs.push_back(b.getAffineConstantExpr(0));


No, there isnt. Just pointing out (minor) downside of the current design

sjarus · 2022-11-22T07:56:37Z

Thanks @MaheshRavishankar for your time and support here!

* [Tcp] Add tanh, add, and broadcast ops * [Tcp] Add tcp-to-linalg pass Co-authored-by: Raghavan Raman <raghavan.raman@getcruise.com> Co-authored-by: Sanjoy Das <sanjoy.das@getcruise.com>

navahgar added 2 commits November 15, 2022 21:26

[Tcp] Add tanh, add, and broadcast ops

74b27c9

[Tcp] Add tcp-to-linalg pass

a95538a

sanjoy requested review from jpienaar and MaheshRavishankar November 16, 2022 05:55

MaheshRavishankar requested changes Nov 17, 2022

View reviewed changes

Sanjoy Das added 4 commits November 18, 2022 22:36

Address review comments

c98b4d6

clang format

29031a9

Add newlines at end of files

0663be9

Minor cleanups to better conform to LLVM style

7720b74

sanjoy requested a review from MaheshRavishankar November 19, 2022 06:38

MaheshRavishankar requested changes Nov 19, 2022

View reviewed changes

Address second round of review comments

ab305e7

sanjoy requested a review from MaheshRavishankar November 20, 2022 06:01

MaheshRavishankar approved these changes Nov 22, 2022

View reviewed changes

Add newline at EOF for CMakeLists.txt

2eb4f7c

sanjoy merged commit f5844cb into llvm:mlir-tcp Nov 22, 2022

navahgar mentioned this pull request Dec 8, 2022

[Tcp] Add lowering from Torch to Tcp for tanh, add, and broadcast ops + Enable e2e tests for Tcp #1695

Merged

sjarus mentioned this pull request Dec 21, 2022

[RFC] Bootstrapping the TCP Dialect in Torch-MLIR #1366

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add tanh, add, and broadcast ops + Lowering to linalg #1595

Add tanh, add, and broadcast ops + Lowering to linalg #1595

sanjoy commented Nov 16, 2022

MaheshRavishankar left a comment

MaheshRavishankar Nov 17, 2022

MaheshRavishankar Nov 17, 2022

MaheshRavishankar Nov 19, 2022

MaheshRavishankar Nov 17, 2022

MaheshRavishankar Nov 17, 2022

MaheshRavishankar Nov 17, 2022

MaheshRavishankar Nov 17, 2022

MaheshRavishankar Nov 17, 2022

MaheshRavishankar Nov 17, 2022

sanjoy Nov 19, 2022

MaheshRavishankar Nov 17, 2022

sanjoy commented Nov 19, 2022

MaheshRavishankar left a comment

MaheshRavishankar Nov 19, 2022

MaheshRavishankar Nov 19, 2022 •

edited

Loading

sanjoy Nov 20, 2022

MaheshRavishankar Nov 19, 2022

sanjoy Nov 20, 2022

MaheshRavishankar Nov 22, 2022

MaheshRavishankar left a comment

MaheshRavishankar Nov 22, 2022

MaheshRavishankar Nov 22, 2022

sjarus commented Nov 22, 2022

		} // namespace mlir


		#endif // TORCH_MLIR_DIALECTS_CONVERSION_PASSES_H


		namespace {

		void getValuesFromIndexArrayAttribute(ArrayAttr attr,

Add tanh, add, and broadcast ops + Lowering to linalg #1595

Add tanh, add, and broadcast ops + Lowering to linalg #1595

Conversation

sanjoy commented Nov 16, 2022

MaheshRavishankar left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sanjoy commented Nov 19, 2022

MaheshRavishankar left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

MaheshRavishankar Nov 19, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

MaheshRavishankar left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sjarus commented Nov 22, 2022

MaheshRavishankar Nov 19, 2022 •

edited

Loading