From 88151dd4285cdd9feeb24ebb1be9cf5252ab0883 Mon Sep 17 00:00:00 2001 From: Finn Plummer <50529406+inbelic@users.noreply.github.com> Date: Thu, 21 Dec 2023 09:24:01 -0800 Subject: [mlir][spirv] Add folding for SNegate, [Logical]Not (#74992) Add missing constant propogation folder for SNegate, [Logical]Not. Implement additional folding when !(!x) for all ops. This helps for readability of lowered code into SPIR-V. Part of work for #70704 --- .../mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td | 2 + mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBitOps.td | 2 + .../mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td | 1 + .../lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp | 55 +++++++++ .../Dialect/SPIRV/Transforms/canonicalize.mlir | 128 +++++++++++++++++++++ 5 files changed, 188 insertions(+) diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td index 51124e1..22d5afc 100644 --- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td +++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td @@ -582,6 +582,8 @@ def SPIRV_SNegateOp : SPIRV_ArithmeticUnaryOp<"SNegate", %3 = spirv.SNegate %2 : vector<4xi32> ``` }]; + + let hasFolder = 1; } // ----- diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBitOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBitOps.td index b460c8e..38639a1 100644 --- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBitOps.td +++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBitOps.td @@ -462,6 +462,8 @@ def SPIRV_NotOp : SPIRV_BitUnaryOp<"Not", [UsableInSpecConstantOp]> { %3 = spirv.Not %1 : vector<4xi32> ``` }]; + + let hasFolder = 1; } #endif // MLIR_DIALECT_SPIRV_IR_BIT_OPS diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td index 2e26c44..e48a56f 100644 --- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td +++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td @@ -534,6 +534,7 @@ def SPIRV_LogicalNotOp : SPIRV_LogicalUnaryOp<"LogicalNot", }]; let hasCanonicalizer = 1; + let hasFolder = 1; } // ----- diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp index 08ddc7c..4c62289 100644 --- a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp +++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp @@ -644,6 +644,45 @@ OpFoldResult spirv::UModOp::fold(FoldAdaptor adaptor) { } //===----------------------------------------------------------------------===// +// spirv.SNegate +//===----------------------------------------------------------------------===// + +OpFoldResult spirv::SNegateOp::fold(FoldAdaptor adaptor) { + // -(-x) = 0 - (0 - x) = x + auto op = getOperand(); + if (auto negateOp = op.getDefiningOp()) + return negateOp->getOperand(0); + + // According to the SPIR-V spec: + // + // Signed-integer subtract of Operand from zero. + return constFoldUnaryOp( + adaptor.getOperands(), [](const APInt &a) { + APInt zero = APInt::getZero(a.getBitWidth()); + return zero - a; + }); +} + +//===----------------------------------------------------------------------===// +// spirv.NotOp +//===----------------------------------------------------------------------===// + +OpFoldResult spirv::NotOp::fold(spirv::NotOp::FoldAdaptor adaptor) { + // !(!x) = x + auto op = getOperand(); + if (auto notOp = op.getDefiningOp()) + return notOp->getOperand(0); + + // According to the SPIR-V spec: + // + // Complement the bits of Operand. + return constFoldUnaryOp(adaptor.getOperands(), [&](APInt a) { + a.flipAllBits(); + return a; + }); +} + +//===----------------------------------------------------------------------===// // spirv.LogicalAnd //===----------------------------------------------------------------------===// @@ -714,6 +753,22 @@ OpFoldResult spirv::LogicalNotEqualOp::fold(FoldAdaptor adaptor) { // spirv.LogicalNot //===----------------------------------------------------------------------===// +OpFoldResult spirv::LogicalNotOp::fold(FoldAdaptor adaptor) { + // !(!x) = x + auto op = getOperand(); + if (auto notOp = op.getDefiningOp()) + return notOp->getOperand(0); + + // According to the SPIR-V spec: + // + // Complement the bits of Operand. + return constFoldUnaryOp(adaptor.getOperands(), + [](const APInt &a) { + APInt zero = APInt::getZero(1); + return a == 1 ? zero : (zero + 1); + }); +} + void spirv::LogicalNotOp::getCanonicalizationPatterns( RewritePatternSet &results, MLIRContext *context) { results diff --git a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir index 871ecd4..1cb6989 100644 --- a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir +++ b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir @@ -1007,6 +1007,102 @@ func.func @umod_fail_fold(%arg0: i32) -> (i32, i32) { // ----- //===----------------------------------------------------------------------===// +// spirv.SNegate +//===----------------------------------------------------------------------===// + +// CHECK-LABEL: @snegate_twice +// CHECK-SAME: (%[[ARG:.*]]: i32) +func.func @snegate_twice(%arg0 : i32) -> i32 { + %0 = spirv.SNegate %arg0 : i32 + %1 = spirv.SNegate %0 : i32 + + // CHECK: return %[[ARG]] : i32 + return %1 : i32 +} + +// CHECK-LABEL: @snegate_min +func.func @snegate_min() -> (i8, i8) { + // CHECK: %[[MIN:.*]] = spirv.Constant -128 : i8 + %cmin = spirv.Constant -128 : i8 + + %0 = spirv.SNegate %cmin : i8 + %1 = spirv.SNegate %0 : i8 + + // CHECK: return %[[MIN]], %[[MIN]] + return %0, %1 : i8, i8 +} + +// CHECK-LABEL: @const_fold_scalar_snegate +func.func @const_fold_scalar_snegate() -> (i32, i32, i32) { + %c0 = spirv.Constant 0 : i32 + %c3 = spirv.Constant 3 : i32 + %cn3 = spirv.Constant -3 : i32 + + // CHECK-DAG: %[[THREE:.*]] = spirv.Constant 3 : i32 + // CHECK-DAG: %[[NTHREE:.*]] = spirv.Constant -3 : i32 + // CHECK-DAG: %[[ZERO:.*]] = spirv.Constant 0 : i32 + %0 = spirv.SNegate %c0 : i32 + %1 = spirv.SNegate %c3 : i32 + %2 = spirv.SNegate %cn3 : i32 + + // CHECK: return %[[ZERO]], %[[NTHREE]], %[[THREE]] + return %0, %1, %2 : i32, i32, i32 +} + +// CHECK-LABEL: @const_fold_vector_snegate +func.func @const_fold_vector_snegate() -> vector<3xi32> { + // CHECK: spirv.Constant dense<[0, 3, -3]> + %cv = spirv.Constant dense<[0, -3, 3]> : vector<3xi32> + %0 = spirv.SNegate %cv : vector<3xi32> + return %0 : vector<3xi32> +} + +// ----- + +//===----------------------------------------------------------------------===// +// spirv.Not +//===----------------------------------------------------------------------===// + +// CHECK-LABEL: @not_twice +// CHECK-SAME: (%[[ARG:.*]]: i32) +func.func @not_twice(%arg0 : i32) -> i32 { + %0 = spirv.Not %arg0 : i32 + %1 = spirv.Not %0 : i32 + + // CHECK: return %[[ARG]] : i32 + return %1 : i32 +} + +// CHECK-LABEL: @const_fold_scalar_not +func.func @const_fold_scalar_not() -> (i32, i32, i32) { + %c0 = spirv.Constant 0 : i32 + %c3 = spirv.Constant 3 : i32 + %cn3 = spirv.Constant -3 : i32 + + // CHECK-DAG: %[[TWO:.*]] = spirv.Constant 2 : i32 + // CHECK-DAG: %[[NFOUR:.*]] = spirv.Constant -4 : i32 + // CHECK-DAG: %[[NONE:.*]] = spirv.Constant -1 : i32 + %0 = spirv.Not %c0 : i32 + %1 = spirv.Not %c3 : i32 + %2 = spirv.Not %cn3 : i32 + + // CHECK: return %[[NONE]], %[[NFOUR]], %[[TWO]] + return %0, %1, %2 : i32, i32, i32 +} + +// CHECK-LABEL: @const_fold_vector_not +func.func @const_fold_vector_not() -> vector<3xi32> { + %cv = spirv.Constant dense<[-1, -4, 2]> : vector<3xi32> + + // CHECK: spirv.Constant dense<[0, 3, -3]> + %0 = spirv.Not %cv : vector<3xi32> + + return %0 : vector<3xi32> +} + +// ----- + +//===----------------------------------------------------------------------===// // spirv.LogicalAnd //===----------------------------------------------------------------------===// @@ -1040,6 +1136,38 @@ func.func @convert_logical_and_true_false_vector(%arg: vector<3xi1>) -> (vector< // spirv.LogicalNot //===----------------------------------------------------------------------===// +// CHECK-LABEL: @logical_not_twice +// CHECK-SAME: (%[[ARG:.*]]: i1) +func.func @logical_not_twice(%arg0 : i1) -> i1 { + %0 = spirv.LogicalNot %arg0 : i1 + %1 = spirv.LogicalNot %0 : i1 + + // CHECK: return %[[ARG]] : i1 + return %1 : i1 +} + +// CHECK-LABEL: @const_fold_scalar_logical_not +func.func @const_fold_scalar_logical_not() -> i1 { + %true = spirv.Constant true + + // CHECK: spirv.Constant false + %0 = spirv.LogicalNot %true : i1 + + return %0 : i1 +} + +// CHECK-LABEL: @const_fold_vector_logical_not +func.func @const_fold_vector_logical_not() -> vector<2xi1> { + %cv = spirv.Constant dense<[true, false]> : vector<2xi1> + + // CHECK: spirv.Constant dense<[false, true]> + %0 = spirv.LogicalNot %cv : vector<2xi1> + + return %0 : vector<2xi1> +} + +// ----- + func.func @convert_logical_not_to_not_equal(%arg0: vector<3xi64>, %arg1: vector<3xi64>) -> vector<3xi1> { // CHECK: %[[RESULT:.*]] = spirv.INotEqual {{%.*}}, {{%.*}} : vector<3xi64> // CHECK-NEXT: spirv.ReturnValue %[[RESULT]] : vector<3xi1> -- cgit v1.1