feat(control-flow): decouple pointer descriptors across SCF

Carry complete block-pointer descriptors across supported for, while, and if boundaries as pointer-free SSA components, and rebuild pointers only at their use sites. Keep tensor-pointer bases outside loop signatures when they are invariant while carrying complete lane offsets.

Teach TritonToLinalg and TritonToUnstructure to lower integer pointer carriers, scalar and opaque pointer joins, descriptor loops, lane offsets, and rebased memref layouts. Preserve exact descriptor producer slots, retain legacy conversion for mixed pointer boundaries, and keep externally typed memref boundaries layout-compatible.

Remove the module-wide addptr-base restriction so unrelated and local make_tensor_ptr operations retain their previous behavior. Add one end-to-end pytest covering dynamic if/for/while block-pointer descriptors, changing bases, ordinary loop results, scalar-base tensor pointers, and opaque lane-wise tensor pointers.

Author: elstainniles

third_party/ascend/include/TritonControlFlowOpt/ControlFlowAnalysis.h

@@ -75,7 +75,6 @@ struct AnalyzedComponent {
 struct AnalyzedValue {
   Type originalType;
   SmallVector<AnalyzedComponent> components;
-  SmallVector<Value> invariants;
   SmallVector<Attribute> attributes;
 };
 

third_party/ascend/include/TritonControlFlowOpt/ControlFlowRewrite.h

@@ -32,19 +32,27 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
 
 namespace mlir::triton::controlflow {
 
+/// Handoff marker for SCF loops whose pointer slots have already been expanded
+/// into policy-owned descriptor components. Its DenseI32ArrayAttr value lists
+/// the loop-carried init/result slots occupied by pointer descriptor state.
+/// TritonToLinalg uses those slots to preserve only the required producer
+/// chains and removes the marker after conversion.
+inline constexpr llvm::StringLiteral kPointerDescriptorBoundaryAttr =
+    "PointerDescriptorBoundary";
+
 /// Policy-owned description of one value crossing a control-flow boundary.
 ///
-/// `components` are runtime values that a policy may place in an expanded SCF
-/// signature. `invariants` and `attributes` are public storage whose layout is
-/// interpreted only by the policy that creates them. The shared rewrite treats
-/// those fields as opaque and only accesses `components` directly.
+/// `components` contain every runtime value needed to rebuild the original
+/// value. A policy may place a selected subset in an expanded SCF signature.
+/// `attributes` retain non-SSA metadata. Both layouts are private to the
+/// policy; the shared rewrite never interprets pointer-specific fields.
 struct DecomposedValue {
   Type originalType;
   SmallVector<Value> components;
-  SmallVector<Value> invariants;
   SmallVector<Attribute> attributes;
 };
 
@@ -71,8 +79,9 @@ class ControlFlowRewriteContext {
 ///
 /// The policy decides how its value is decomposed and rebuilt, which components
 /// cross loop/if boundaries, and whether two decompositions share a compatible
-/// invariant schema. It is not an IR marker and carries no state between
-/// policy invocations.
+/// non-carried schema. It carries no mutable state between policy invocations;
+/// a capability hook tells the shared rewrite whether expanded loop slots must
+/// be recorded for downstream conversion.
 class ControlFlowRewritePolicy : public ControlFlowAnalysisPolicy {
 public:
   virtual ~ControlFlowRewritePolicy() = default;
@@ -81,6 +90,13 @@ class ControlFlowRewritePolicy : public ControlFlowAnalysisPolicy {
   /// after cloning so later operations can reuse their exact component state.
   virtual bool shouldDecomposeOperation(Operation *op) const = 0;
 
+  /// Whether rewritten loops owned by this policy must expose their descriptor
+  /// slots to downstream conversion. The shared rewrite owns the positional
+  /// marker because it alone knows both the previous and expanded signatures.
+  virtual bool requiresPointerDescriptorBoundaryMarker() const {
+    return false;
+  }
+
   virtual FailureOr<DecomposedValue>
   decompose(Value value, const ControlFlowRewriteContext &context,
             OpBuilder &builder, Location loc) const = 0;

third_party/ascend/include/TritonToLinalg/BlockPtrAnalysis.h

@@ -45,6 +45,16 @@ namespace triton {
 
 enum class MemAccVal { Undefined = 0, StrucMemAcc = 1, UnstrucMemAcc = 2 };
 
+/// Creates a verifier-valid HIVM pointer cast for a scalar Triton pointer
+/// represented by an integer address. Triton scalar pointers do not carry an
+/// extent, while their downstream carrier is normally `memref<?xT>` and every
+/// dynamic memref dimension requires a size operand. Use one element as the
+/// conservative carrier extent; reinterpret-cast lowering replaces it with a
+/// precise access range when a larger descriptor is materialized.
+hivm::PointerCastOp createScalarPointerCast(OpBuilder &builder, Location loc,
+                                            MemRefType resultType,
+                                            Value address);
+
 struct MemAccType {
 
   MemAccVal value;
@@ -288,8 +298,8 @@ class BlockDataParser {
                             ConversionPatternRewriter &rewriter,
                             llvm::SmallDenseMap<Value, BlockData> &known);
 
-  static void
-  rewriteMakeTensorPtrOp(triton::MakeTensorPtrOp op, Value base,
+  static LogicalResult
+  rewriteMakeTensorPtrOp(triton::MakeTensorPtrOp op, Value convertedBase,
                          ConversionPatternRewriter &rewriter,
                          llvm::SmallDenseMap<Value, BlockData> &known);
 
@@ -312,9 +322,9 @@ class BlockDataParser {
 
   /// @param known is mainly designed for `rewriteLoop`, and is just non-const
   /// in `rewriteLoop`, `rewriteAddPtr` and `rewriteAdvance`
-  static void rewriteLoopOp(LoopLikeOpInterface op,
-                            ConversionPatternRewriter &rewriter,
-                            llvm::SmallDenseMap<Value, BlockData> &known);
+  static LogicalResult
+  rewriteLoopOp(LoopLikeOpInterface op, ConversionPatternRewriter &rewriter,
+                llvm::SmallDenseMap<Value, BlockData> &known);
 
   static void rewriteAddPtrToUnstrucMemAcc(triton::AddPtrOp op,
                                            triton::AddPtrOp::Adaptor &adaptor,

third_party/ascend/include/TritonToLinalg/TritonOpConverter.h

@@ -37,6 +37,7 @@
 #include "mlir/Transforms/DialectConversion.h"
 
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 
@@ -576,6 +577,66 @@ class GatherConverter : public OpConversionPattern<triton::GatherOp> {
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+// These predicates select only scalar !tt.ptr<T> transports. A
+// tensor<...x!tt.ptr<T>> follows the separate tensor-pointer lowering.
+bool hasScalarPointerResult(scf::IfOp op);
+bool isScalarPointerSelect(arith::SelectOp op);
+
+// Marks an scf.if temporarily rebuilt by IfConverter. Its scalar-pointer
+// results are represented as complete i64 addresses, so only its own yields
+// require the matching pointer-to-address conversion.
+inline constexpr llvm::StringLiteral kScalarPointerCarrierBoundaryAttr =
+    "ScalarPointerCarrierBoundary";
+
+// Rebuild an scf.if with scalar-pointer results so the boundary carries
+// complete i64 addresses and reconstructs memrefs only after the join.
+// The original branch regions are moved into the new operation, preserving
+// side effects and allowing the conversion driver to rewrite each scf.yield
+// operand in place.
+//
+// Example:
+//   %base = scf.if %cond -> !tt.ptr<f32> {
+//     scf.yield %lhs : !tt.ptr<f32>
+//   } else {
+//     scf.yield %rhs : !tt.ptr<f32>
+//   }
+//   %ptr = tt.make_tensor_ptr %base, ...
+// becomes an scf.if returning i64 plus one hivm.pointer_cast after the if.
+class IfConverter : public OpConversionPattern<scf::IfOp> {
+public:
+  using OpConversionPattern<scf::IfOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(scf::IfOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
+// Convert a scalar-pointer select into a select over complete integer addresses
+// and reconstruct one memref after the selection. This handles both BlockPtr
+// bases and ordinary scalar pointers without asking the backend to merge two
+// memory objects.
+//
+// Example:
+//   %base = arith.select %cond, %lhs, %rhs : !tt.ptr<f32>
+//   %ptr = tt.make_tensor_ptr %base, ...
+// becomes:
+//   %lhs_addr = memref.extract_aligned_pointer_as_index %lhs
+//   %rhs_addr = memref.extract_aligned_pointer_as_index %rhs
+//   %selected_addr = arith.select %cond, %lhs_addr, %rhs_addr : i64
+//   %base = hivm.pointer_cast %selected_addr : i64 to memref<?xf32>
+class PointerSelectConverter : public OpConversionPattern<arith::SelectOp> {
+public:
+  using OpConversionPattern<arith::SelectOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::SelectOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
+// Convert the yields of an IfConverter-created scf.if to its carrier result
+// types. In particular, a yielded scalar pointer becomes its complete i64
+// address. Yields belonging to ordinary ifs or loops are intentionally left to
+// their owning conversions.
 class YieldConverter : public OpConversionPattern<scf::YieldOp> {
 public:
   using OpConversionPattern<scf::YieldOp>::OpConversionPattern;
@@ -596,11 +657,16 @@ class LoopConverter : public OpConversionPattern<LoopOpTy> {
   matchAndRewrite(LoopOpTy op,
                   typename OpConversionPattern<LoopOpTy>::OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    // CFO-expanded descriptor loops already carry pointer-free policy values
+    // and remain structurally unchanged. This legacy BlockData rewrite is only
+    // valid for explicitly marked loops.
+    if (!op->hasAttr("UnhandledLoopOp"))
+      return failure();
     llvm::SmallDenseMap<Value, BlockData> known;
 
-    op->removeAttr("UnhandledLoopOp");
-    BlockDataParser::rewriteLoopOp(op, rewriter, known);
-    return success();
+    rewriter.modifyOpInPlace(
+        op, [&]() { op->removeAttr("UnhandledLoopOp"); });
+    return BlockDataParser::rewriteLoopOp(op, rewriter, known);
   }
 };
 
@@ -736,6 +802,14 @@ class PtrToIntConverter : public OpConversionPattern<triton::PtrToIntOp> {
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+class IntToPtrConverter : public OpConversionPattern<triton::IntToPtrOp> {
+public:
+  using OpConversionPattern<triton::IntToPtrOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(triton::IntToPtrOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 class IndexPutConverter
     : public OpConversionPattern<triton::ascend::IndexPutOp> {
 public:

third_party/ascend/include/TritonToUnstructure/BubbleUpOperation.h

@@ -72,6 +72,11 @@ class BubbleUpExtract : public OpRewritePattern<ExtractOpTy> {
                          PatternRewriter &rewriter) const;
   void bubbleUpOperation(ExtractOpTy op, arith::CmpIOp parentOp, Location loc,
                          PatternRewriter &rewriter) const;
+  // Pushes extract(select(condition, lhs, rhs)) through the select. A shaped
+  // condition is extracted at the same position while a scalar condition is
+  // reused directly.
+  void bubbleUpOperation(ExtractOpTy op, arith::SelectOp parentOp, Location loc,
+                         PatternRewriter &rewriter) const;
   void bubbleUpOperation(ExtractOpTy op, arith::TruncFOp parentOp, Location loc,
                          PatternRewriter &rewriter) const;
   void bubbleUpOperation(ExtractOpTy op, arith::ExtFOp parentOp, Location loc,