Raise llvm.intr.abs to math.absi

src/enzyme_ad/jax/Passes/ArithRaising.cpp

@@ -682,6 +682,7 @@ struct ArithRaisingPass
         RaiseUnary<math::CountLeadingZerosOp, stablehlo::ClzOp, mhlo::ClzOp>,
         RaiseUnary<math::CtPopOp,       stablehlo::PopulationCountOp, mhlo::PopulationCountOp>,
         RaiseUnary<math::AbsFOp,        stablehlo::AbsOp,      mhlo::AbsOp>,
+        RaiseUnary<math::AbsIOp,        stablehlo::AbsOp,      mhlo::AbsOp>,
         RaiseUnary<math::IsFiniteOp,    stablehlo::IsFiniteOp, mhlo::IsFiniteOp>,
         RaiseUnary<math::CeilOp,        stablehlo::CeilOp,     mhlo::CeilOp>,
         RaiseUnary<math::FloorOp,       stablehlo::FloorOp,    mhlo::FloorOp>,

src/enzyme_ad/jax/Passes/LibDeviceFuncsRaisingPass.cpp

@@ -700,6 +700,17 @@ using ConvertFMFMathFromLLVMPattern =
 
 using AbsFOpLowering =
     ConvertFMFMathFromLLVMPattern<math::AbsFOp, LLVM::FAbsOp>;
+
+// llvm.intr.abs carries an is_int_min_poison flag arith has no place for;
+// drop it and raise to math.absi.
+struct AbsIOpRaising : public OpRewritePattern<LLVM::AbsOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(LLVM::AbsOp op,
+                                PatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<math::AbsIOp>(op, op.getIn());
+    return success();
+  }
+};
 using CeilOpLowering =
     ConvertFMFMathFromLLVMPattern<math::CeilOp, LLVM::FCeilOp>;
 using CopySignOpLowering =
@@ -1415,6 +1426,7 @@ void populateLLVMToMathPatterns(MLIRContext *context,
   // From
   // https://github.com/llvm/llvm-project/blob/7d8b4eb0ead277f41ff69525ed807f9f6e227f37/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp#L306
   // patterns.add<FTruncOpLowering>(converter);
+  patterns.add<AbsIOpRaising>(patterns.getContext());
   patterns.add<AbsFOpLowering,
                // AbsIOpLowering,
                CeilOpLowering, CopySignOpLowering, CosOpLowering,

test/lit_tests/raising/absi.mlir

@@ -0,0 +1,20 @@
+// RUN: enzymexlamlir-opt --libdevice-funcs-raise %s | FileCheck %s --check-prefix=RAISE
+// RUN: enzymexlamlir-opt --arith-raise %s | FileCheck %s --check-prefix=HLO
+
+module {
+  // RAISE-LABEL: @intr_absi
+  // RAISE: math.absi %arg0 : i32
+  // RAISE-NOT: llvm.intr.abs
+  func.func @intr_absi(%arg0: i32) -> i32 {
+    %res = "llvm.intr.abs"(%arg0) <{is_int_min_poison = false}> : (i32) -> i32
+    func.return %res : i32
+  }
+
+  // HLO-LABEL: @tensor_absi
+  // HLO: stablehlo.abs %arg0 : tensor<20xi32>
+  // HLO-NOT: math.absi
+  func.func @tensor_absi(%arg0: tensor<20xi32>) -> tensor<20xi32> {
+    %res = math.absi %arg0 : tensor<20xi32>
+    func.return %res : tensor<20xi32>
+  }
+}