feat(vpto): canonicalize integer-backed castptr for post-update addressing (#591)

include/PTO/Analysis/PTOAddressAnalysis.h

@@ -32,6 +32,13 @@ class AnalysisManager;
 
 namespace pto {
 
+/// Element storage size in bytes for an integer/float element type, or for a
+/// pointer/view-typed SSA value. Shared by address analysis and the integer
+/// address canonicalization rewrite so the element-size rule lives in one
+/// place.
+std::optional<int64_t> getPTOElementBytes(Type elementType);
+std::optional<int64_t> getPTOElementBytes(Value pointer);
+
 struct PTOTypedAddressOffset {
   Value sourceValue;
   PTOTypedExprRef value;

include/PTO/IR/VPTOOps.td

@@ -3659,7 +3659,8 @@ def PTO_VmulconvOp : PTO_VectorMicroOp<"vmulconv", [Pure]> {
 }
 
 def PTO_Vstsx2Op : PTO_VectorMicroOp<"vstsx2", [
-    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
+    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+    VPTOAddressSemanticsOpInterface
   ]> {
   let arguments = (ins
     PTO_VectorType:$low,

include/PTO/Transforms/Passes.h

@@ -102,6 +102,8 @@ std::unique_ptr<Pass> createPTOInferVPTOVecScopePass();
 std::unique_ptr<Pass> createVPTOExpandWrapperOpsPass();
 std::unique_ptr<Pass> createVPTOSoftPostUpdatePass();
 std::unique_ptr<Pass> createPTOPrintAddressAnalysisPass();
+std::unique_ptr<Pass> createPTOIntegerAddressCanonicalizationPass();
+std::unique_ptr<Pass> createPTOAbsorbAddPtrPass();
 std::unique_ptr<Pass> createPTOVPTOPtrBoundaryPass();
 std::unique_ptr<Pass>
 createPTOLowLevelLoopFusionPass(const PTOLowLevelLoopFusionOptions &options = {});

include/PTO/Transforms/Passes.td

@@ -1141,17 +1141,17 @@ def VMILowerUnifiedToLegacy : Pass<"vmi-lower-unified-to-legacy", "ModuleOp"> {
     Ops lowered (Category A–C6):
       A:  vci, vinterpret_cast, vsel, vbrc  →  iota, bitcast, select, broadcast/group_broadcast
       B:  vadd/vsub/vmul/vdiv/vmin/vmax/vand/vor/vxor/vshl/vshr (masked binary)
-          vneg/vabs/vsqrt/vexp/vln/vrelu/vnot (masked unary)
-          →  binary ops discard mask/pmode; unary ops preserve zero mode with select;
+          vneg/vsqrt/vexp/vln/vrelu/vnot (masked unary)
+          →  mask-less legacy operations;
              vshr selects shrui for explicit unsigned elements and shrsi otherwise
       C1: vcmp/vcmps  →  legacy cmp + select
       C2: vcvt        →  legacy extf/truncf/fptosi/sitofp/extsi/extui/trunci
       C3: vload/vstore  →  legacy load/store variants
       C4: pset/pge    →  create_mask/create_group_mask
       C6: vcadd/vcmax/vcmin  →  legacy reduce variants
 
-    Ops NOT lowered (no legacy equivalent — require direct VMIToVPTO 1:N patterns):
-      plt, vadds, vmuls, vmaxs, vmins, vshls, vshrs, vhist, vintlv, vdintlv, vselr,
+    Ops kept unified for direct VMIToVPTO 1:N patterns:
+      plt, vabs, vadds, vmuls, vmaxs, vmins, vshls, vshrs, vhist, vintlv, vdintlv, vselr,
       vgather, vgatherb, vscatter,
       vexpdif, vaxpy, vlrelu, vprelu, vmull, vmula
   }];
@@ -1300,6 +1300,44 @@ def PTOPrintAddressAnalysis
                            "mlir::scf::SCFDialect"];
 }
 
+def PTOIntegerAddressCanonicalization
+    : Pass<"pto-canonicalize-integer-address", "func::FuncOp"> {
+  let summary = "Canonicalize integer-backed pto.castptr into castptr(root) + addptr(quotient)";
+  let description = [{
+    Rewrites `castptr(B)` with an integer byte-address input `B` into the
+    canonical pointer form `addptr(castptr(R), Q)` where `R` is the canonical
+    root (constant 0, or a single non-divisible atom leaf such as a runtime
+    base-address kernel parameter) and `Q` is the exact element quotient
+    `(B - R) / sizeof(element)`. The rewrite only matches zero-origin integral
+    address spaces (A5 UB). It refuses non-linear inputs, non-exact quotients,
+    non-unit-coefficient or multiple non-divisible atoms, and inputs whose
+    width does not round-trip into index. The rule requires a non-trivial
+    quotient (constant-zero and pure-atom inputs stay untouched) so the
+    rewrite converges to the normal form in one pass.
+  }];
+  let constructor = "mlir::pto::createPTOIntegerAddressCanonicalizationPass()";
+  let dependentDialects = ["mlir::func::FuncDialect",
+                           "mlir::pto::PTODialect",
+                           "mlir::arith::ArithDialect"];
+}
+
+def PTOAbsorbAddPtr
+    : Pass<"pto-absorb-addptr", "func::FuncOp"> {
+  let summary = "Absorb addptr element offsets into VPTO memory op offsets";
+  let description = [{
+    Canonical fold `op(addptr(base, A), O) -> op(base, A + O)` for VPTO memory
+    operations whose offset unit is Element and whose element type matches the
+    addptr. Legality comes only from VPTOAddressSemanticsOpInterface:
+    current access with an Element offset, no updated-base post-update form,
+    and a no-loss index addition. This is the backend-shape fold that lets the
+    post-update consumer see the affine offset directly on the operation.
+  }];
+  let constructor = "mlir::pto::createPTOAbsorbAddPtrPass()";
+  let dependentDialects = ["mlir::func::FuncDialect",
+                           "mlir::pto::PTODialect",
+                           "mlir::arith::ArithDialect"];
+}
+
 def PTOVPTOPtrBoundary
     : Pass<"pto-vpto-ptr-boundary", "ModuleOp"> {
   let summary =

lib/PTO/Analysis/PTOAddressAnalysis.cpp

@@ -21,25 +21,6 @@ namespace {
 
 static constexpr int64_t kBlockSizeBytes = 32;
 
-static std::optional<int64_t> getElementBytes(Value pointer) {
-  Type elementType;
-  if (auto pointerType = dyn_cast<PtrType>(pointer.getType())) {
-    elementType = pointerType.getElementType();
-  } else if (auto memrefType = dyn_cast<BaseMemRefType>(pointer.getType())) {
-    elementType = memrefType.getElementType();
-  } else {
-    return std::nullopt;
-  }
-  if (!elementType || !elementType.isIntOrFloat()) {
-    return std::nullopt;
-  }
-  unsigned bitWidth = elementType.getIntOrFloatBitWidth();
-  if (bitWidth == 0 || bitWidth % mlir::pto::kValue8 != 0) {
-    return std::nullopt;
-  }
-  return static_cast<int64_t>(bitWidth / mlir::pto::kValue8);
-}
-
 static std::optional<int64_t> getUnitBytes(Operation *operation,
                                            VPTOAddressUnit unit,
                                            int64_t elementBytes) {
@@ -91,6 +72,29 @@ static bool isZero(const PTOTypedExprRef &expression) {
 
 } // namespace
 
+std::optional<int64_t> mlir::pto::getPTOElementBytes(Type elementType) {
+  if (!elementType || !elementType.isIntOrFloat()) {
+    return std::nullopt;
+  }
+  unsigned bitWidth = elementType.getIntOrFloatBitWidth();
+  if (bitWidth == 0 || bitWidth % mlir::pto::kValue8 != 0) {
+    return std::nullopt;
+  }
+  return static_cast<int64_t>(bitWidth / mlir::pto::kValue8);
+}
+
+std::optional<int64_t> mlir::pto::getPTOElementBytes(Value pointer) {
+  Type elementType;
+  if (auto pointerType = dyn_cast<PtrType>(pointer.getType())) {
+    elementType = pointerType.getElementType();
+  } else if (auto memrefType = dyn_cast<BaseMemRefType>(pointer.getType())) {
+    elementType = memrefType.getElementType();
+  } else {
+    return std::nullopt;
+  }
+  return getPTOElementBytes(elementType);
+}
+
 PTOAddressAnalysis::PTOAddressAnalysis(func::FuncOp func,
                                        AnalysisManager &analysisManager)
     : func(func),
@@ -118,7 +122,7 @@ PTOAddressAnalysis::getAddresses(Operation *operation) {
   VPTOAddressSemantics contract = semantics.getVPTOAddressSemantics();
   for (const VPTOAddressAccess &access : contract.currentAccesses) {
     Value base = access.baseOperand->get();
-    auto elementBytes = getElementBytes(base);
+    auto elementBytes = getPTOElementBytes(base);
     if (!elementBytes) {
       return PTOAnalysisResult<SmallVector<PTOAddressExpr>>::unknown(
           PTOAnalysisUnknownReason::UnknownElementSize);
@@ -132,7 +136,7 @@ PTOAddressAnalysis::getAddresses(Operation *operation) {
 
     while (auto addPointer =
                address.rootOrBase.getDefiningOp<AddPtrOp>()) {
-      auto parentElementBytes = getElementBytes(addPointer.getPtr());
+      auto parentElementBytes = getPTOElementBytes(addPointer.getPtr());
       if (!parentElementBytes || *parentElementBytes != *elementBytes) {
         break;
       }

lib/PTO/IR/VPTOAddressSemantics.cpp

@@ -95,6 +95,15 @@ mlir::pto::getDefaultVPTOAddressSemantics(Operation *operation) {
             postUpdate(base, &offset, VPTOAddressUnit::Element,
                        op.getUpdatedBase())};
       })
+      .Case<Vstsx2Op>([](Vstsx2Op op) {
+        // vstsx2 has no result, so it cannot carry a post-update updated base;
+        // only the current access (destination + element offset) is modeled.
+        OpOperand &base = op.getDestinationMutable();
+        OpOperand &offset = op.getOffsetMutable();
+        return VPTOAddressSemantics{
+            {oneAccess(base, offset, VPTOAddressUnit::Element)},
+            std::nullopt};
+      })
       .Case<VstusOp>([](VstusOp op) {
         OpOperand &base = op.getBaseMutable();
         return VPTOAddressSemantics{

lib/PTO/Transforms/CMakeLists.txt

@@ -48,6 +48,8 @@ add_mlir_dialect_library(PTOTransforms
   VPTOMaskSimplify.cpp
   VPTOExpandWrapperOps.cpp
   VPTOSoftPostUpdate.cpp
+  PTOIntegerAddressCanonicalization.cpp
+  PTOAbsorbAddPtr.cpp
   PTOPrintAddressAnalysis.cpp
   VPTOScheduler/VPTORegPressureTracker.cpp
   VPTOScheduler/VPTOSchedBoundary.cpp

lib/PTO/Transforms/PTOAbsorbAddPtr.cpp

@@ -0,0 +1,123 @@
+// Copyright (c) 2026 Huawei Technologies Co., Ltd.
+// This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+// CANN Open Software License Agreement Version 2.0 (the "License").
+// Please refer to the License for details. You may not use this file except in compliance with the License.
+// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+// See LICENSE in the root of the software repository for the full text of the License.
+
+//===- PTOAbsorbAddPtr.cpp -----------------------------------------------===//
+//
+// Backend-shape canonical fold for VPTO memory operations:
+//
+//   op(addptr(base, A), O)  ->  op(base, A + O)
+//
+// Bisheng only emits a post-update load/store when the affine offset sits on
+// the operation itself; an offset hidden inside an `addptr` is re-materialized
+// as `VLDI + SADD` per iteration. This fold is the "addptr absorption" rule of
+// docs/designs/vpto-integer-address-canonicalization-design-zh.md and must run
+// before the post-update consumer (VPTOSoftPostUpdate).
+//
+// Legality comes only from VPTOAddressSemanticsOpInterface: a current access
+// with an Element-unit offset, a base that is an `addptr` result, no
+// updated-base post-update form, and a no-loss index addition.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PTO/IR/VPTOAddressSemantics.h"
+#include "PTO/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+namespace pto {
+#define GEN_PASS_DEF_PTOABSORBADDPTR
+#include "PTO/Transforms/Passes.h.inc"
+} // namespace pto
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::pto;
+
+namespace {
+
+struct AbsorbAddPtrIntoOpOffset final : public RewritePattern {
+  AbsorbAddPtrIntoOpOffset(MLIRContext *context)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    auto semantics = dyn_cast<pto::VPTOAddressSemanticsOpInterface>(op);
+    if (!semantics) {
+      return rewriter.notifyMatchFailure(op, "no address semantics");
+    }
+    VPTOAddressSemantics contract = semantics.getVPTOAddressSemantics();
+    if (contract.currentAccesses.empty()) {
+      return rewriter.notifyMatchFailure(op, "no current access");
+    }
+    const VPTOAddressAccess &access = contract.currentAccesses.front();
+    if (!access.offset || access.offset->unit != VPTOAddressUnit::Element) {
+      return rewriter.notifyMatchFailure(op, "offset is not element-unit");
+    }
+    // Post-update form: the offset operand denotes the after-access advance,
+    // not a current access; never fold those.
+    if (contract.postUpdate && contract.postUpdate->updatedBase) {
+      return rewriter.notifyMatchFailure(op, "already in post-update form");
+    }
+
+    Value base = access.baseOperand->get();
+    auto addptr = base.getDefiningOp<pto::AddPtrOp>();
+    if (!addptr) {
+      return rewriter.notifyMatchFailure(op, "base is not an addptr");
+    }
+    // Only fold addptr whose base is itself an integer-backed castptr — the
+    // canonical shape produced by pto-canonicalize-integer-address. Arbitrary
+    // addptr chains over user pointers are left alone: VPTOSoftPostUpdate has
+    // its own sequential base-chain handling for those, and folding them here
+    // would change (or destroy) that post-update structure.
+    if (!addptr.getPtr().getDefiningOp<pto::CastPtrOp>()) {
+      return rewriter.notifyMatchFailure(
+          op, "addptr base is not a castptr (sequential chain handled by "
+              "soft post-update)");
+    }
+    if (!addptr.getOffset().getType().isIndex() ||
+        !access.offset->operand->get().getType().isIndex()) {
+      return rewriter.notifyMatchFailure(op, "offsets are not index-typed");
+    }
+    // addptr and the op share the same pointer element type because the op's
+    // base operand *is* the addptr result (AllTypesMatch on AddPtrOp).
+
+    Value combined = rewriter.create<arith::AddIOp>(
+        op->getLoc(), addptr.getOffset(), access.offset->operand->get());
+
+    rewriter.modifyOpInPlace(op, [&]() {
+      access.baseOperand->set(addptr.getPtr());
+      access.offset->operand->set(combined);
+    });
+    if (addptr->use_empty()) {
+      rewriter.eraseOp(addptr);
+    }
+    return success();
+  }
+};
+
+struct PTOAbsorbAddPtrPass final
+    : public pto::impl::PTOAbsorbAddPtrBase<PTOAbsorbAddPtrPass> {
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+    RewritePatternSet patterns(&getContext());
+    patterns.add<AbsorbAddPtrIntoOpOffset>(&getContext());
+    if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
+      signalPassFailure();
+    }
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> mlir::pto::createPTOAbsorbAddPtrPass() {
+  return std::make_unique<PTOAbsorbAddPtrPass>();
+}