[MLIR][TORCH] Add -convert-torch-to-tosa-custom pass for undecomposed selective ops #1514

AmosLewis · 2022-10-20T18:42:16Z

Use tosa::custom op to represent softmax op which we deliberately choose not to decompose into simple ops.The reason to do this is in this RFC here: #1519

ramiro050 · 2022-10-20T19:15:26Z

The CI error is due to this patch. The program is crashing when creating the tosa::CustomOp:

https://github.com/llvm/torch-mlir/actions/runs/3292098619/jobs/5427059644#step:6:6218

ramiro050 · 2022-10-20T19:33:24Z

I don't think this is the right approach for what this patch is trying to achieve. As a user of torch-mlir, I would expect that running the torch-to-tosa pass on any set of supported ops will result in standard tosa ops. Users should not be expected by default to handle arbitrary custom ops. To use custom ops should be a choice made by the user.

I think this requires a more general solution where users can specify which torch-ops to pass down as custom ops, and have the conversion to tosa::CustomOp happen automatically. Since this will require some non-trivial changes to several parts of torch-mlir, can you open an RFC issue with what you're tryinig to achieve, and we can brainstorm on the best solution.

AmosLewis · 2022-10-20T19:35:33Z

The CI error is due to this patch. The program is crashing when creating the tosa::CustomOp:

https://github.com/llvm/torch-mlir/actions/runs/3292098619/jobs/5427059644#step:6:6218

The werid thing is when I run https://github.com/llvm/torch-mlir/actions/runs/3292098619/jobs/5427059644#step:6:6218:~:text=6191-,%2D%2D,%2D%2D,-6194 on my local machine, everything looks good. Here is my local run cmd print:

➜  torch-mlir git:(tosa-custom-softmax) ✗ torch-mlir-opt test/Conversion/TorchToTosa/basic.mlir -convert-torch-to-tosa -split-input-file -verify-diagnostics | FileCheck  test/Conversion/TorchToTosa/basic.mlir
➜  torch-mlir git:(tosa-custom-softmax) ✗

AmosLewis · 2022-10-20T19:42:34Z

I don't think this is the right approach for what this patch is trying to achieve. As a user of torch-mlir, I would expect that running the torch-to-tosa pass on any set of supported ops will result in standard tosa ops. Users should not be expected by default to handle arbitrary custom ops. To use custom ops should be a choice made by the user.

I think this requires a more general solution where users can specify which torch-ops to pass down as custom ops, and have the conversion to tosa::CustomOp happen automatically. Since this will require some non-trivial changes to several parts of torch-mlir, can you open an RFC issue with what you're tryinig to achieve, and we can brainstorm on the best solution.

Yes. The RFC is on the way by my partner. I just try to do something before the RFC is done. This is the initial step for demonstrating custom op could be used to handle complex ops that we don't want to decompose. There should be a pass/list where we add ops and only ops from this list should go through the custom op route, the rest of op should be lower normally.

silvasean · 2022-10-21T08:36:14Z

Let's wait for the RFC to discuss more. Per-op custom patterns are definitely not the right way to do this. Maybe we could have a pass torch-to-tosa-convert-to-custom{op-list=aten.softmax.int,...} which systematically converts any Torch op to the appropriate tosa.custom.

sjarus · 2022-10-22T04:20:24Z

This is an interesting PR. Integer domain softmax could be done using a custom op, but need not be. There is a working and bit accurate decomposition of TFLite softmax that is fairly detailed but works: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tosa/transforms/legalize_common.cc#L1375

Please note that the legalize_common.cc there is the same as TosaLegalizeCommon.cc - the goal here is to eventually move these builders into the llvm-project TOSA codebase.

The straightforward solution here is to copy the convertSoftMaxOp() into TosaLegalizeCommon.cpp, then just call that. It should do all the work for you. This approach is used by several other TorchToTosa legalizations that invoke calls to TosaLegalizeCommon to implement legalizations using previously solved pattern replacements.

Note: the goal is to keep TosaLegalizeCommon.* and TosaLegalizeUtils.* files essentially identical to their TF counterparts legalize_common.* and legalize_utils.* so that a followon PR can move them upstream since multiple projects can all demonstrably depend on the same conversions.

sjarus

Please see previous comment.

AmosLewis · 2022-10-24T19:56:23Z

Let's wait for the RFC to discuss more. Per-op custom patterns are definitely not the right way to do this. Maybe we could have a pass torch-to-tosa-convert-to-custom{op-list=aten.softmax.int,...} which systematically converts any Torch op to the appropriate tosa.custom.

The RFC is ready here. #1519 @ramiro050 @sjarus

sjarus · 2022-10-24T23:29:57Z

Adding @eric-k256

AmosLewis · 2022-10-25T19:44:47Z

Instead of add a new pass torch-to-tosa-convert-to-custom{op-list=aten.softmax.int,...}. I add an options in the exist pass -convert-torch-to-tosa. Here is my most recent local result:

➜  torch-mlir git:(tosa-custom-softmax) ✗ torch-mlir-opt -convert-torch-to-tosa="custom-ops=torch.aten.softmax.int"    test/Conversion/TorchToTosa/custom.mlir  | externals/llvm-project/mlir/utils/generate-test-checks.py

// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py

// The script is designed to make adding checks to
// a test case fast, it is *not* designed to be authoritative
// about what constitutes a good test! The CHECK should be
// minimized and named to reflect the test intent.



// CHECK-LABEL:   func.func @torch.aten.softmax.int$cst_dim(
// CHECK-SAME:                                              %[[VAL_0:.*]]: !torch.vtensor<[2,3],f32>) -> !torch.vtensor<[2,3],f32> {
// CHECK:           %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[2,3],f32> -> tensor<2x3xf32>
// CHECK:           %[[VAL_2:.*]] = torch.constant.none
// CHECK:           %[[VAL_3:.*]] = torch.constant.int 1
// CHECK:           %[[VAL_4:.*]] = "tosa.custom"(%[[VAL_1]]) {dim = 1 : i64, identifier = "softmax"} : (tensor<2x3xf32>) -> tensor<2x3xf32>
// CHECK:           %[[VAL_5:.*]] = torch_c.from_builtin_tensor %[[VAL_4]] : tensor<2x3xf32> -> !torch.vtensor<[2,3],f32>
// CHECK:           return %[[VAL_5]] : !torch.vtensor<[2,3],f32>
// CHECK:         }

AmosLewis · 2022-10-27T01:10:01Z

Instead of add a new pass torch-to-tosa-convert-to-custom{op-list=aten.softmax.int,...}. I add an options in the exist pass -convert-torch-to-tosa. Here is my most recent local result:

Based on Svoch requirements, move dim into input tensor instead of attributes, Here is the new output:

➜  torch-mlir git:(tosa-custom-softmax) ✗ torch-mlir-opt -convert-torch-to-tosa="custom-ops=torch.aten.softmax.int"    test/Conversion/TorchToTosa/custom.mlir  | externals/llvm-project/mlir/utils/generate-test-checks.py
// CHECK-LABEL:   func.func @torch.aten.softmax.int$cst_dim(
// CHECK-SAME:                                              %[[VAL_0:.*]]: !torch.vtensor<[2,3],f32>) -> !torch.vtensor<[2,3],f32> {
// CHECK:           %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[2,3],f32> -> tensor<2x3xf32>
// CHECK:           %[[VAL_2:.*]] = torch.constant.none
// CHECK:           %[[VAL_3:.*]] = torch.constant.int 1
// CHECK:           %[[VAL_4:.*]] = "tosa.const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64>
// CHECK:           %[[VAL_5:.*]] = "tosa.custom"(%[[VAL_1]], %[[VAL_4]]) {identifier = "softmax"} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x3xf32>
// CHECK:           %[[VAL_6:.*]] = torch_c.from_builtin_tensor %[[VAL_5]] : tensor<2x3xf32> -> !torch.vtensor<[2,3],f32>
// CHECK:           return %[[VAL_6]] : !torch.vtensor<[2,3],f32>
// CHECK:         }

AmosLewis · 2022-10-27T01:13:37Z

Just got an E2E test bug about the option I added:

python -m e2e_testing.main --config=tosa -v
: CommandLine Error: Option 'custom-ops' registered more than once!
LLVM ERROR: inconsistency in registered CommandLine options

ramiro050 · 2022-10-27T23:03:51Z

include/torch-mlir/Dialect/TorchConversion/Transforms/Passes.h

+  Option<bool> custom{*this, "custom-ops",
+                      llvm::cl::desc("custom complex operations."),
+                      llvm::cl::init(true)};
+  ListOption<std::string> customOps{
+      *this, "custom-ops",
+      llvm::cl::desc("List of ops to be converted to the backend."),
+      llvm::cl::ZeroOrMore};


I think the build issue is here. There are two options both with the same name "custom-ops"

Make sense. Thanks, Ramiro. I have this fixed, could you review and approved this patch if everything else makes sense?

AmosLewis · 2022-10-28T18:52:35Z

Please see previous comment.

Please review this patch again. The purpose of this patch is to work with the op that we choose to not decompose.

eric-k256

This feels halfway between being a one-off solution for torch.aten.softmax.int, and providing a flexible way to handle many ops in this manner. With some relatively small changes, I think it can be an example use of the flexible mechanism.

eric-k256 · 2022-10-28T19:00:55Z

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -3854,6 +3902,9 @@ class ConvertTorchToTosa : public ConvertTorchToTosaBase<ConvertTorchToTosa> {
    INSERT_ATENOP_PATTERN(AtenSliceTensorOp);
    INSERT_ATENOP_PATTERN(AtenBroadcastToOp);
    INSERT_ATENOP_PATTERN(AtenWhereSelfOp);
+    if(!customOps.empty()){


This feels awkward, and isn't checking the value of the customOps to make sure that softmax is in the customOp list. I think a better option is to always insert the pattern, and then at the top of the softmax matchAndRewrite, search for torch.aten.softmax.int in the string, and use that to determine whether to create the custom op.

Some checking that the customOp list is valid would be good. With this change, someone could put any op in the customOp list and get no warning that the pass will not do what is being asked. At least having a list of ops that support legalization to tosa.custom could be used to warn if an op was passed that wasn't expected.

eric-k256 · 2022-10-28T19:01:43Z

include/torch-mlir/Dialect/TorchConversion/Transforms/Passes.h

+    : public PassPipelineOptions<TosaBackendPipelineOptions> {
+  // If this option is true, custom complex operations.
+  // If this option is false, skip decomposition of complex operations.
+  Option<bool> enableCustomOps{*this, "enableCustomOps",


I don't see any other reference to this, is this option needed, or is the custom-ops ListOption sufficient?

ramiro050 · 2022-11-01T17:27:00Z

This feels halfway between being a one-off solution for torch.aten.softmax.int, and providing a flexible way to handle many ops in this manner. With some relatively small changes, I think it can be an example use of the flexible mechanism.

Also, isn't custom() being deprecated or at least being replaced with interface() aligned with ml_program ?

I think that torch softmax quantized is better off leveraging the existing functional lowering into TosaLegalizeCommon that is bit accurate for tfl.softmax.

@AmosLewis, can you address this?

AmosLewis · 2022-11-01T19:21:22Z

interface()

I didn't find out what does the interface() means. Could you point out where is the code for it?

eric-k256 · 2022-11-01T19:48:24Z

We had thoughts of making bigger changes to tosa.custom, and calling it tosa.interface. That hasn't happened, but I do have a proposed change to tosa.custom in the dialect here: https://reviews.llvm.org/D137133
I think that makes the generic Operation* mapping simpler, as you can now encode attributes into the implementation_attrs

The other point Suraj was making is that we do have a legalization for softmax int into a series of TOSA ops that could be used to work aorund this specific problem. The larger problem of conversion from torch to tosa.custom remains, which the pass described by Ramiro aims to solve.

AmosLewis · 2022-11-04T08:33:53Z

@ramiro050 I just added a generic MatchAndRewrite for torch-to-tosa-custom pass. Now, there is no need to write new MatchAndRewrite for new ops. It worked for the 2 examples in custom.mlir.We can add more type conversion support into it later. Please review.

This patch is not a special fix for softmax, please review the new update.

ramiro050 · 2022-11-04T16:50:37Z

lib/Conversion/TorchToTosa/TorchToTosaCustom.cpp

+
+namespace {
+
+template <typename AtenOpT>


This shouldn't be a template. It should be a class pattern that works on Operation * inputs. Here's an example:

torch-mlir/lib/Conversion/TorchToLinalg/Uncategorized.cpp

Lines 1031 to 1038 in fedf8c0

class ConvertElementwiseOp : public ConversionPattern {

public:

ConvertElementwiseOp(TypeConverter &typeConverter, MLIRContext *context)

: ConversionPattern(typeConverter, MatchAnyOpTypeTag(), /*benefit=*/1,

context) {}

LogicalResult

matchAndRewrite(Operation *op, ArrayRef<Value> operands,

ramiro050 · 2022-11-04T16:51:02Z

lib/Conversion/TorchToTosa/TorchToTosaCustom.cpp

+  LogicalResult
+  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    ValueRange adaptor_operands = adaptor.getOperands();


In torch-mlir, we use camel-case

ramiro050 · 2022-11-04T16:51:26Z

lib/Conversion/TorchToTosa/TorchToTosaCustom.cpp

+    int num_operands = adaptor_operands.size();
+    std::vector<mlir::Value> inputs_vec;
+    for (int i = 0; i < num_operands; i++) {
+      auto operand = *op.getODSOperands(i).begin();


You can you this by calling operand = op.getOperand(i);

ramiro050 · 2022-11-04T16:52:49Z

lib/Conversion/TorchToTosa/TorchToTosaCustom.cpp

+      auto operand = *op.getODSOperands(i).begin();
+      auto adaptor_operand_type = adaptor_operands[i].getType();
+      // type convert for operands
+      if (adaptor_operand_type


There shouldn't be a need to check the adaptor types. matchPattern will check if the input is of the expected type or not.

ramiro050 · 2022-11-04T16:54:11Z

lib/Conversion/TorchToTosa/TorchToTosaCustom.cpp

+        if (!matchPattern(operand, m_TorchConstantFloat(&operand_tosa)))
+          return rewriter.notifyMatchFailure(
+              op, "unimplemented: operand should be a torch.constant.float");
+        auto operand_tensor_float = tosa::getConstTensor<int64_t>(


The type of this tensor should be double, right?

ramiro050 · 2022-11-04T17:00:32Z