triton-lang · elstainniles · Aug 13, 2026
@@ -75,7 +75,6 @@ struct AnalyzedComponent {
 struct AnalyzedValue {
   Type originalType;
   SmallVector<AnalyzedComponent> components;
-  SmallVector<Value> invariants;
   SmallVector<Attribute> attributes;
 };
 

@@ -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,11 @@ 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;

@@ -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,

@@ -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,15 @@ 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 +801,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:

@@ -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,