[aievec] Lowers divf to peano intrinsic (#2851)

erwei-xilinx · web-flow · commit dda2755d9cca · 2026-02-06T03:46:29.000Z
diff --git a/include/aie/Dialect/AIEVec/IR/AIEVecOps.td b/include/aie/Dialect/AIEVec/IR/AIEVecOps.td
@@ -516,6 +516,25 @@ def AIEVec_ExpOp:
   let hasVerifier = 0;
 }
 
+def AIEVec_InvOp:
+  AIEVec_Op<"inv", [
+    Pure,
+    AllTypesMatch<["source", "result"]>
+  ]>,
+  Arguments<(ins AnyTypeOf<[F32, VectorOfLengthAndType<[16, 32], [F32]>]>:$source)>,
+  Results<(outs AnyTypeOf<[F32, VectorOfLengthAndType<[16, 32], [F32]>]>:$result)> {
+  let summary = "AIE inverse";
+  let description = [{
+    AMD-specific intrinsic that computes the inverse (1/x) of a scalar f32
+    or a vector of f32 elements.
+    For AIE2P, scalar f32 is lowered to the inv intrinsic, and vector f32
+    is unrolled into scalar inv intrinsic calls.
+    `$result = inv(`$source`) = 1.0 / $source`.
+  }];
+  let assemblyFormat = "$source attr-dict `:` type($result)";
+  let hasVerifier = 0;
+}
+
 def AIEVec_BxorOp:
   AIEVec_Op<"bxor", [
     Pure,
diff --git a/include/aie/Dialect/XLLVM/IR/XLLVMAIE2IntrOps.td b/include/aie/Dialect/XLLVM/IR/XLLVMAIE2IntrOps.td
@@ -626,6 +626,13 @@ def InvsqrtAIE2pIntrOp :
         [Pure, TypeIs<"res", F32>]>,
     Arguments<(ins F32:$src)>;
 
+// ----- INV (Reciprocal) -----
+
+def InvAIE2pIntrOp :
+    AIEVec2p_IntrOp<"inv",
+        [Pure, TypeIs<"res", F32>]>,
+    Arguments<(ins F32:$src)>;
+
 // ----- SHIFT -----
 
 def VectorShiftI512I512IntrOp :
diff --git a/lib/Conversion/AIEVecToLLVM/AIEVecToLLVM.cpp b/lib/Conversion/AIEVecToLLVM/AIEVecToLLVM.cpp
@@ -4598,6 +4598,58 @@ class ShuffleOpConversion
   }
 };
 
+// Convert aievec.inv to xllvm.intr.aie2p.inv intrinsic for AIE2P
+// Scalar f32: direct conversion to xllvm.intr.aie2p.inv
+// Vector f32: unroll into scalar xllvm.intr.aie2p.inv operations
+class InvOpAIE2pConversion
+    : public mlir::ConvertOpToLLVMPattern<aievec::InvOp> {
+public:
+  using ConvertOpToLLVMPattern<aievec::InvOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(aievec::InvOp invOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = invOp.getLoc();
+    auto operandType = adaptor.getSource().getType();
+
+    // Handle scalar f32 inverse
+    if (operandType.isF32()) {
+      auto invResult = xllvm::InvAIE2pIntrOp::create(
+          rewriter, loc, rewriter.getF32Type(), adaptor.getSource());
+      rewriter.replaceOp(invOp, invResult);
+      return success();
+    }
+
+    // Handle vector<N x f32> inverse
+    auto vecType = dyn_cast<VectorType>(operandType);
+    if (!vecType || !vecType.getElementType().isF32())
+      return failure();
+
+    // Unroll vector inverse into scalar operations
+    int numElements = getVectorLaneSize(vecType);
+    Value result = LLVM::PoisonOp::create(rewriter, loc, vecType);
+
+    for (int i = 0; i < numElements; ++i) {
+      // Extract element i
+      auto indexCst = LLVM::ConstantOp::create(
+          rewriter, loc, rewriter.getI64Type(), rewriter.getI64IntegerAttr(i));
+      auto extractedElem = LLVM::ExtractElementOp::create(
+          rewriter, loc, adaptor.getSource(), indexCst);
+
+      // Call xllvm.intr.aie2p.inv on the scalar
+      auto invResult = xllvm::InvAIE2pIntrOp::create(
+          rewriter, loc, rewriter.getF32Type(), extractedElem);
+
+      // Insert result back into vector
+      result = LLVM::InsertElementOp::create(rewriter, loc, vecType, result,
+                                             invResult, indexCst);
+    }
+
+    rewriter.replaceOp(invOp, result);
+    return success();
+  }
+};
+
 // Convert aievec.exp to xllvm.exp2 intrinsic for AIE2P
 // Uses the identity: exp(x) = exp2(x * log2(e))
 // Supports both lane-16 and lane-32 bf16 vectors
@@ -4941,6 +4993,7 @@ void populateAIEVecToLLVMAIE2pConversionPatterns(
   patterns.add<ExtractElemOpAIE2pConversion>(converter);
   patterns.add<ConcatOpAIE2pConversion>(converter);
   patterns.add<ExpOpAIE2pConversion>(converter);
+  patterns.add<InvOpAIE2pConversion>(converter);
   patterns.add<BroadcastScalarOpAIE2pConversion>(converter);
   patterns.add<RsqrtOpAIE2pConversion>(converter);
   patterns.add<FdivOpAIE2pConversion>(converter);
diff --git a/lib/Dialect/AIEVec/Transforms/VectorToAIEVecConversions.cpp b/lib/Dialect/AIEVec/Transforms/VectorToAIEVecConversions.cpp
@@ -2372,7 +2372,7 @@ struct ComputeExpOpByLUTPattern : OpConversionPattern<math::ExpOp> {
   }
 };
 
-// Lower the inverse of a float to a function call
+// Lower the inverse of a float to a function call (CPP backend)
 // Convert the pattern-
 //  %cst = arith.constant 1.000000e+00 : f32
 //  %0 = arith.divf %cst, %arg1 : f32
@@ -2426,6 +2426,77 @@ struct ComputeInvOpByLUTPattern : OpConversionPattern<arith::DivFOp> {
   }
 };
 
+// Lower the inverse of a float to aievec.inv (LLVMIR backend for AIE2P)
+// Supports both scalar f32 and vector<Nxf32> types.
+// Convert the pattern-
+//  %cst = arith.constant 1.000000e+00 : f32
+//  %0 = arith.divf %cst, %arg1 : f32
+// to -
+//  %0 = aievec.inv %arg1 : f32
+// Also supports:
+//  %cst = arith.constant dense<1.0> : vector<16xf32>
+//  %0 = arith.divf %cst, %arg1 : vector<16xf32>
+// to -
+//  %0 = aievec.inv %arg1 : vector<16xf32>
+struct ConvertDivFToAIEVecInvOpPattern : OpConversionPattern<arith::DivFOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::DivFOp divOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type srcType = adaptor.getLhs().getType();
+
+    // Check if LHS is defined by an operation
+    auto *defOp = divOp.getLhs().getDefiningOp();
+    if (!defOp)
+      return failure();
+
+    auto constOp = dyn_cast<arith::ConstantOp>(defOp);
+    if (!constOp)
+      return failure();
+
+    // Handle scalar f32 case
+    if (auto fType = dyn_cast<FloatType>(srcType)) {
+      if (fType.getWidth() != 32)
+        return failure();
+
+      auto floatAttr = dyn_cast<FloatAttr>(constOp.getValue());
+      if (!floatAttr || !floatAttr.getValue().isExactlyValue(1.0))
+        return failure();
+
+      rewriter.replaceOpWithNewOp<aievec::InvOp>(divOp, srcType,
+                                                 adaptor.getRhs());
+      return success();
+    }
+
+    // Handle vector f32 case
+    if (auto vecType = dyn_cast<VectorType>(srcType)) {
+      auto elemType = vecType.getElementType();
+      if (!elemType.isF32())
+        return failure();
+
+      // Check for supported vector sizes (16 or 32 lanes)
+      unsigned laneSize = getVectorLaneSize(vecType);
+      if (laneSize != 16 && laneSize != 32)
+        return failure();
+
+      // Check if it's a splat of 1.0
+      auto denseAttr = dyn_cast<DenseFPElementsAttr>(constOp.getValue());
+      if (!denseAttr || !denseAttr.isSplat())
+        return failure();
+
+      if (!denseAttr.getSplatValue<APFloat>().isExactlyValue(1.0))
+        return failure();
+
+      rewriter.replaceOpWithNewOp<aievec::InvOp>(divOp, vecType,
+                                                 adaptor.getRhs());
+      return success();
+    }
+
+    return failure();
+  }
+};
+
 // Convert math.tanh to a function call to compute tanh(x) by look up tables
 struct ComputeTanhOpByLUTPattern : OpConversionPattern<math::TanhOp> {
   using OpConversionPattern::OpConversionPattern;
@@ -3508,6 +3579,7 @@ populateAIEVecV2CommonConversionPatterns(RewritePatternSet &patterns,
       >(patterns.getContext(), 128, 1024, 256, 1024);
     patterns.add<
         ComputeExpOpByLUTPattern,
+        ComputeInvOpByLUTPattern,
         ComputeRsqrtOpPattern,
         LowerVectorAddFOpToAIEVecAddElemOp,
         LowerVectorSubFOpToAIEVecSubElemOp,
@@ -3521,7 +3593,6 @@ populateAIEVecV2CommonConversionPatterns(RewritePatternSet &patterns,
       >(patterns.getContext());
   }
   patterns.add<
-      ComputeInvOpByLUTPattern,
       ComputeTanhOpByLUTPattern,
       ComputeSqrtOpAIE2Pattern,
       ComputeErfOpPattern,
@@ -3629,10 +3700,11 @@ static void populateAIEVecV2PConversionPatterns(RewritePatternSet &patterns,
   // AIE2p-specific broadcast pattern that handles 256-bit directly
   patterns.add<ConvertSplatToAIEBroadcastAIE2p>(patterns.getContext());
   patterns.add<LowerVectorReductionAddBfloat16OpAIE2P>(patterns.getContext());
-  // For AIE2P with LLVMIR backend, use aievec.exp
+  // For AIE2P with LLVMIR backend, use aievec.exp and aievec.inv
   // math.rsqrt is kept legal and will be lowered in AIEVecToLLVM pass
   if (backend == TargetBackend::LLVMIR) {
-    patterns.add<ConvertMathExpToAIEVecExpOpPattern>(patterns.getContext());
+    patterns.add<ConvertMathExpToAIEVecExpOpPattern,
+                 ConvertDivFToAIEVecInvOpPattern>(patterns.getContext());
   }
 }
 
@@ -4066,6 +4138,41 @@ static void configureAIEVecV2PLegalizations(ConversionTarget &target,
 
       return false;
     });
+
+    // AIE2P-specific legalization for divf 1.0/x pattern with LLVMIR backend
+    // Scalar f32 or vector<Nxf32> divf with constant 1.0 LHS is illegal
+    target.addDynamicallyLegalOp<arith::DivFOp>([](arith::DivFOp divfOp) {
+      Type srcType = divfOp.getLhs().getType();
+
+      // Check if LHS is defined by a constant operation
+      auto constOp =
+          dyn_cast_or_null<arith::ConstantOp>(divfOp.getLhs().getDefiningOp());
+      if (!constOp)
+        return true;
+
+      // Scalar f32 case - check for exactly 1.0
+      if (srcType.isF32()) {
+        auto floatAttr = dyn_cast<FloatAttr>(constOp.getValue());
+        if (floatAttr && floatAttr.getValue().isExactlyValue(1.0))
+          return false; // illegal - will be converted to aievec.inv
+        return true;
+      }
+
+      // Vector f32 case - check for splat of exactly 1.0
+      if (auto vecType = dyn_cast<VectorType>(srcType)) {
+        if (vecType.getElementType().isF32()) {
+          unsigned laneSize = getVectorLaneSize(vecType);
+          if (laneSize == 16 || laneSize == 32) {
+            auto denseAttr = dyn_cast<DenseFPElementsAttr>(constOp.getValue());
+            if (denseAttr && denseAttr.isSplat() &&
+                denseAttr.getSplatValue<APFloat>().isExactlyValue(1.0))
+              return false; // illegal - will be converted to aievec.inv
+          }
+        }
+      }
+
+      return true;
+    });
   }
   // For CPP backend, exp remains legal (uses LUT pattern from common patterns)
 
diff --git a/test/Conversion/AIEVecToLLVM/test-inv-aie2p.mlir b/test/Conversion/AIEVecToLLVM/test-inv-aie2p.mlir
@@ -0,0 +1,36 @@
+//===- test-inv-aie2p.mlir -----------------------------------------*- MLIR -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// Copyright (C) 2026, Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: aie-opt %s -split-input-file --convert-aievec-to-llvm="aie-target=aie2p" | FileCheck %s
+
+// Test: aievec.inv → xllvm.intr.aie2p.inv for AIE2P
+
+// CHECK-LABEL: @test_inv_f32
+// CHECK-SAME: %[[ARG:.*]]: f32
+func.func @test_inv_f32(%arg0 : f32) -> f32 {
+  // CHECK: %[[INV:.*]] = "xllvm.intr.aie2p.inv"(%[[ARG]]) : (f32) -> f32
+  %0 = aievec.inv %arg0 : f32
+  // CHECK: return %[[INV]] : f32
+  return %0 : f32
+}
+
+// -----
+
+// CHECK-LABEL: @test_inv_v16f32
+// CHECK-SAME: %[[ARG:.*]]: vector<16xf32>
+func.func @test_inv_v16f32(%arg0 : vector<16xf32>) -> vector<16xf32> {
+  // CHECK: %[[POISON:.*]] = llvm.mlir.poison : vector<16xf32>
+  // CHECK: %[[IDX0:.*]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[ELEM0:.*]] = llvm.extractelement %[[ARG]][%[[IDX0]] : i64] : vector<16xf32>
+  // CHECK: %[[INV0:.*]] = "xllvm.intr.aie2p.inv"(%[[ELEM0]]) : (f32) -> f32
+  // CHECK: %[[RES0:.*]] = llvm.insertelement %[[INV0]], {{.*}}[%[[IDX0]] : i64] : vector<16xf32>
+  %0 = aievec.inv %arg0 : vector<16xf32>
+  return %0 : vector<16xf32>
+}
diff --git a/test/Conversion/VectorToAIEVec/test-inv-aie2p.mlir b/test/Conversion/VectorToAIEVec/test-inv-aie2p.mlir
