Fix TCmp (partially)

Author: Zhendong404

include/PTO/IR/PTOOps.td

@@ -1438,6 +1438,35 @@ def SimdReductionOp : PTO_Op<"simd.reduction", [
   }];
 }
 
+def SimdStorePredicateOp : PTO_Op<"simd.store_predicate"> {
+  let summary = "A5 OP-Lib marker for packed predicate compare+store.";
+  let description = [{
+    Backend-only OP-Lib bridge op that compares a SIMD vector against either
+    another SIMD vector or a scalar under an active-lane count and writes the
+    packed predicate result to the destination tile/memref using the hardware
+    packed-mask store path.
+  }];
+
+  let arguments = (ins
+    AnyVector:$lhs,
+    AnyType:$rhs,
+    TileBufOrMemRef:$dst,
+    Index:$linear_offset,
+    Index:$active_count,
+    PTO_CmpModeAttr:$cmpMode
+  );
+
+  let results = (outs);
+
+  let hasVerifier = 1;
+
+  let assemblyFormat = [{
+    $lhs `,` $rhs `,` $dst `,` $linear_offset `,` $active_count attr-dict
+      `:` type($lhs) `,` type($rhs) `,` qualified(type($dst)) `,`
+         type($linear_offset) `,` type($active_count)
+  }];
+}
+
 // High-Level Mov: 值语义
 def MovOp : PTO_Op<"mov", [
     SameOperandsAndResultType, // 输入输出类型相同 (!pto.tile -> !pto.tile)

lib/PTO/IR/PTO.cpp

@@ -3829,6 +3829,43 @@ mlir::LogicalResult mlir::pto::SimdReductionOp::verify() {
 
   return success();
 }
+
+static LogicalResult verifySimdPackedPredicateDst(Operation *op, Type dstTy) {
+  Type elemTy;
+  if (auto tileTy = dyn_cast<TileBufType>(dstTy)) {
+    elemTy = tileTy.getElementType();
+  } else if (auto memTy = dyn_cast<MemRefType>(dstTy)) {
+    elemTy = memTy.getElementType();
+  } else {
+    return op->emitOpError("expects dst to be !pto.tile_buf or memref");
+  }
+
+  auto intTy = dyn_cast<IntegerType>(elemTy);
+  if (!intTy || intTy.getWidth() != 8)
+    return op->emitOpError("expects dst element type to be i8/ui8");
+  return success();
+}
+
+mlir::LogicalResult mlir::pto::SimdStorePredicateOp::verify() {
+  auto lhsTy = dyn_cast<VectorType>(getLhs().getType());
+  if (!lhsTy || lhsTy.isScalable())
+    return emitOpError("expects lhs to be a fixed-width vector type");
+  if (!lhsTy.getElementType().isIntOrFloat())
+    return emitOpError("expects lhs element type to be integer or float");
+
+  Type rhsTy = getRhs().getType();
+  if (auto rhsVecTy = dyn_cast<VectorType>(rhsTy)) {
+    if (rhsVecTy.isScalable())
+      return emitOpError("expects rhs vector to be fixed-width");
+    if (rhsVecTy != lhsTy)
+      return emitOpError("expects rhs vector type to match lhs");
+  } else if (rhsTy != lhsTy.getElementType()) {
+    return emitOpError(
+        "expects rhs to be either a matching vector type or lhs element type");
+  }
+
+  return verifySimdPackedPredicateDst(*this, getDst().getType());
+}
 //===----------------------------------------------------------------------===//
 // PTO.cpp  (add TSYNC DPS/tilebuf implementation)
 //===----------------------------------------------------------------------===//

lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp

@@ -1000,7 +1000,12 @@ static FailureOr<int64_t> getFixedVectorLanes(Type ty) {
 
 static bool isSimdBridgeOp(Operation *op) {
   return isa<pto::SimdPredicateOp, pto::SimdLoadOp, pto::SimdStoreOp,
-             pto::SimdLoadPUOp, pto::SimdStorePUOp>(op);
+             pto::SimdLoadPUOp, pto::SimdStorePUOp,
+             pto::SimdStorePredicateOp>(op);
+}
+
+static bool isPackedCmpSimdBridgeOp(Operation *op) {
+  return isa<pto::SimdStorePredicateOp>(op);
 }
 
 static bool isAllowedTemplateBodyOp(Operation *op) {
@@ -1462,6 +1467,22 @@ struct TemplateRegistry {
     std::string inferredVstDist;
     std::string inferredExecMode;
     bool requiresUnifiedLevel3Simd = isLevel3TemplateKind(entry.kind);
+    bool hasPackedCmpBridgeOnly = false;
+    if (entry.hasSimdBridgeOps) {
+      bool sawPackedCmp = false;
+      bool sawOtherBridge = false;
+      imported.walk([&](Operation *op) {
+        if (!isSimdBridgeOp(op))
+          return WalkResult::advance();
+        if (isPackedCmpSimdBridgeOp(op)) {
+          sawPackedCmp = true;
+          return WalkResult::advance();
+        }
+        sawOtherBridge = true;
+        return WalkResult::interrupt();
+      });
+      hasPackedCmpBridgeOnly = sawPackedCmp && !sawOtherBridge;
+    }
 
     llvm::DenseMap<Operation *, int64_t> preorder;
     int64_t seq = 0;
@@ -1738,23 +1759,26 @@ struct TemplateRegistry {
       return emitFailureWithCode(imported.getLoc(), kErrCoreSlot,
                                  "seed template must contain exactly one core slot op");
     }
-    if (entry.hasSimdBridgeOps && coreCount != 1) {
+    if (entry.hasSimdBridgeOps && !hasPackedCmpBridgeOnly && coreCount != 1) {
       return emitFailureWithCode(imported.getLoc(), kErrCoreSlot,
                                  "template using pto.simd.* must contain exactly one core slot op");
     }
     if (coreCount > 1) {
       return emitFailureWithCode(imported.getLoc(), kErrCoreSlot,
                                  "template must not contain multiple core slot ops");
     }
-    if (entry.hasSimdBridgeOps && firstLoad == std::numeric_limits<int64_t>::max()) {
+    if (entry.hasSimdBridgeOps && !hasPackedCmpBridgeOnly &&
+        firstLoad == std::numeric_limits<int64_t>::max()) {
       return emitFailureWithCode(imported.getLoc(), kErrCoreSlot,
                                  "template must contain simd.load/simd.load_pu");
     }
-    if (entry.hasSimdBridgeOps && firstStore == std::numeric_limits<int64_t>::max()) {
+    if (entry.hasSimdBridgeOps && !hasPackedCmpBridgeOnly &&
+        firstStore == std::numeric_limits<int64_t>::max()) {
       return emitFailureWithCode(imported.getLoc(), kErrCoreSlot,
                                  "template must contain simd.store/simd.store_pu");
     }
-    if (entry.hasSimdBridgeOps && !(firstLoad < coreSeq && coreSeq < firstStore)) {
+    if (entry.hasSimdBridgeOps && !hasPackedCmpBridgeOnly &&
+        !(firstLoad < coreSeq && coreSeq < firstStore)) {
       return emitFailureWithCode(imported.getLoc(), kErrCoreSlot,
                                  "template ordering must satisfy load -> core -> store");
     }

lib/PTO/Transforms/PTOToEmitC.cpp

@@ -2441,6 +2441,21 @@ buildA5PredicateFromActiveCount(ConversionPatternRewriter &rewriter, Location lo
   return pred.getResult(0);
 }
 
+static FailureOr<Value> buildA5PredicateFromActiveCountAndElemToken(
+    ConversionPatternRewriter &rewriter, Location loc,
+    llvm::StringRef elemTok, Value activeCount) {
+  auto *ctx = rewriter.getContext();
+  Value activeCountLValue =
+      materializeA5PredicateScalarLValue(rewriter, loc, activeCount);
+  auto maskRegTy = emitc::OpaqueType::get(ctx, "MaskReg");
+  auto templateArgs =
+      rewriter.getArrayAttr({emitc::OpaqueAttr::get(ctx, elemTok)});
+  auto pred = rewriter.create<emitc::CallOpaqueOp>(
+      loc, TypeRange{maskRegTy}, "CreatePredicate",
+      ValueRange{activeCountLValue}, ArrayAttr{}, templateArgs);
+  return pred.getResult(0);
+}
+
 template <typename ArithOp> static llvm::StringRef getA5VectorBinaryCallee();
 
 template <> llvm::StringRef getA5VectorBinaryCallee<arith::AddFOp>() {
@@ -2590,6 +2605,8 @@ getA5ReductionCallee(vector::CombiningKind kind, Type elemTy) {
   }
 }
 
+static std::string cmpModeTok(pto::CmpModeAttr a);
+
 static FailureOr<llvm::StringRef> getA5SimdReductionCallee(StringRef kind) {
   if (kind == "add")
     return llvm::StringRef("vcadd");
@@ -2942,6 +2959,81 @@ struct SimdReductionToEmitC : public OpConversionPattern<pto::SimdReductionOp> {
   }
 };
 
+struct SimdStorePredicateToEmitC
+    : public OpConversionPattern<pto::SimdStorePredicateOp> {
+  using OpConversionPattern<pto::SimdStorePredicateOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(pto::SimdStorePredicateOp op, OpAdaptor adaptor,
+                                ConversionPatternRewriter &rewriter) const override {
+    auto vecTy = dyn_cast<VectorType>(op.getLhs().getType());
+    if (!vecTy)
+      return failure();
+    if (failed(
+            validateA5OplibVectorType(op, vecTy, op->getName().getStringRef())))
+      return failure();
+
+    auto elemTokOr =
+        getA5VectorElemToken(op, vecTy, op->getName().getStringRef());
+    if (failed(elemTokOr))
+      return failure();
+
+    auto predOr = buildA5PredicateFromActiveCountAndElemToken(
+        rewriter, op.getLoc(), *elemTokOr, adaptor.getActiveCount());
+    if (failed(predOr))
+      return failure();
+
+    auto *ctx = rewriter.getContext();
+    auto maskTy = emitc::OpaqueType::get(ctx, "MaskReg");
+    auto maskVar = rewriter.create<emitc::VariableOp>(
+        op.getLoc(), maskTy, emitc::OpaqueAttr::get(ctx, ""));
+
+    auto modeTy = emitc::OpaqueType::get(ctx, "CmpMode");
+    auto modeVal = rewriter.create<emitc::ConstantOp>(
+        op.getLoc(), modeTy,
+        emitc::OpaqueAttr::get(ctx, cmpModeTok(op.getCmpModeAttr())));
+
+    auto cmpArgs = rewriter.getArrayAttr(
+        {rewriter.getIndexAttr(0), rewriter.getIndexAttr(1),
+         rewriter.getIndexAttr(2), rewriter.getIndexAttr(3),
+         rewriter.getIndexAttr(4)});
+    Value rhs = adaptor.getRhs();
+    if (isa<VectorType>(op.getRhs().getType())) {
+      auto rhsTy = cast<VectorType>(op.getRhs().getType());
+      if (failed(validateA5OplibVectorType(op, rhsTy,
+                                           op->getName().getStringRef())))
+        return failure();
+      rewriter.create<emitc::CallOpaqueOp>(
+          op.getLoc(), TypeRange{}, "ptoas_vcmp",
+          ValueRange{maskVar.getResult(), adaptor.getLhs(), rhs,
+                     modeVal.getResult(), *predOr},
+          cmpArgs, ArrayAttr{});
+    } else {
+      rewriter.create<emitc::CallOpaqueOp>(
+          op.getLoc(), TypeRange{}, "ptoas_vcmps",
+          ValueRange{maskVar.getResult(), adaptor.getLhs(), rhs,
+                     modeVal.getResult(), *predOr},
+          cmpArgs, ArrayAttr{});
+    }
+
+    Value dst = peelUnrealized(adaptor.getDst());
+    if (!isEmitCTileOpaqueType(dst.getType()))
+      return rewriter.notifyMatchFailure(
+          op,
+          "simd.store_predicate currently requires tile-like dst in EmitC lowering");
+
+    auto i32Ty = emitc::OpaqueType::get(ctx, "int32_t");
+    Value linearIndex = emitCCast(rewriter, op.getLoc(), i32Ty,
+                                  adaptor.getLinearOffset());
+    rewriter.create<emitc::CallOpaqueOp>(
+        op.getLoc(), TypeRange{}, "ptoas_pstore",
+        ValueRange{maskVar.getResult(), dst, linearIndex}, ArrayAttr{},
+        ArrayAttr{});
+
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 struct OplibMemRefLoadToEmitC : public OpConversionPattern<memref::LoadOp> {
   using OpConversionPattern<memref::LoadOp>::OpConversionPattern;
 
@@ -9124,6 +9216,8 @@ struct SimdVecScopeToEmitC : public OpConversionPattern<pto::SimdVecScopeOp> {
   matchAndRewrite(pto::SimdVecScopeOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
+    rewriter.create<emitc::VerbatimOp>(
+        loc, "#if defined(__CCE_AICORE__) || defined(__CPU_SIM)");
     rewriter.create<emitc::VerbatimOp>(loc, "__VEC_SCOPE__ {");
 
     Block &innerBlock = op.getBody().front();
@@ -9132,6 +9226,8 @@ struct SimdVecScopeToEmitC : public OpConversionPattern<pto::SimdVecScopeOp> {
     }
 
     rewriter.create<emitc::VerbatimOp>(loc, "}");
+    rewriter.create<emitc::VerbatimOp>(
+        loc, "#endif // __CCE_AICORE__ || __CPU_SIM");
     rewriter.eraseOp(op);
     return success();
   }
@@ -9771,6 +9867,7 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns,
                OplibVectorUnaryToEmitC<math::SqrtOp>,
                OplibVectorUnaryToEmitC<math::RsqrtOp>,
                OplibVectorCmpFToEmitC, OplibVectorCmpIToEmitC,
+               SimdStorePredicateToEmitC,
                OplibVectorSelectToEmitC<arith::SelectOp>,
                SimdReductionToEmitC,
                OplibVectorReductionToEmitC,
@@ -10465,7 +10562,8 @@ struct EmitPTOManualPass
       if (callee == "ptoas_vreduce_add" || callee == "ptoas_vreduce_max" ||
           callee == "ptoas_vreduce_min")
         needsVectorReductionHelper = true;
-      if (callee == "ptoas_vcmp")
+      if (callee == "ptoas_vcmp" || callee == "ptoas_vcmps" ||
+          callee == "ptoas_pstore")
         needsVectorCmpHelper = true;
       if (callee == "ptoas_vrem")
         needsVectorRemfHelper = true;
@@ -10494,6 +10592,7 @@ struct EmitPTOManualPass
     if (needsVectorReductionHelper) {
       helperBuilder.create<emitc::VerbatimOp>(
           loc, helperBuilder.getStringAttr(R"cpp(
+    #if defined(__CCE_AICORE__) || defined(__CPU_SIM)
     template <typename T>
     PTO_INTERNAL T ptoas_vreduce_add(RegTensor<T> src, MaskReg pred) {
       RegTensor<T> dst;
@@ -10546,11 +10645,13 @@ struct EmitPTOManualPass
       T red = ptoas_vreduce_min(src, pred);
       return red < acc ? red : acc;
     }
+    #endif
     )cpp"));
     }
     if (needsVectorCmpHelper) {
       helperBuilder.create<emitc::VerbatimOp>(
           loc, helperBuilder.getStringAttr(R"cpp(
+    #if defined(__CCE_AICORE__) || defined(__CPU_SIM)
     template <typename T>
     PTO_INTERNAL void ptoas_vcmp(MaskReg &dst, RegTensor<T> src0, RegTensor<T> src1,
                                  CmpMode mode, MaskReg pred) {
@@ -10578,11 +10679,47 @@ struct EmitPTOManualPass
         break;
       }
     }
+    template <typename RegT, typename T>
+    PTO_INTERNAL void ptoas_vcmps(MaskReg &dst, RegT src0, T src1,
+                                  CmpMode mode, MaskReg pred) {
+      switch (mode) {
+      case CmpMode::EQ:
+        vcmps_eq(dst, src0, src1, pred);
+        break;
+      case CmpMode::NE:
+        vcmps_ne(dst, src0, src1, pred);
+        break;
+      case CmpMode::LT:
+        vcmps_lt(dst, src0, src1, pred);
+        break;
+      case CmpMode::LE:
+        vcmps_le(dst, src0, src1, pred);
+        break;
+      case CmpMode::GT:
+        vcmps_gt(dst, src0, src1, pred);
+        break;
+      case CmpMode::GE:
+        vcmps_ge(dst, src0, src1, pred);
+        break;
+      default:
+        vcmps_eq(dst, src0, src1, pred);
+        break;
+      }
+    }
+    template <typename TileT>
+    PTO_INTERNAL void ptoas_pstore(MaskReg src, TileT &dst, int32_t linearIndex) {
+      __ubuf__ uint32_t *dstWords =
+          reinterpret_cast<__ubuf__ uint32_t *>(dst.data());
+      int32_t wordOffset = linearIndex / 32;
+      psts(src, dstWords + wordOffset, 0, PK);
+    }
+    #endif
     )cpp"));
     }
     if (needsVectorRemfHelper) {
       helperBuilder.create<emitc::VerbatimOp>(
           loc, helperBuilder.getStringAttr(R"cpp(
+    #if defined(__CCE_AICORE__) || defined(__CPU_SIM)
     template <typename T>
     PTO_INTERNAL void ptoas_vrem(RegTensor<T> &dst, RegTensor<T> src0,
                                  RegTensor<T> src1, MaskReg pred) {
@@ -10611,6 +10748,7 @@ struct EmitPTOManualPass
         vor(dst, dstEven, dstOdd, pred);
       }
     }
+    #endif
     )cpp"));
     }
     if (needsBitcastHelper) {