diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp index 186e0b7eff..2450f568a4 100644 --- a/lib/PTO/Transforms/VMILayoutAssignment.cpp +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -664,6 +664,13 @@ struct LayoutSolver { failed(unite(fma.getLhs(), fma.getResult(), op)); return constraintResult(failure(failedToUnite)); }) + .Case([this, op](VMIVmulaOp vmula) { + bool failedToUnite = + failed(uniteDataEquivalent(vmula.getLhs(), vmula.getRhs(), op)) || + failed(uniteDataEquivalent(vmula.getLhs(), vmula.getAcc(), op)) || + failed(uniteDataEquivalent(vmula.getLhs(), vmula.getResult(), op)); + return constraintResult(failure(failedToUnite)); + }) .Case([this, op](auto compareOp) { return constraintResult( unite(compareOp.getLhs(), compareOp.getRhs(), op)); @@ -1718,6 +1725,18 @@ struct LayoutSolver { if (failed(requestDataUseSeeds(propagator, phase, /*late=*/false))) { return failure(); } + if (phase == DataLayoutSeedPhase::Store) { + if (failed(propagator.run())) + return failure(); + module.walk([&](VMIVmulaOp vmula) { + if (vmula.getMask().empty()) + return; + VMILayoutAttr layout = + propagator.getRequestedOrCurrentLayout(vmula.getLhs()); + if (layout) + (void)propagator.request(*vmula.getMaskMutable().begin(), layout); + }); + } if (failed(requestMaskUseSeeds(propagator, phase))) { return failure(); } diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index 3c31ab3b5d..67599324bd 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -943,7 +943,9 @@ static LogicalResult lowerVcmin(VMIvcminOp op, OpBuilder &builder) { /// Legacy fma is floating-point only; integer vmula has no legacy equivalent /// and is skipped (falls through to VMIToVPTO). static LogicalResult lowerVmula(VMIVmulaOp op, OpBuilder &builder) { - if (hasMergePmode(op)) { + // Legacy FMA has no predicate operand. Keep masked vmula in the VMI + // pipeline so VMIToVPTO can preserve its predicate and zero-mode semantics. + if (hasMergePmode(op) || !op.getMask().empty()) { return failure(); } diff --git a/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp index 6c66aa87e0..e70120bc40 100644 --- a/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp +++ b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp @@ -297,6 +297,15 @@ struct MaskGranularitySolver { return constraintResult( requestMaskUse(histogram.getMaskMutable(), "b8", op)); }) + .Case([this, op](VMIVmulaOp vmula) { + if (vmula.getMask().empty()) { + return constraintResult(success()); + } + auto accType = cast(vmula.getAcc().getType()); + return constraintResult(requestMaskUse( + *vmula.getMaskMutable().begin(), + getMaskGranularityForElement(accType.getElementType()), op)); + }) .Default([](Operation *) { return std::nullopt; }); } diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 75873dcb7a..be94309292 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -360,11 +360,10 @@ static FailureOr getVMIMaskPhysicalGranularity(VMIMaskType type) { if (bits == 0) return failure(); - VMILayoutAttr layout = type.getLayoutAttr(); - int64_t laneStride = layout && layout.hasLaneStride() ? layout.getLaneStride() - : 1; - int64_t physicalBits = bits * laneStride; - StringRef physicalGranularity = getMaskGranularityForBits(physicalBits); + // VPTO masks are typed by the data element width. VMI layouts may carry a + // lane stride for packed data, but that stride does not widen the predicate + // consumed by a vector instruction (e.g. i16 data still requires mask). + StringRef physicalGranularity = getMaskGranularityForBits(bits); if (physicalGranularity.empty()) return failure(); return physicalGranularity; @@ -5992,6 +5991,30 @@ struct OneToNVMICreateMaskOpPattern SmallVector results; results.reserve(resultTypes.size()); + // A contiguous packed VMI vector maps to one regular VPTO mask chunk. + // Materialize the prefix in physical-lane space for any positive stride; + // do not special-case a particular packing factor (e.g. lane_stride=2). + if (layout.isContiguous() && layout.hasLaneStride() && + layout.getLaneStride() > 0 && factor == 1 && resultTypes.size() == 1) { + auto maskType = dyn_cast(resultTypes.front()); + if (!maskType) + return rewriter.notifyMatchFailure(op, "create_mask result must be mask"); + // Each logical lane occupies lane_stride physical lanes in a contiguous + // packed layout, so predicate coverage scales with that stride. + int64_t activeInChunk = std::min( + activeLanes * layout.getLaneStride(), *lanesPerPart); + std::optional pattern = + getPrefixPattern(activeInChunk, *lanesPerPart); + if (!pattern) + return rewriter.notifyMatchFailure(op, "unsupported packed create_mask prefix"); + FailureOr mask = createPrefixMask(op.getLoc(), maskType, *pattern, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure(op, "failed to materialize packed create_mask"); + results.push_back(*mask); + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + for (int64_t part = 0; part < factor; ++part) { for (int64_t chunk = 0;; ++chunk) { bool anyLane = false; @@ -9798,6 +9821,67 @@ struct OneToNVMIFmaOpPattern : OneToNOpConversionPattern { } }; +struct OneToNVMIVmulaOpPattern : OneToNOpConversionPattern { + using OneToNOpConversionPattern::OneToNOpConversionPattern; + LogicalResult matchAndRewrite(VMIVmulaOp op, OpAdaptor adaptor, + OneToNPatternRewriter &rewriter) const override { + ValueRange acc = adaptor.getAcc(); + ValueRange lhs = adaptor.getLhs(); + ValueRange rhs = adaptor.getRhs(); + ArrayRef maskParts = adaptor.getMask(); + FailureOr> converted = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(converted) || acc.size() != lhs.size() || + acc.size() != rhs.size() || acc.size() != converted->size()) + return rewriter.notifyMatchFailure(op, "vmula physical arity mismatch"); + if (maskParts.size() > 1 || + (!maskParts.empty() && maskParts.front().size() != acc.size())) + return rewriter.notifyMatchFailure(op, "vmula mask physical arity mismatch"); + SmallVector results; + for (unsigned i = 0; i < acc.size(); ++i) { + auto type = dyn_cast((*converted)[i]); + if (!type || acc[i].getType() != (*converted)[i] || + lhs[i].getType() != (*converted)[i] || + rhs[i].getType() != (*converted)[i]) + return rewriter.notifyMatchFailure(op, "vmula requires matching physical vreg parts"); + Value mask; + if (maskParts.empty()) { + FailureOr allTrue = createAllTrueMaskForVReg(op.getLoc(), type, rewriter); + if (failed(allTrue)) + return rewriter.notifyMatchFailure(op, + "unsupported element type for vmula"); + mask = *allTrue; + } else { + mask = maskParts.front()[i]; + } + auto vregType = dyn_cast((*converted)[i]); + if (!vregType) + return rewriter.notifyMatchFailure(op, + "vmula requires a vector result type"); + Value result = rewriter + .create(op.getLoc(), (*converted)[i], acc[i], + lhs[i], rhs[i], mask) + .getResult(); + // VPTO vmula preserves the accumulator on inactive lanes. VMI's + // explicit zero mode instead requires inactive lanes to be cleared, so + // materialize that semantic difference after the fused operation. + if (op.getPmode().has_value() && *op.getPmode() == "zero") { + FailureOr zero = createZeroVector(op.getLoc(), vregType, rewriter); + if (failed(zero)) + return rewriter.notifyMatchFailure( + op, "failed to materialize vmula zero-mode value"); + result = rewriter + .create(op.getLoc(), (*converted)[i], result, *zero, + mask) + .getResult(); + } + results.push_back(result); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + struct OneToNVMIVexpdifOpPattern : OneToNOpConversionPattern { using OneToNOpConversionPattern::OneToNOpConversionPattern; @@ -13256,6 +13340,7 @@ void populateVMIConversionPatterns( OneToNVMIVecScalarOpPattern, OneToNVMIVecScalarOpPattern, OneToNVMIVecScalarOpPattern, OneToNVMIVmullOpPattern, + OneToNVMIVmulaOpPattern, OneToNVMIFmaOpPattern, OneToNVMIVexpdifOpPattern, OneToNVMIBinaryOpPattern, OneToNVMIBinaryOpPattern,