[mlir][tosa] : adding folder and canonicalizer for select

[lipracer]

Summary

define canonicalizer and folder for tosa::select

[lipracer]

update patch rebase on latest

[mehdi_amini]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

624

What is the issue with clang-format here? It is fairly rare that we have to disable clang-format, and in general it is always about some large struct initializers or things similar.

653

Nit: remove trivial braces, here and elsewhere.

661

Shouldn't this be part of a verifier?

[lipracer]

refine

[lipracer]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

624

There is an indenter on line 593 on the left of difference. I don't know why. I feel the format is strange.

[Mogball]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

195

how about using hasCanonicalizeMethod?

204–206

[jpienaar]

mlir/test/Dialect/Tosa/canonicalize.mlir

256

smae -> same

[mehdi_amini]

LG, after you look into @Mogball's comments

[lipracer]

wrap setOperands with updateRootInPlace

[lipracer]

fix typo

[lipracer]

format code style with clang-format

[Mogball]

LGTM

[lipracer]

Thanks,Can you help me commit this? user.name: lipracer user.email: [e-mail]

[Mogball]

gotchu

[hpmorgan]

Hi @lipracer, this change has caused a build error when building with clang:

/usr/bin/clang++ -DGTEST_HAS_RTTI=0 -DMLIR_CUDA_CONVERSIONS_ENABLED=0 -DMLIR_ROCM_CONVERSIONS_ENABLED=0 -D_DEBUG -D_GNU_SOURCE -DSTDC_CONSTANT_MACROS -DSTDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -I/mnt/vss/_work/1/b/llvm/Release/tools/mlir/lib/Dialect/Tosa -I/mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Tosa -I/mnt/vss/_work/1/b/llvm/Release/include -I/mnt/vss/_work/1/llvm-project/llvm/include -I/mnt/vss/_work/1/llvm-project/mlir/include -I/mnt/vss/_work/1/b/llvm/Release/tools/mlir/include -fPIC -fvisibility-inlines-hidden -Werror -Werror=date-time -Werror=unguarded-availability-new -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wmissing-field-initializers -pedantic -Wno-long-long -Wc++98-compat-extra-semi -Wimplicit-fallthrough -Wcovered-switch-default -Wno-noexcept-type -Wnon-virtual-dtor -Wdelete-non-virtual-dtor -Wsuggest-override -Wstring-conversion -Wmisleading-indentation -fdiagnostics-color -ffunction-sections -fdata-sections -Werror=mismatched-tags -Werror=global-constructors -O3 -DNDEBUG -fno-exceptions -fno-rtti -UNDEBUG -std=c++14 -MD -MT tools/mlir/lib/Dialect/Tosa/CMakeFiles/obj.MLIRTosa.dir/IR/TosaOps.cpp.o -MF tools/mlir/lib/Dialect/Tosa/CMakeFiles/obj.MLIRTosa.dir/IR/TosaOps.cpp.o.d -o tools/mlir/lib/Dialect/Tosa/CMakeFiles/obj.MLIRTosa.dir/IR/TosaOps.cpp.o -c /mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
/mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp:599:26: error: extra ';' outside of a function is incompatible with C++98 [-Werror,-Wc++98-compat-extra-semi]
ReduceFolder(ReduceAllOp);

[lipracer]

This comment was removed by lipracer.

[lipracer]

In D121513#3411651, @hpmorgan wrote:

Hi @lipracer, this change has caused a build error when building with clang:

/usr/bin/clang++ -DGTEST_HAS_RTTI=0 -DMLIR_CUDA_CONVERSIONS_ENABLED=0 -DMLIR_ROCM_CONVERSIONS_ENABLED=0 -D_DEBUG -D_GNU_SOURCE -DSTDC_CONSTANT_MACROS -DSTDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -I/mnt/vss/_work/1/b/llvm/Release/tools/mlir/lib/Dialect/Tosa -I/mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Tosa -I/mnt/vss/_work/1/b/llvm/Release/include -I/mnt/vss/_work/1/llvm-project/llvm/include -I/mnt/vss/_work/1/llvm-project/mlir/include -I/mnt/vss/_work/1/b/llvm/Release/tools/mlir/include -fPIC -fvisibility-inlines-hidden -Werror -Werror=date-time -Werror=unguarded-availability-new -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wmissing-field-initializers -pedantic -Wno-long-long -Wc++98-compat-extra-semi -Wimplicit-fallthrough -Wcovered-switch-default -Wno-noexcept-type -Wnon-virtual-dtor -Wdelete-non-virtual-dtor -Wsuggest-override -Wstring-conversion -Wmisleading-indentation -fdiagnostics-color -ffunction-sections -fdata-sections -Werror=mismatched-tags -Werror=global-constructors -O3 -DNDEBUG -fno-exceptions -fno-rtti -UNDEBUG -std=c++14 -MD -MT tools/mlir/lib/Dialect/Tosa/CMakeFiles/obj.MLIRTosa.dir/IR/TosaOps.cpp.o -MF tools/mlir/lib/Dialect/Tosa/CMakeFiles/obj.MLIRTosa.dir/IR/TosaOps.cpp.o.d -o tools/mlir/lib/Dialect/Tosa/CMakeFiles/obj.MLIRTosa.dir/IR/TosaOps.cpp.o -c /mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
/mnt/vss/_work/1/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp:599:26: error: extra ';' outside of a function is incompatible with C++98 [-Werror,-Wc++98-compat-extra-semi]
ReduceFolder(ReduceAllOp);

I have fixed https://reviews.llvm.org/D122599, Could you please review it, It's very strange why we get normal clang-format results when we change the macro definition to uppercase.

[void]

This change is resulting in these warnings:

/usr/local/google/home/morbo/llvm/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp:599:26: warning: extra ';' outside of a function is incompatible with C++98 [-Wc++98-compat-extra-semi]
ReduceFolder(ReduceAllOp);
                         ^
/usr/local/google/home/morbo/llvm/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp:600:26: warning: extra ';' outside of a function is incompatible with C++98 [-Wc++98-compat-extra-semi]
ReduceFolder(ReduceAnyOp);
                         ^
/usr/local/google/home/morbo/llvm/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp:601:26: warning: extra ';' outside of a function is incompatible with C++98 [-Wc++98-compat-extra-semi]
ReduceFolder(ReduceMaxOp);
                         ^
/usr/local/google/home/morbo/llvm/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp:602:26: warning: extra ';' outside of a function is incompatible with C++98 [-Wc++98-compat-extra-semi]
ReduceFolder(ReduceMinOp);
                         ^
/usr/local/google/home/morbo/llvm/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp:603:27: warning: extra ';' outside of a function is incompatible with C++98 [-Wc++98-compat-extra-semi]
ReduceFolder(ReduceProdOp);
                          ^
/usr/local/google/home/morbo/llvm/llvm-project/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp:604:26: warning: extra ';' outside of a function is incompatible with C++98 [-Wc++98-compat-extra-semi]
ReduceFolder(ReduceSumOp);
                         ^
6 warnings generated.

[mehdi_amini]

In D121513#3412417, @void wrote:

This change is resulting in these warnings:

I think this is what was mention right above? Can you check if the follow-up revision linked above fixed it for you as well?

[void]

In D121513#3412675, @mehdi_amini wrote:

In D121513#3412417, @void wrote:

This change is resulting in these warnings:

I think this is what was mention right above? Can you check if the follow-up revision linked above fixed it for you as well?

It looks like it did the trick. Thanks!

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>, NoSideEffect]> {
   let summary = "Elementwise select operator";
 
   let description = [{
     Elementwise select of the output based on a condition.
   }];
 
   let arguments = (ins
-    I1Tensor:$input1,
-    Tosa_Tensor:$input2,
-    Tosa_Tensor:$input3
+    I1Tensor:$pred,
+    Tosa_Tensor:$on_true,
+    Tosa_Tensor:$on_false
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
+  let hasCanonicalizeMethod = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.7
 // Operator Class: Logical Operations.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

 };
 
 void ReshapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
   results.add<ReshapeReshapeOptimization>(context);
   results.add<ReshapeConstOptimization>(context);
 }
 
+LogicalResult SelectOp::canonicalize(SelectOp op, PatternRewriter &rewriter) {
+  auto notOp = op.pred().getDefiningOp<tosa::LogicalNotOp>();
+  if (!notOp)
+    return failure();
+  rewriter.updateRootInPlace(op, [&]() {
+    op.getOperation()->setOperands(
+        {notOp.input1(), op.on_false(), op.on_true()});
+  });
+  return success();
+}
+
 struct ConstantTransposeOptimization
     : public OpRewritePattern<tosa::TransposeOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tosa::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto outputType = op.getType().cast<ShapedType>();
     ArrayRef<int64_t> outputShape = outputType.getShape();
     // TOSA supports quantized types.
OpFoldResult OP::fold(ArrayRef<Attribute> operands) {                        \
     ShapedType inputTy = input().getType().cast<ShapedType>();                 \
     if (!inputTy.hasRank())                                                    \
       return {};                                                               \
     if (inputTy.getDimSize(axis()) == 1)                                       \
       return input();                                                          \
     return {};                                                                 \
   }
 
-ReduceFolder(ReduceAllOp) ReduceFolder(ReduceAnyOp) ReduceFolder(ReduceMaxOp)
-    ReduceFolder(ReduceMinOp) ReduceFolder(ReduceProdOp)
-        ReduceFolder(ReduceSumOp)
+ReduceFolder(ReduceAllOp);
+ReduceFolder(ReduceAnyOp);
+ReduceFolder(ReduceMaxOp);
+ReduceFolder(ReduceMinOp);
+ReduceFolder(ReduceProdOp);
+ReduceFolder(ReduceSumOp);
 #undef ReduceFolder
 
-            OpFoldResult ReshapeOp::fold(ArrayRef<Attribute> operands) {
   auto inputTy = input1().getType().dyn_cast<RankedTensorType>();
   auto outputTy = getType().dyn_cast<RankedTensorType>();
 
   if (!inputTy || !outputTy || inputTy != outputTy)
     return {};
   return input1();
 }
 
 OpFoldResult PadOp::fold(ArrayRef<Attribute> operands) {
   // If the pad is all zeros we can fold this operation away.
   if (operands[1]) {
     auto densePad = operands[1].cast<DenseElementsAttr>();
     if (densePad.isSplat() && densePad.getSplatValue<APInt>().isZero()) {
       return input1();
     }
   }
 
   return {};
 }
 
 OpFoldResult SliceOp::fold(ArrayRef<Attribute> operands) {
   auto inputTy = input().getType().dyn_cast<RankedTensorType>();
   auto outputTy = getType().dyn_cast<RankedTensorType>();
 
   if (!inputTy || !outputTy || inputTy != outputTy)
     return {};
   if (inputTy.hasStaticShape())
     return input();
 
   return {};
 }
 
+OpFoldResult tosa::SelectOp::fold(ArrayRef<Attribute> operands) {
+  if (on_true() == on_false())
+    return on_true();
+
+  auto predicate = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
+  if (!predicate)
+    return {};
+
+  if (!predicate.isSplat())
+    return {};
+  return predicate.getSplatValue<APInt>().getBoolValue() ? on_true()
+                                                         : on_false();
+}
+
 OpFoldResult TileOp::fold(ArrayRef<Attribute> operands) {
   bool allOnes = true;
   for (Attribute val : multiples().getValue()) {
     allOnes = allOnes && val.cast<IntegerAttr>().getValue().getSExtValue() == 1;
   }
 
   if (allOnes && input1().getType() == getType())
     return input1();
   return {};
 }
 
 OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
   if (!operands[1])
     return {};
 
   // Transposing splat values just means reshaping.
if (resultKnowledge.size() != yieldOp.getNumOperands())
       return failure();
 
     for (const auto &it : llvm::enumerate(yieldOp.getOperands())) {
       int32_t index = it.index();
       if (auto meet = ValueKnowledge::meet(
               resultKnowledge[index],
               ValueKnowledge::getKnowledgeFromType(it.value().getType()))) {
         resultKnowledge[index] = meet;
-      };
+      }
     }
   }
 
   for (const ValueKnowledge &result : resultKnowledge) {
     inferredReturnShapes.push_back(result.getShapedTypeComponents());
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Operator Definitions.
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"

mlir/test/Dialect/Tosa/canonicalize.mlir

func @mul_one_int(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> {
   // CHECK-NOT: tosa.mul
   %ones = "tosa.const"() {value = dense<1> : tensor<2x3xi32>} : () -> tensor<2x3xi32>
   %1 = "tosa.mul"(%arg0, %ones) {shift = 0 : i32} : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
   return %1 : tensor<2x3xi32>
 }
 
 // -----
 
+// CHECK-LABEL: @select_same_value
+func @select_same_value(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
+  %0 = "tosa.select"(%arg0, %arg1, %arg1) : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
+  // CHECK: return %arg1
+  // CHECK-NOT: tosa.select
+  return %0 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @select_true_value
+func @select_true_value(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
+  %c1 = "tosa.const"() {value = dense<1> : tensor<2x3xi1>} : () -> tensor<2x3xi1>
+  %0 = "tosa.select"(%c1, %arg0, %arg1) : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
+  // CHECK: return %arg0
+  // CHECK-NOT: tosa.select
+  return %0 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @select_false_value
+func @select_false_value(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
+  %c0 = "tosa.const"() {value = dense<0> : tensor<2x3xi1>} : () -> tensor<2x3xi1>
+  %0 = "tosa.select"(%c0, %arg0, %arg1) : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
+  // CHECK: return %arg1
+  // CHECK-NOT: tosa.select
+  return %0 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @select_not_pred
+func @select_not_pred(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>, %arg2: tensor<2x3xi32>) -> tensor<2x3xi32> {
+  %0 = "tosa.logical_not"(%arg0) : (tensor<2x3xi1>) -> tensor<2x3xi1>
+  %1 = "tosa.select"(%0, %arg1, %arg2) : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
+  // CHECK: "tosa.select"(%arg0, %arg2, %arg1)
+  return %1 : tensor<2x3xi32>
+}
+
+// -----
+
 // CHECK-LABEL: @reduce_all_fold
 func @reduce_all_fold(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
   // CHECK: return %arg0
   %0 = "tosa.reduce_all"(%arg0) {axis = 1 : i64}: (tensor<?x1xf32>) -> tensor<?x1xf32>
   return %0 : tensor<?x1xf32>
 }
 
 // -----