[SIMT][OPT] Add persistent SIMT fragment materialization

include/PTO/Transforms/Passes.h

@@ -92,6 +92,9 @@ std::unique_ptr<Pass> createPTOFusionRegionGenPass();
 LogicalResult validateIntToPtrUses(func::FuncOp func);
 
 std::unique_ptr<Pass> createPTOUnrollSIMTForPass();
+std::unique_ptr<Pass> createPTOAnalyzeSIMTPersistentFragmentPass();
+std::unique_ptr<Pass> createPTOMaterializeSIMTPersistentFragmentPass();
+std::unique_ptr<Pass> createPTOOutlineSIMTSectionsPass();
 std::unique_ptr<Pass> createPTOInferVPTOVecScopePass();
 std::unique_ptr<Pass> createVPTOExpandWrapperOpsPass();
 std::unique_ptr<Pass> createPTOVPTOPtrBoundaryPass();

include/PTO/Transforms/Passes.td

@@ -813,10 +813,12 @@ def PTOResolveBufferSelect : Pass<"pto-resolve-buffer-select", "ModuleOp"> {
 
 def PTOUnrollSIMTFor : Pass<"pto-unroll-simt-for", "func::FuncOp"> {
   let summary =
-      "Unroll explicitly annotated scf.for loops in pto.simt_entry functions";
+      "Unroll explicitly annotated scf.for loops in SIMT contexts";
   let description = [{
-    Walks functions with `pto.simt_entry` attribute and fully unrolls `scf.for`
-    loops that carry the `{pto.unroll = "full"}` attribute.
+    Walks functions with the `pto.simt_entry` attribute or inline
+    `pto.section.simt` regions and fully unrolls `scf.for` loops that carry the
+    `{pto.unroll = "full"}` attribute. Annotated loops outside an inline SIMT
+    section in an ordinary function are left unchanged.
 
     Only loops with constant lower bound, upper bound, and positive step are
     unrolled.  The annotation is opt-in: loops without the attribute are
@@ -838,6 +840,70 @@ def PTOUnrollSIMTFor : Pass<"pto-unroll-simt-for", "func::FuncOp"> {
   ];
 }
 
+def PTOAnalyzeSIMTPersistentFragment
+    : Pass<"pto-analyze-simt-persistent-fragment", "func::FuncOp"> {
+  let summary = "Analyze SIMT persistent fragments without changing IR";
+  let description = [{
+    Validates persistent fragment definitions, access graphs, init/carry
+    sections, resident elements, and keep/resume slot assignments. This pass
+    is analysis-only; it does not insert or remove operations.
+  }];
+  let constructor =
+      "mlir::pto::createPTOAnalyzeSIMTPersistentFragmentPass()";
+  let dependentDialects = [
+    "mlir::func::FuncDialect",
+    "mlir::LLVM::LLVMDialect",
+    "mlir::pto::PTODialect"
+  ];
+}
+
+def PTOMaterializeSIMTPersistentFragment
+    : Pass<"pto-materialize-simt-persistent-fragment", "func::FuncOp"> {
+  let summary =
+      "Materialize SIMT persistent fragments as keep/resume slot traffic";
+  let description = [{
+    Consumes `llvm.alloca` values marked with `pto.persistent` in a kernel body
+    before inline SIMT sections are outlined. The first implementation handles
+    single-block `pto.section.simt` bodies with scalar and fixed-width vector
+    LLVM load/store accesses through statically-resolved offsets. Vector
+    accesses are split into scalar lanes before the section-local state is
+    rewritten into `pto.resume` at entry and `pto.keep` at exit.
+
+    Unsupported persistent fragment shapes fail fast because the allocation is
+    only a lexical carrier for register-resident state, not real stack memory.
+  }];
+  let constructor =
+      "mlir::pto::createPTOMaterializeSIMTPersistentFragmentPass()";
+  let dependentDialects = [
+    "mlir::func::FuncDialect",
+    "mlir::LLVM::LLVMDialect",
+    "mlir::pto::PTODialect",
+    "mlir::arith::ArithDialect"
+  ];
+}
+
+def PTOOutlineSIMTSections
+    : Pass<"pto-outline-simt-sections", "ModuleOp"> {
+  let summary =
+      "Outline inline pto.section.simt regions into pto.simt_entry functions";
+  let description = [{
+    Converts each inline `pto.section.simt` region into a private helper
+    function marked with `pto.simt_entry`, captures external SSA values as
+    helper arguments, and replaces the original section with a
+    `pto.simt_launch` using the section launch dimensions.
+
+    This is the late outline step used by TileLang/PTODSL style frontends that
+    keep SIMT code inline until PTOAS has had a chance to materialize
+    persistent SIMT fragments.
+  }];
+  let constructor = "mlir::pto::createPTOOutlineSIMTSectionsPass()";
+  let dependentDialects = [
+    "mlir::func::FuncDialect",
+    "mlir::pto::PTODialect",
+    "mlir::arith::ArithDialect"
+  ];
+}
+
 def PTOInferVPTOVecScope
     : Pass<"pto-infer-vpto-vecscope", "func::FuncOp"> {
   let summary =

lib/PTO/IR/PTO.cpp

@@ -17696,18 +17696,6 @@ static LogicalResult verifyUniqueKeepGroupSlots(KeepOp current,
   return success();
 }
 
-static LogicalResult verifySimtKeepResumeCommon(Operation *op, int64_t slot) {
-  func::FuncOp func = getParentFunc(op);
-  if (!func || !func->hasAttr(pto::kPTOSimtEntryAttrName))
-    return op->emitOpError("must appear inside a function marked with '")
-           << pto::kPTOSimtEntryAttrName << "'";
-  if (slot < 0 || slot >= kSimtKeepResumeSlotLimit) {
-    return op->emitOpError("requires slot in range [0, ")
-           << (kSimtKeepResumeSlotLimit - 1) << "]";
-  }
-  return success();
-}
-
 static bool isSupportedSimtKeepResumeType(Type type) {
   if (auto intType = dyn_cast<IntegerType>(type))
     return intType.getWidth() <= 64;
@@ -17728,6 +17716,17 @@ static LogicalResult verifyInsideSimtExecutionScope(Operation *op) {
   return success();
 }
 
+static LogicalResult verifySimtKeepResumeCommon(Operation *op, int64_t slot) {
+  if (!isInsideSimtExecutionScope(op))
+    return op->emitOpError("must appear inside a function marked with '")
+           << pto::kPTOSimtEntryAttrName << "' or inside pto.section.simt";
+  if (slot < 0 || slot >= kSimtKeepResumeSlotLimit) {
+    return op->emitOpError("requires slot in range [0, ")
+           << (kSimtKeepResumeSlotLimit - 1) << "]";
+  }
+  return success();
+}
+
 LogicalResult SyncthreadsOp::verify() {
   return verifyInsideSimtExecutionScope(getOperation());
 }
@@ -17778,20 +17777,31 @@ LogicalResult KeepOp::verify() {
     return failure();
 
   Block *block = getOperation()->getBlock();
-  Operation *terminator = block ? block->getTerminator() : nullptr;
-  if (!terminator || !isa<func::ReturnOp>(terminator))
+  bool insideSection =
+      getOperation()->getParentOfType<SectionSimtOp>() != nullptr;
+  Operation *lastPayloadOp = nullptr;
+  if (insideSection) {
+    if (!block->empty())
+      lastPayloadOp = &block->back();
+  } else {
+    Operation *terminator = block->getTerminator();
+    if (isa<func::ReturnOp>(terminator))
+      lastPayloadOp = terminator->getPrevNode();
+  }
+  if (!lastPayloadOp)
     return emitOpError(
-        "must be placed in the SIMT epilogue before func.return");
+        "must be placed in the SIMT epilogue before func.return or the end "
+        "of pto.section.simt");
 
-  Operation *cur = terminator->getPrevNode();
+  Operation *cur = lastPayloadOp;
   while (cur && isa<SyncthreadsOp>(cur))
     cur = cur->getPrevNode();
   Operation *lastKeep = cur;
   if (!lastKeep || !isa<KeepOp>(lastKeep))
     return emitOpError()
            << "must be placed in the SIMT epilogue before func.return; only "
               "'pto.syncthreads' may appear between the final 'pto.keep' group "
-              "and func.return";
+              "and func.return or the end of pto.section.simt";
 
   Operation *firstKeep = lastKeep;
   while (Operation *prev = firstKeep->getPrevNode()) {
@@ -17802,7 +17812,8 @@ LogicalResult KeepOp::verify() {
   if (!isOpInRange(getOperation(), firstKeep, lastKeep))
     return emitOpError()
            << "must be in the contiguous SIMT keep epilogue group immediately "
-              "before optional 'pto.syncthreads' and func.return";
+              "before optional 'pto.syncthreads' and func.return or the end "
+              "of pto.section.simt";
   if (failed(verifyUniqueKeepGroupSlots(*this, firstKeep, lastKeep)))
     return failure();
   return success();

lib/PTO/Transforms/CMakeLists.txt

@@ -35,6 +35,10 @@ add_mlir_dialect_library(PTOTransforms
   VPTOBufferMaterialization.cpp
   PTOValidateVPTOIR.cpp
   PTOUnrollSIMTForPass.cpp
+  SIMTPersistentFragmentAnalysis.cpp
+  PTOAnalyzeSIMTPersistentFragment.cpp
+  PTOMaterializeSIMTPersistentFragment.cpp
+  PTOOutlineSIMTSections.cpp
   PTOInferVPTOVecScope.cpp
 
   InsertSync/PTOInsertSync.cpp

lib/PTO/Transforms/PTOAnalyzeSIMTPersistentFragment.cpp

@@ -0,0 +1,52 @@
+// 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.
+
+#include "PTO/Transforms/Passes.h"
+#include "SIMTPersistentFragmentAnalysis.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace pto {
+#define GEN_PASS_DEF_PTOANALYZESIMTPERSISTENTFRAGMENT
+#include "PTO/Transforms/Passes.h.inc"
+} // namespace pto
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+struct PTOAnalyzeSIMTPersistentFragmentPass
+    : public pto::impl::PTOAnalyzeSIMTPersistentFragmentBase<
+          PTOAnalyzeSIMTPersistentFragmentPass> {
+  using pto::impl::PTOAnalyzeSIMTPersistentFragmentBase<
+      PTOAnalyzeSIMTPersistentFragmentPass>::
+      PTOAnalyzeSIMTPersistentFragmentBase;
+
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+    if (func.isExternal())
+      return;
+
+    const auto &analysis = getAnalysis<pto::SIMTPersistentFragmentAnalysis>();
+    if (!analysis.isValid()) {
+      signalPassFailure();
+      return;
+    }
+
+    markAllAnalysesPreserved();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> mlir::pto::createPTOAnalyzeSIMTPersistentFragmentPass() {
+  return std::make_unique<PTOAnalyzeSIMTPersistentFragmentPass>();
+}