[aievec] Add scalar compound shift+clamp+truncate to SRS pattern (#2894)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: erwei-xilinx

lib/Conversion/AIEVecToLLVM/AIEVecToLLVM.cpp

@@ -4076,7 +4076,18 @@ class ExtractElemOpConversion
     // create truncation op (and bitcast op)
     if (llvm::isa<IntegerType>(resultType)) {
       if (resultBitWidth < 32) {
-        rewriter.replaceOpWithNewOp<LLVM::TruncOp>(op, resultType, extElemOp);
+        // Two-step truncation to avoid direct i32→i8 which the AIE2
+        // backend cannot legalize after SLP vectorization.
+        if (resultBitWidth < 16) {
+          auto i16Ty = rewriter.getI16Type();
+          auto trunc16 =
+              LLVM::TruncOp::create(rewriter, loc, i16Ty, extElemOp);
+          rewriter.replaceOpWithNewOp<LLVM::TruncOp>(op, resultType,
+                                                     trunc16.getResult());
+        } else {
+          rewriter.replaceOpWithNewOp<LLVM::TruncOp>(op, resultType,
+                                                     extElemOp);
+        }
       } else {
         rewriter.replaceOp(op, extElemOp);
       }
@@ -4207,34 +4218,48 @@ class MatMulOpConversion
                   /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0,
                   /*sub_mask=*/0)};
 
+    // Helper: look through vector.shape_cast ops to find the defining op.
+    // The VectorToAIEVec pass inserts shape_casts (via reshapeLeadingUnitDims)
+    // between the extension ops and the matmul, which hides the signedness.
+    auto lookThroughShapeCasts = [](Value v) -> Value {
+      while (auto castOp = v.getDefiningOp<vector::ShapeCastOp>())
+        v = castOp.getSource();
+      return v;
+    };
+
     int signX = 0, signY = 0;
     auto lhsVecTy = cast<VectorType>(lhs.getType());
     auto lhsScaTy = cast<IntegerType>(lhsVecTy.getElementType());
-    if (auto extSIOp = lhs.getDefiningOp<arith::ExtSIOp>()) {
-      lhs = extSIOp.getIn();
+    Value lhsOrig = lookThroughShapeCasts(lhs);
+    if (auto extSIOp = lhsOrig.getDefiningOp<arith::ExtSIOp>()) {
+      lhs = lookThroughShapeCasts(extSIOp.getIn());
       lhsVecTy = cast<VectorType>(lhs.getType());
       lhsScaTy = cast<IntegerType>(lhsVecTy.getElementType());
       signX = 1;
-    } else if (auto extUIOp = lhs.getDefiningOp<arith::ExtUIOp>()) {
-      lhs = extUIOp.getIn();
+    } else if (auto extUIOp = lhsOrig.getDefiningOp<arith::ExtUIOp>()) {
+      lhs = lookThroughShapeCasts(extUIOp.getIn());
       lhsVecTy = cast<VectorType>(lhs.getType());
       lhsScaTy = cast<IntegerType>(lhsVecTy.getElementType());
     } else {
-      // NOTE: We're choosing 'signed' by default
-      if (!lhsScaTy.isUnsigned())
-        signX = 1;
+      // Default to unsigned for lhs (activation input is typically uint8).
+      // The VectorToAIEVec pass strips extsi/extui before creating
+      // aievec.matmul, so sign info is not available here. Using unsigned
+      // for A matches the common use case of uint8 activations × int8 weights.
+      if (lhsScaTy.isUnsigned())
+        signX = 0;
     }
     auto lhsShape = lhsVecTy.getShape();
 
     auto rhsVecTy = cast<VectorType>(rhs.getType());
     auto rhsScaTy = cast<IntegerType>(rhsVecTy.getElementType());
-    if (auto extSIOp = rhs.getDefiningOp<arith::ExtSIOp>()) {
-      rhs = extSIOp.getIn();
+    Value rhsOrig = lookThroughShapeCasts(rhs);
+    if (auto extSIOp = rhsOrig.getDefiningOp<arith::ExtSIOp>()) {
+      rhs = lookThroughShapeCasts(extSIOp.getIn());
       rhsVecTy = cast<VectorType>(rhs.getType());
       rhsScaTy = cast<IntegerType>(rhsVecTy.getElementType());
       signY = 1;
-    } else if (auto extUIOp = rhs.getDefiningOp<arith::ExtUIOp>()) {
-      rhs = extUIOp.getIn();
+    } else if (auto extUIOp = rhsOrig.getDefiningOp<arith::ExtUIOp>()) {
+      rhs = lookThroughShapeCasts(extUIOp.getIn());
       rhsVecTy = cast<VectorType>(rhs.getType());
       rhsScaTy = cast<IntegerType>(rhsVecTy.getElementType());
     } else {

lib/Dialect/AIE/Transforms/AIECoreToStandard.cpp

@@ -109,6 +109,9 @@ static auto getAIE2Intrinsics(OpBuilder &builder) {
       {"llvm.aie2.release",
        {int32Type, int32Type},
        {}}, //(%lock_id, %lock_val) -> ()
+      {"llvm.aie2.set.ctrl.reg",
+       {int32Type, int32Type},
+       {}}, //(%reg_id, %value) -> ()
   };
   return functions;
 }
@@ -137,6 +140,9 @@ static auto getAIE2pIntrinsics(OpBuilder &builder) {
       {"llvm.aie2p.release",
        {int32Type, int32Type},
        {}}, //(%lock_id, %lock_val) -> ()
+      {"llvm.aie2p.set.ctrl.reg",
+       {int32Type, int32Type},
+       {}}, //(%reg_id, %value) -> ()
   };
   return functions;
 }
@@ -542,6 +548,50 @@ struct AIECoreToStandardFunc : OpConversionPattern<CoreOp> {
     rewriter.cloneRegionBefore(op.getBody(), coreFunc.getBody(),
                                coreFunc.getBody().begin(), mapper);
 
+    // Set saturation and rounding modes at core entry for AIE2/AIE2p, but
+    // only if the core body contains aievec.srs ops (narrowing SRS needs
+    // saturation mode enabled). Skip for cores with only lock/stream ops
+    // to avoid breaking existing test SSA naming.
+    bool hasSRS = false;
+    coreFunc.walk([&](Operation *childOp) {
+      if (childOp->getName().getStringRef() == "aievec.srs")
+        hasSRS = true;
+    });
+    if (hasSRS) {
+      auto device = op->getParentOfType<DeviceOp>();
+      if (device) {
+        AIEArch arch = device.getTargetModel().getTargetArch();
+        if (arch == AIEArch::AIE2 || arch == AIEArch::AIE2p) {
+          std::string ctrlRegFuncName = (arch == AIEArch::AIE2p)
+                                            ? "llvm.aie2p.set.ctrl.reg"
+                                            : "llvm.aie2.set.ctrl.reg";
+          auto ctrlRegFunc = module.lookupSymbol<func::FuncOp>(ctrlRegFuncName);
+          if (ctrlRegFunc) {
+            Block &entryBlock = coreFunc.getBody().front();
+            rewriter.setInsertionPointToStart(&entryBlock);
+            Location loc = op.getLoc();
+            // saturation_mode::saturate (register 9 = 1)
+            auto c9 = arith::ConstantOp::create(rewriter, loc,
+                                                rewriter.getI32IntegerAttr(9));
+            auto c1 = arith::ConstantOp::create(rewriter, loc,
+                                                rewriter.getI32IntegerAttr(1));
+            func::CallOp::create(rewriter, loc, ctrlRegFunc,
+                                 ValueRange{c9, c1});
+            // rounding_mode::floor (register 6 = 0)
+            // Use floor (truncation) to avoid double-rounding when user
+            // code already performs explicit rounding via arith.addi
+            // before shrsi.
+            auto c6 = arith::ConstantOp::create(rewriter, loc,
+                                                rewriter.getI32IntegerAttr(6));
+            auto c0 = arith::ConstantOp::create(rewriter, loc,
+                                                rewriter.getI32IntegerAttr(0));
+            func::CallOp::create(rewriter, loc, ctrlRegFunc,
+                                 ValueRange{c6, c0});
+          }
+        }
+      }
+    }
+
     // Rewrite the AIE.end() op
     coreFunc.getBody().walk([&](Operation *childOp) {
       rewriter.setInsertionPointAfter(childOp);