diff --git a/third_party/ascend/include/TritonControlFlowOpt/BlockPtrDecompose.h b/third_party/ascend/include/TritonControlFlowOpt/BlockPtrDecompose.h deleted file mode 100644 index 678ccde0db..0000000000 --- a/third_party/ascend/include/TritonControlFlowOpt/BlockPtrDecompose.h +++ /dev/null @@ -1,37 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#ifndef TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_BLOCK_PTR_DECOMPOSE_H -#define TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_BLOCK_PTR_DECOMPOSE_H - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Support/LogicalResult.h" - -namespace mlir::triton::controlflow { - -/// Internal decomposition entry point used by TritonControlFlowOpt. -/// Decomposes block-pointer control-flow slots into descriptor components. -LogicalResult runBlockPtrDecompose(ModuleOp module); - -} // namespace mlir::triton::controlflow - -#endif // TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_BLOCK_PTR_DECOMPOSE_H diff --git a/third_party/ascend/include/TritonControlFlowOpt/CFGStructuring.h b/third_party/ascend/include/TritonControlFlowOpt/CFGStructuring.h deleted file mode 100644 index 029ae71c4a..0000000000 --- a/third_party/ascend/include/TritonControlFlowOpt/CFGStructuring.h +++ /dev/null @@ -1,40 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#ifndef TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_CFG_STRUCTURING_H -#define TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_CFG_STRUCTURING_H - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Support/LogicalResult.h" - -namespace mlir::triton::controlflow { - -/// Internal entry point used by the composite TritonControlFlowOpt pass. -/// -/// Converts supported acyclic, tree-like function CFGs in `module` to SCF. It -/// is intentionally not registered as an independent pass because pointer -/// decomposition requires structured control flow. -LogicalResult structureCFG(ModuleOp module); - -} // namespace mlir::triton::controlflow - -#endif // TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_CFG_STRUCTURING_H diff --git a/third_party/ascend/include/TritonControlFlowOpt/ControlFlowAnalysis.h b/third_party/ascend/include/TritonControlFlowOpt/ControlFlowAnalysis.h deleted file mode 100644 index 1b09ca56a7..0000000000 --- a/third_party/ascend/include/TritonControlFlowOpt/ControlFlowAnalysis.h +++ /dev/null @@ -1,195 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#ifndef TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_CONTROL_FLOW_ANALYSIS_H -#define TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_CONTROL_FLOW_ANALYSIS_H - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/Operation.h" -#include "mlir/Support/LogicalResult.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/SmallVector.h" - -namespace mlir::triton::controlflow { - -/// Classification used by the signature planner. Recomputed components are -/// intentionally not modeled yet; they can be added when IV extraction becomes -/// part of the analysis rather than the existing block-pointer rewrite helper. -enum class ComponentTransferKind { Invariant, Transferred }; - -/// Symbolic identity of a component during read-only analysis. -/// -/// A component produced by an ordinary SSA value is identified by that value -/// and its policy-defined component index. Zero represents a canonical zero -/// offsets value without materializing an arith.constant in the source IR. -struct ComponentIdentity { - enum class Kind { Value, Zero } kind = Kind::Value; - Value value; - unsigned componentIndex = 0; - - static ComponentIdentity fromValue(Value value, unsigned componentIndex) { - return {Kind::Value, value, componentIndex}; - } - static ComponentIdentity zero(unsigned componentIndex = 0) { - return {Kind::Zero, {}, componentIndex}; - } - - bool operator==(const ComponentIdentity &other) const { - return kind == other.kind && value == other.value && - componentIndex == other.componentIndex; - } - bool operator!=(const ComponentIdentity &other) const { - return !(*this == other); - } -}; - -struct AnalyzedComponent { - Type type; - ComponentIdentity identity; -}; - -/// Policy-owned abstract pointer state. Unlike DecomposedValue, this structure -/// never contains newly created IR values and is therefore safe to compute -/// before a rewrite starts. -struct AnalyzedValue { - Type originalType; - SmallVector components; - SmallVector invariants; - SmallVector attributes; -}; - -/// Signature decision for one original result/iter-argument position. -struct ControlFlowSlotAnalysis { - unsigned oldIndex = 0; - SmallVector componentKinds; - SmallVector componentIndices; - SmallVector componentTypes; -}; - -/// Cached decision for one structured control-flow operation. -struct ControlFlowOpAnalysis { - SmallVector slots; - bool hasNestedRewrite = false; - - bool rewritesOwnSignature() const { return !slots.empty(); } - bool needsRewrite() const { - return rewritesOwnSignature() || hasNestedRewrite; - } -}; - -/// Immutable operation-level contract consumed after read-only analysis. -/// -/// Temporary Value -> AnalyzedValue state is deliberately excluded: those -/// values may be erased while earlier roots are rewritten, whereas a root and -/// its nested operations remain valid until that root is processed. -struct ControlFlowRewritePlan { - llvm::DenseMap operations; - - const ControlFlowOpAnalysis *lookup(Operation *op) const; -}; - -class ControlFlowAnalysisContext; - -/// Pointer-specific part of the read-only analysis. Implementations describe -/// their component layout and merge rules; the common analyzer owns SCF region -/// traversal, argument/result correspondence, and nested-op caching. -class ControlFlowAnalysisPolicy { -public: - virtual ~ControlFlowAnalysisPolicy() = default; - - virtual bool matches(Type type) const = 0; - - /// Returns whether a value belongs to this decomposition stage. Pointer - /// policies use the type-based default. A future StructuredOffsets policy - /// can override this with the result of a backward address-demand analysis. - virtual bool isDecompositionTarget(Value value) const { - return matches(value.getType()); - } - - virtual FailureOr - analyzeValue(Value value, ControlFlowAnalysisContext &context) const = 0; - - /// Components which may legally be carried by a loop. Other components must - /// remain symbolically identical across the backedge. - virtual FailureOr> - getLoopCandidateComponents(const AnalyzedValue &value) const = 0; - - virtual FailureOr> - getLoopTransferredComponents(const AnalyzedValue &initial, - const AnalyzedValue ®ionArgument, - const AnalyzedValue &next) const = 0; - - virtual FailureOr> - getIfTransferredComponents(const AnalyzedValue &thenValue, - const AnalyzedValue &elseValue) const = 0; - - /// Selects the component type used in the replacement SCF signature. - virtual FailureOr joinComponentTypes(Type lhs, Type rhs) const = 0; -}; - -/// Stage-scoped transient cache used while computing a rewrite plan. One -/// context analyzes all outermost roots before any IR mutation; only the -/// operation-level decisions survive in ControlFlowRewritePlan. -class ControlFlowAnalysisContext { -public: - explicit ControlFlowAnalysisContext(const ControlFlowAnalysisPolicy &policy) - : policy(policy) {} - - FailureOr analyzeValue(Value value); - FailureOr analyzeControlFlowOp(Operation *op); - - ControlFlowRewritePlan takeRewritePlan() &&; - - const AnalyzedValue *lookupValue(Value value) const; - const ControlFlowOpAnalysis *lookup(Operation *op) const; - -private: - LogicalResult analyzeNestedOperations(Block *block, bool &hasNestedRewrite); - FailureOr analyzeFor(Operation *op); - FailureOr analyzeWhile(Operation *op); - FailureOr analyzeIf(Operation *op); - - void bindRegionArgument(Value argument, const AnalyzedValue &initial, - ArrayRef componentIndices); - FailureOr> - getTransferredTypes(const AnalyzedValue &lhs, const AnalyzedValue &rhs, - ArrayRef componentIndices) const; - - const ControlFlowAnalysisPolicy &policy; - llvm::DenseMap analyzedValues; - llvm::DenseMap analyzedOps; - llvm::DenseSet operationsBeingAnalyzed; -}; - -/// Returns SCF roots that are not nested in another supported SCF operation. -SmallVector collectOutermostControlFlowOps(ModuleOp module); - -/// Analyzes every control-flow root for one decomposition stage before any IR -/// mutation and freezes only the operation-level rewrite decisions. -FailureOr -analyzeControlFlow(ModuleOp module, const ControlFlowAnalysisPolicy &policy); - -} // namespace mlir::triton::controlflow - -#endif // TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_CONTROL_FLOW_ANALYSIS_H diff --git a/third_party/ascend/include/TritonControlFlowOpt/ControlFlowRewrite.h b/third_party/ascend/include/TritonControlFlowOpt/ControlFlowRewrite.h deleted file mode 100644 index e77c2b104c..0000000000 --- a/third_party/ascend/include/TritonControlFlowOpt/ControlFlowRewrite.h +++ /dev/null @@ -1,112 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#ifndef TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_CONTROL_FLOW_REWRITE_H -#define TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_CONTROL_FLOW_REWRITE_H - -#include "TritonControlFlowOpt/ControlFlowAnalysis.h" - -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/Support/LogicalResult.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/SmallVector.h" - -namespace mlir::triton::controlflow { - -/// 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. -struct DecomposedValue { - Type originalType; - SmallVector components; - SmallVector invariants; - SmallVector attributes; -}; - -/// Read-only view of the SSA mapping and decompositions visible at the current -/// rewrite point. Policies use it while recursively analyzing pointer -/// producers; the state exists only for one control-flow rewrite attempt. The -/// referenced mappings are not owned and must outlive the context. -class ControlFlowRewriteContext { -public: - ControlFlowRewriteContext( - const IRMapping &valueMapping, - const llvm::DenseMap &decomposedValues) - : valueMapping(valueMapping), decomposedValues(decomposedValues) {} - - Value remap(Value value) const; - const DecomposedValue *lookup(Value value) const; - -private: - const IRMapping &valueMapping; - const llvm::DenseMap &decomposedValues; -}; - -/// Pointer-semantics interface implemented by each decomposition policy. -/// -/// 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. -class ControlFlowRewritePolicy : public ControlFlowAnalysisPolicy { -public: - virtual ~ControlFlowRewritePolicy() = default; - - /// Whether results of this ordinary operation need immediate decomposition - /// after cloning so later operations can reuse their exact component state. - virtual bool shouldDecomposeOperation(Operation *op) const = 0; - - virtual FailureOr - decompose(Value value, const ControlFlowRewriteContext &context, - OpBuilder &builder, Location loc) const = 0; - - virtual Value recompose(const DecomposedValue &value, OpBuilder &builder, - Location loc) const = 0; -}; - -/// Rewrites supported SCF operations from outermost to innermost. -/// -/// Applies a previously frozen plan without running value analysis again. -/// Signature expansion, region cloning, terminator rewriting, nested recursion -/// and result replacement are driven solely by operation-level decisions in -/// `plan`. -LogicalResult -applyControlFlowRewritePlan(ModuleOp module, - const ControlFlowRewritePolicy &policy, - const ControlFlowRewritePlan &plan); - -/// Analyzes the complete decomposition stage before mutating the IR, then -/// applies the frozen plan from outermost to innermost. Pointer semantics -/// remain selected by `policy` so different decompositions share the same SCF -/// plumbing. -LogicalResult rewriteControlFlow(ModuleOp module, - const ControlFlowRewritePolicy &policy); - -} // namespace mlir::triton::controlflow - -#endif // TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_CONTROL_FLOW_REWRITE_H diff --git a/third_party/ascend/include/TritonControlFlowOpt/TensorPtrDecompose.h b/third_party/ascend/include/TritonControlFlowOpt/TensorPtrDecompose.h deleted file mode 100644 index 4fd63928d9..0000000000 --- a/third_party/ascend/include/TritonControlFlowOpt/TensorPtrDecompose.h +++ /dev/null @@ -1,39 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#ifndef TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_TENSOR_PTR_DECOMPOSE_H -#define TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_TENSOR_PTR_DECOMPOSE_H - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Support/LogicalResult.h" - -namespace mlir::triton::controlflow { - -/// Internal decomposition entry point used by TritonControlFlowOpt. -/// Replaces common-base tensor-of-pointers control-flow slots with complete -/// offsets. The common base remains outside the rewritten control-flow -/// signature and is used only to rebuild the original pointer value. -LogicalResult runTensorPtrDecompose(ModuleOp module); - -} // namespace mlir::triton::controlflow - -#endif // TRITON_ASCEND_TRITON_CONTROL_FLOW_OPT_TENSOR_PTR_DECOMPOSE_H diff --git a/third_party/ascend/include/Utils/Utils.h b/third_party/ascend/include/Utils/Utils.h index 51e0266f0e..82b5c5f3b3 100644 --- a/third_party/ascend/include/Utils/Utils.h +++ b/third_party/ascend/include/Utils/Utils.h @@ -31,7 +31,6 @@ #include "mlir/IR/OpDefinition.h" #include "mlir/IR/Operation.h" #include "mlir/IR/PatternMatch.h" -#include "mlir/IR/TypeRange.h" #include "mlir/Transforms/DialectConversion.h" #include "triton/Dialect/Triton/IR/Dialect.h" #include "llvm/ADT/ArrayRef.h" @@ -219,31 +218,14 @@ inline constexpr unsigned kDotAccIntWidth = 32; class OpBuilder; -enum class IntegerExtensionKind { - Signed, - Unsigned, -}; - -/// Returns true when both ranges have identical ordered type signatures. -bool haveSameTypes(TypeRange lhs, TypeRange rhs); - -FailureOr -castIntegerLike(OpBuilder &builder, Location loc, Value value, Type targetType, - IntegerExtensionKind extension = IntegerExtensionKind::Signed); - -/// Without an explicit result type, preserve equal types and widen signless -/// integers. Mixed index/integer operands require an explicit result type. OpFoldResult addOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b, - Type resultType = {}); + const Location &loc, OpBuilder &b); OpFoldResult subOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, const Location &loc, OpBuilder &b); -/// Uses the same result-type rules as addOpFoldResult. OpFoldResult mulOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b, - Type resultType = {}); + const Location &loc, OpBuilder &b); OpFoldResult divOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, const Location &loc, OpBuilder &b); @@ -291,6 +273,8 @@ bool isZero(const OpFoldResult ofr); bool isOne(const OpFoldResult ofr); +Value convertToIndexIfNeeded(Value intValue, const Location &loc, OpBuilder &b); + RankedTensorType getExtractSlicedType(ArrayRef shape, const llvm::SmallBitVector &droppedDims, Type elemType); diff --git a/third_party/ascend/lib/TritonControlFlowOpt/BlockPtrDecompose.cpp b/third_party/ascend/lib/TritonControlFlowOpt/BlockPtrDecompose.cpp deleted file mode 100644 index ae9f6a247f..0000000000 --- a/third_party/ascend/lib/TritonControlFlowOpt/BlockPtrDecompose.cpp +++ /dev/null @@ -1,317 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#include "TritonControlFlowOpt/BlockPtrDecompose.h" - -#include "TritonControlFlowOpt/ControlFlowRewrite.h" -#include "Utils/Utils.h" - -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "llvm/ADT/STLExtras.h" - -using namespace mlir; -using namespace mlir::triton; -using namespace mlir::triton::controlflow; - -namespace { - -/// Block-pointer component layout used only by this policy: -/// -/// components = [shape..., strides..., offsets...] -/// invariants = [base] -/// attributes = [order] -/// -/// Loops currently carry only `offsets`; shape and strides must remain -/// invariant across a backedge. An scf.if may select any component whose SSA -/// value differs between its branches. -/// -/// Keep rank/layout checks local to this file: the generic control-flow -/// machinery intentionally does not know the descriptor format of a policy. -// Extracts the block-pointer rank from the original type shared by analysis -// and rewrite states. Invalid pointer and pointee types fail consistently. -static FailureOr getRank(Type originalType) { - auto pointerType = dyn_cast(originalType); - if (!pointerType) - return failure(); - auto tensorType = dyn_cast(pointerType.getPointeeType()); - if (!tensorType) - return failure(); - return tensorType.getRank(); -} - -// Validates the common block-pointer schema for either state representation. -// Their component element types differ, but this check only needs field sizes. -template static bool hasValidLayout(const StateT &state) { - FailureOr rank = getRank(state.originalType); - return succeeded(rank) && state.components.size() == 3 * *rank && - state.invariants.size() == 1 && state.attributes.size() == 1 && - isa(state.attributes.front()); -} - -class BlockPtrPolicy final : public ControlFlowRewritePolicy { -public: - bool matches(Type type) const override { - // A Triton block pointer is a scalar !tt.ptr whose pointee is a ranked - // tensor. Tensor-of-pointers are handled by their own decomposition. - auto pointerType = dyn_cast(type); - return pointerType && isa(pointerType.getPointeeType()); - } - - FailureOr - analyzeValue(Value value, - ControlFlowAnalysisContext &context) const override { - // Region arguments and control-flow results are installed by the generic - // analysis after merging their incoming abstract component states. - if (const AnalyzedValue *known = context.lookupValue(value)) { - if (!matches(known->originalType)) - return failure(); - return *known; - } - - if (auto makePtr = value.getDefiningOp()) { - // make_tensor_ptr exposes the complete descriptor directly. Record only - // types and symbolic identities here; this phase must not create IR. - AnalyzedValue result; - result.originalType = value.getType(); - unsigned componentIndex = 0; - auto appendComponents = [&](ValueRange values) { - for (Value component : values) { - result.components.push_back( - {component.getType(), - ComponentIdentity::fromValue(component, componentIndex++)}); - } - }; - appendComponents(makePtr.getShape()); - appendComponents(makePtr.getStrides()); - appendComponents(makePtr.getOffsets()); - result.invariants.push_back(makePtr.getBase()); - result.attributes.push_back(makePtr.getOrderAttr()); - if (!hasValidLayout(result)) - return failure(); - return result; - } - - auto advance = value.getDefiningOp(); - if (!advance) - return failure(); - // tt.advance preserves base/shape/strides/order and produces new offsets. - // Give those offsets identities tied to the result so a loop/if merge can - // detect that they differ without materializing arith.addi operations. - FailureOr result = context.analyzeValue(advance.getPtr()); - if (failed(result) || !hasValidLayout(*result)) - return failure(); - unsigned rank = *getRank(result->originalType); - if (advance.getOffsets().size() != rank) - return failure(); - for (unsigned dimension = 0; dimension < rank; ++dimension) { - unsigned componentIndex = 2 * rank + dimension; - result->components[componentIndex].identity = - ComponentIdentity::fromValue(value, componentIndex); - } - result->originalType = value.getType(); - return *result; - } - - FailureOr> - getLoopCandidateComponents(const AnalyzedValue &value) const override { - if (!hasValidLayout(value)) - return failure(); - unsigned rank = *getRank(value.originalType); - SmallVector indices; - // Only the final rank entries (offsets) are legal loop-carried state in the - // current block-pointer model. - for (unsigned dimension = 0; dimension < rank; ++dimension) - indices.push_back(2 * rank + dimension); - return indices; - } - - FailureOr> - getLoopTransferredComponents(const AnalyzedValue &initial, - const AnalyzedValue ®ionArgument, - const AnalyzedValue &next) const override { - if (!hasValidLayout(initial) || !hasValidLayout(regionArgument) || - !hasValidLayout(next) || - initial.originalType != regionArgument.originalType || - initial.originalType != next.originalType || - initial.invariants != regionArgument.invariants || - initial.invariants != next.invariants || - initial.attributes != regionArgument.attributes || - initial.attributes != next.attributes) - return failure(); - - unsigned rank = *getRank(initial.originalType); - // The current implementation does not expand shape or stride iter_args. - // Reject a loop that changes either instead of silently reconstructing a - // descriptor with stale values. - for (unsigned index = 0; index < 2 * rank; ++index) { - if (initial.components[index].type != next.components[index].type || - initial.components[index].identity != next.components[index].identity) - return failure(); - } - - SmallVector transferred; - // An offset is carried only if the backedge state depends on the region - // argument. Constant/invariant offsets remain outside the loop signature. - for (unsigned dimension = 0; dimension < rank; ++dimension) { - unsigned index = 2 * rank + dimension; - if (failed(joinComponentTypes(initial.components[index].type, - next.components[index].type))) - return failure(); - if (regionArgument.components[index].identity != - next.components[index].identity) - transferred.push_back(index); - } - return transferred; - } - - FailureOr> - getIfTransferredComponents(const AnalyzedValue &thenValue, - const AnalyzedValue &elseValue) const override { - if (!hasValidLayout(thenValue) || !hasValidLayout(elseValue) || - thenValue.originalType != elseValue.originalType || - thenValue.invariants != elseValue.invariants || - thenValue.attributes != elseValue.attributes || - thenValue.components.size() != elseValue.components.size()) - return failure(); - - // Unlike loops, an if can select shape, stride, or offset components. Base - // and order remain invariants because AdapterIR cannot represent a runtime - // selection between heterogeneous pointer descriptors. - SmallVector transferred; - for (unsigned index = 0; index < thenValue.components.size(); ++index) { - if (failed(joinComponentTypes(thenValue.components[index].type, - elseValue.components[index].type))) - return failure(); - if (thenValue.components[index].identity != - elseValue.components[index].identity) - transferred.push_back(index); - } - return transferred; - } - - FailureOr joinComponentTypes(Type lhs, Type rhs) const override { - // Block-pointer descriptor operands must have identical types on all - // incoming paths. Tensor-pointer offsets use a more permissive integer - // width join in their own policy. - if (lhs != rhs) - return failure(); - return lhs; - } - - bool shouldDecomposeOperation(Operation *op) const override { - // Recording each cloned advance lets downstream advances reuse its - // flattened descriptor rather than walking through the rebuilt pointer. - return isa(op); - } - - FailureOr decompose(Value value, - const ControlFlowRewriteContext &context, - OpBuilder &builder, - Location loc) const override { - // Prefer decompositions recorded while rebuilding the enclosing region; - // this is how analysis results cross nested SCF boundaries at rewrite time. - if (const DecomposedValue *known = context.lookup(value)) { - if (!matches(known->originalType)) - return failure(); - return *known; - } - - value = context.remap(value); - if (auto makePtr = value.getDefiningOp()) { - // Materialize the concrete counterpart of analyzeValue's descriptor. - DecomposedValue result; - result.originalType = value.getType(); - result.components.append(makePtr.getShape().begin(), - makePtr.getShape().end()); - result.components.append(makePtr.getStrides().begin(), - makePtr.getStrides().end()); - result.components.append(makePtr.getOffsets().begin(), - makePtr.getOffsets().end()); - result.invariants.push_back(makePtr.getBase()); - result.attributes.push_back(makePtr.getOrderAttr()); - if (!hasValidLayout(result)) - return failure(); - return result; - } - - auto advance = value.getDefiningOp(); - if (!advance) - return failure(); - - FailureOr result = - decompose(advance.getPtr(), context, builder, loc); - if (failed(result) || !hasValidLayout(*result)) - return failure(); - FailureOr rank = getRank(result->originalType); - if (advance.getOffsets().size() != *rank) - return failure(); - - // Flatten an advance into offset arithmetic at the original operation's - // position. Base, shape, strides, and order are inherited unchanged. - OpBuilder::InsertionGuard guard(builder); - builder.setInsertionPoint(advance); - for (auto [dim, delta] : llvm::enumerate(advance.getOffsets())) { - unsigned component = 2 * *rank + dim; - Value currentOffset = result->components[component]; - Value remappedDelta = context.remap(delta); - if (!remappedDelta) - return failure(); - Value offset = dyn_cast( - addOpFoldResult(currentOffset, remappedDelta, advance.getLoc(), - builder, currentOffset.getType())); - if (!offset) - return failure(); - result->components[component] = offset; - } - result->originalType = value.getType(); - return *result; - } - - Value recompose(const DecomposedValue &value, OpBuilder &builder, - Location loc) const override { - // Rebuild the original pointer type immediately inside/after the rewritten - // control-flow boundary so ordinary users remain untouched. - if (!hasValidLayout(value)) - return nullptr; - unsigned rank = *getRank(value.originalType); - auto order = cast(value.attributes.front()); - return builder.create( - loc, value.originalType, value.invariants.front(), - ValueRange(value.components).take_front(rank), - ValueRange(value.components).slice(rank, rank), - ValueRange(value.components).take_back(rank), order); - } -}; - -} // namespace - -namespace mlir::triton::controlflow { - -LogicalResult runBlockPtrDecompose(ModuleOp module) { - // Make the explicit descriptor carried by a block pointer cross each - // supported SCF boundary as ordinary SSA components. - BlockPtrPolicy policy; - return rewriteControlFlow(module, policy); -} - -} // namespace mlir::triton::controlflow diff --git a/third_party/ascend/lib/TritonControlFlowOpt/CFGStructuring.cpp b/third_party/ascend/lib/TritonControlFlowOpt/CFGStructuring.cpp deleted file mode 100644 index a8c9553f15..0000000000 --- a/third_party/ascend/lib/TritonControlFlowOpt/CFGStructuring.cpp +++ /dev/null @@ -1,917 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#include "TritonControlFlowOpt/CFGStructuring.h" -#include "Utils/Utils.h" - -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/IR/Visitors.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallPtrSet.h" -#include "llvm/ADT/SmallVector.h" - -#include - -using namespace mlir; -using namespace triton; - -namespace { - -// This transformation handles acyclic CFGs made from cf.br/cf.cond_br and -// return-like terminators. A branch with a common successor becomes an scf.if -// whose results model that successor's block arguments. A branch whose arms -// terminate independently is handled by the terminal-path builder below. - -//===----------------------------------------------------------------------===// -// CFG discovery and join selection -//===----------------------------------------------------------------------===// - -static bool isSupportedReturn(Operation *op) { - return isa( - op); -} - -static SmallVector getCfgSuccessors(Block *block) { - Operation *term = block->getTerminator(); - if (auto br = dyn_cast(term)) - return {br.getDest()}; - if (auto condBr = dyn_cast(term)) - return {condBr.getTrueDest(), condBr.getFalseDest()}; - return {}; -} - -static DenseMap computeDistances(Block *start) { - DenseMap distances; - SmallVector worklist; - - distances[start] = 0; - worklist.push_back(start); - - for (unsigned i = 0; i < worklist.size(); ++i) { - Block *block = worklist[i]; - unsigned nextDistance = distances[block] + 1; - for (Block *successor : getCfgSuccessors(block)) { - if (successor->getParent() != start->getParent()) - continue; - if (distances.count(successor)) - continue; - distances[successor] = nextDistance; - worklist.push_back(successor); - } - } - - return distances; -} - -/// Finds the closest block reachable from both branch arms. Minimizing the -/// maximum arm distance prefers the earliest balanced convergence point; the -/// total distance provides deterministic tie-breaking. -static FailureOr findNearestCommonBlock(Block *lhs, Block *rhs, - Location loc, - bool emitDiagnostic = true) { - DenseMap lhsDistances = computeDistances(lhs); - DenseMap rhsDistances = computeDistances(rhs); - - Block *best = nullptr; - unsigned bestMaxDistance = std::numeric_limits::max(); - unsigned bestTotalDistance = std::numeric_limits::max(); - - for (auto &entry : lhsDistances) { - Block *candidate = entry.first; - auto rhsIt = rhsDistances.find(candidate); - if (rhsIt == rhsDistances.end()) - continue; - - unsigned lhsDistance = entry.second; - unsigned rhsDistance = rhsIt->second; - unsigned maxDistance = std::max(lhsDistance, rhsDistance); - unsigned totalDistance = lhsDistance + rhsDistance; - if (maxDistance < bestMaxDistance || - (maxDistance == bestMaxDistance && totalDistance < bestTotalDistance)) { - best = candidate; - bestMaxDistance = maxDistance; - bestTotalDistance = totalDistance; - } - } - - if (!best && emitDiagnostic) { - emitError(loc) << "unsupported non-tree control flow: branch arms do not " - "reach a common convergence block"; - return failure(); - } - if (!best) - return failure(); - - return best; -} - -/// Replaces a destination block's SSA arguments with the values carried by the -/// incoming branch. Callers erase the original CFG blocks after their bodies -/// have been moved or cloned. -static LogicalResult replaceBlockArguments(Block *block, ValueRange incoming, - Location loc) { - if (block->getNumArguments() != incoming.size()) { - emitError(loc) << "invalid branch operand count while structuring " - "control flow: " - << incoming.size() << " operands for " - << block->getNumArguments() << " block arguments"; - return failure(); - } - - for (auto [arg, value] : llvm::zip(block->getArguments(), incoming)) - arg.replaceAllUsesWith(value); - return success(); -} - -/// Moves non-terminator operations into the currently constructed SCF region. -/// The original terminator remains available to drive recursive CFG traversal. -static void moveBlockBodyBefore(Block *block, OpBuilder &builder) { - SmallVector movedOps = llvm::map_to_vector( - block->without_terminator(), [](Operation &op) { return &op; }); - for (Operation *op : movedOps) - op->moveBefore(builder.getInsertionBlock(), builder.getInsertionPoint()); -} - -struct ReturnPathResult { - // A terminal path cannot place a return inside an scf.if region. Bubble the - // return operands outward so the caller can yield them and emit one return - // after the structured conditional. - SmallVector operands; -}; - -//===----------------------------------------------------------------------===// -// Converging branch construction -//===----------------------------------------------------------------------===// - -static FailureOr> buildRegionPath(Block *block, - ValueRange incoming, - Block *stopBlock, - OpBuilder &builder); - -static FailureOr -buildReturnPath(Block *block, ValueRange incoming, OpBuilder &builder); - -static FailureOr buildTerminalValueIf(cf::CondBranchOp condBr, - OpBuilder &builder); - -static FailureOr buildStructuredIf(cf::CondBranchOp condBr, - Block *joinBlock, - OpBuilder &builder) { - // The join block arguments define the exact result contract of the new if. - // Each arm is recursively consumed up to that join and yields its incoming - // values in the same order. - SmallVector resultTypes; - resultTypes.reserve(joinBlock->getNumArguments()); - for (BlockArgument arg : joinBlock->getArguments()) - resultTypes.push_back(arg.getType()); - - auto ifOp = builder.create(condBr.getLoc(), resultTypes, - condBr.getCondition(), - /*withElseRegion=*/true); - - { - OpBuilder::InsertionGuard guard(builder); - Operation *autoYield = - resultTypes.empty() ? ifOp.thenBlock()->getTerminator() : nullptr; - if (autoYield) - builder.setInsertionPoint(autoYield); - else - builder.setInsertionPointToStart(ifOp.thenBlock()); - FailureOr> thenYield = buildRegionPath( - condBr.getTrueDest(), condBr.getTrueDestOperands(), joinBlock, builder); - if (failed(thenYield)) - return failure(); - if (thenYield->size() != resultTypes.size()) { - condBr.emitError("then branch yields ") - << thenYield->size() << " values, expected " << resultTypes.size(); - return failure(); - } - if (!autoYield) - builder.create(condBr.getLoc(), *thenYield); - } - - { - OpBuilder::InsertionGuard guard(builder); - Operation *autoYield = - resultTypes.empty() ? ifOp.elseBlock()->getTerminator() : nullptr; - if (autoYield) - builder.setInsertionPoint(autoYield); - else - builder.setInsertionPointToStart(ifOp.elseBlock()); - FailureOr> elseYield = - buildRegionPath(condBr.getFalseDest(), condBr.getFalseDestOperands(), - joinBlock, builder); - if (failed(elseYield)) - return failure(); - if (elseYield->size() != resultTypes.size()) { - condBr.emitError("else branch yields ") - << elseYield->size() << " values, expected " << resultTypes.size(); - return failure(); - } - if (!autoYield) - builder.create(condBr.getLoc(), *elseYield); - } - - return ifOp; -} - -static Operation *createReturnLike(OpBuilder &builder, Location loc, - Operation *sampleReturn, - ValueRange operands) { - // Preserve whether the containing callable uses tt.return or func.return, - // along with any dialect-specific attributes on that terminator. - OperationState state(loc, sampleReturn->getName()); - state.addOperands(operands); - state.addAttributes(sampleReturn->getAttrs()); - return builder.create(state); -} - -//===----------------------------------------------------------------------===// -// Terminal branch construction -//===----------------------------------------------------------------------===// - -/// Computes the value types returned by a terminal path without mutating the -/// CFG. Both arms of a terminal conditional must return the same signature so -/// they can become the results of a value-producing scf.if. -static FailureOr> -collectReturnPathTypes(Block *block, SmallPtrSetImpl &visiting) { - if (!visiting.insert(block).second) - return block->getTerminator()->emitError() - << "unsupported cyclic terminal control flow"; - - Operation *term = block->getTerminator(); - if (auto br = dyn_cast(term)) { - FailureOr> result = - collectReturnPathTypes(br.getDest(), visiting); - visiting.erase(block); - return result; - } - - if (auto condBr = dyn_cast(term)) { - FailureOr nestedJoin = findNearestCommonBlock( - condBr.getTrueDest(), condBr.getFalseDest(), condBr.getLoc(), - /*emitDiagnostic=*/false); - if (succeeded(nestedJoin)) { - FailureOr> result = - collectReturnPathTypes(*nestedJoin, visiting); - visiting.erase(block); - return result; - } - - FailureOr> thenTypes = - collectReturnPathTypes(condBr.getTrueDest(), visiting); - FailureOr> elseTypes = - collectReturnPathTypes(condBr.getFalseDest(), visiting); - visiting.erase(block); - if (failed(thenTypes) || failed(elseTypes)) - return failure(); - if (!haveSameTypes(TypeRange{*thenTypes}, TypeRange{*elseTypes})) { - condBr.emitError("terminal branch return types do not match"); - return failure(); - } - return *thenTypes; - } - - if (isSupportedReturn(term)) { - SmallVector types; - for (Value operand : term->getOperands()) - types.push_back(operand.getType()); - visiting.erase(block); - return types; - } - - visiting.erase(block); - return term->emitError() - << "unsupported terminator while analyzing terminal control flow"; -} - -static Operation *findReturnOnPath(Block *block, - SmallPtrSetImpl &visited) { - // The operation name/attributes of one reachable return are used as the - // template after terminal paths have been converted to yielded values. - if (!visited.insert(block).second) - return nullptr; - - Operation *term = block->getTerminator(); - if (isSupportedReturn(term)) - return term; - for (Block *successor : getCfgSuccessors(block)) { - if (successor->getParent() != block->getParent()) - continue; - if (Operation *returnOp = findReturnOnPath(successor, visited)) - return returnOp; - } - return nullptr; -} - -static SmallVector mapValues(ValueRange values, IRMapping &mapping) { - // lookupOrDefault is intentional for values captured from outside the cloned - // path; only block arguments and locally cloned results require mappings. - SmallVector mapped; - mapped.reserve(values.size()); - for (Value value : values) - mapped.push_back(mapping.lookupOrDefault(value)); - return mapped; -} - -static FailureOr> -buildClonedTerminalPath(Block *block, ValueRange incoming, OpBuilder &builder, - IRMapping mapping, SmallPtrSetImpl &visiting); - -static FailureOr> -buildClonedTerminalTerminator(Operation *term, OpBuilder &builder, - IRMapping mapping, - SmallPtrSetImpl &visiting) { - if (auto br = dyn_cast(term)) { - SmallVector incoming = mapValues(br.getDestOperands(), mapping); - return buildClonedTerminalPath(br.getDest(), incoming, builder, mapping, - visiting); - } - - if (auto condBr = dyn_cast(term)) { - SmallPtrSet thenVisiting; - FailureOr> thenTypes = - collectReturnPathTypes(condBr.getTrueDest(), thenVisiting); - SmallPtrSet elseVisiting; - FailureOr> elseTypes = - collectReturnPathTypes(condBr.getFalseDest(), elseVisiting); - if (failed(thenTypes) || failed(elseTypes)) - return failure(); - if (!haveSameTypes(TypeRange{*thenTypes}, TypeRange{*elseTypes})) { - condBr.emitError("terminal branch return types do not match"); - return failure(); - } - - auto ifOp = builder.create( - condBr.getLoc(), *thenTypes, - mapping.lookupOrDefault(condBr.getCondition()), - /*withElseRegion=*/true); - - { - OpBuilder::InsertionGuard guard(builder); - Operation *autoYield = - thenTypes->empty() ? ifOp.thenBlock()->getTerminator() : nullptr; - if (autoYield) - builder.setInsertionPoint(autoYield); - else - builder.setInsertionPointToStart(ifOp.thenBlock()); - SmallVector incoming = - mapValues(condBr.getTrueDestOperands(), mapping); - FailureOr> thenReturn = buildClonedTerminalPath( - condBr.getTrueDest(), incoming, builder, mapping, visiting); - if (failed(thenReturn)) - return failure(); - if (!haveSameTypes(TypeRange{ValueRange{*thenReturn}}, - TypeRange{*thenTypes})) { - condBr.emitError("then terminal branch returns incompatible values"); - return failure(); - } - if (!autoYield) - builder.create(condBr.getLoc(), *thenReturn); - } - - { - OpBuilder::InsertionGuard guard(builder); - Operation *autoYield = - thenTypes->empty() ? ifOp.elseBlock()->getTerminator() : nullptr; - if (autoYield) - builder.setInsertionPoint(autoYield); - else - builder.setInsertionPointToStart(ifOp.elseBlock()); - SmallVector incoming = - mapValues(condBr.getFalseDestOperands(), mapping); - FailureOr> elseReturn = buildClonedTerminalPath( - condBr.getFalseDest(), incoming, builder, mapping, visiting); - if (failed(elseReturn)) - return failure(); - if (!haveSameTypes(TypeRange{ValueRange{*elseReturn}}, - TypeRange{*thenTypes})) { - condBr.emitError("else terminal branch returns incompatible values"); - return failure(); - } - if (!autoYield) - builder.create(condBr.getLoc(), *elseReturn); - } - - return SmallVector(ifOp->getResults().begin(), - ifOp->getResults().end()); - } - - if (isSupportedReturn(term)) - return mapValues(term->getOperands(), mapping); - - return term->emitError() - << "unsupported terminator while structuring terminal control flow"; -} - -static FailureOr> -buildClonedTerminalPath(Block *block, ValueRange incoming, OpBuilder &builder, - IRMapping mapping, SmallPtrSetImpl &visiting) { - // Terminal paths are cloned rather than moved. Analysis and construction can - // therefore fail without partially consuming the original CFG; the old - // blocks are erased only after the complete replacement return is built. - if (!visiting.insert(block).second) - return block->getTerminator()->emitError() - << "unsupported cyclic terminal control flow"; - - if (block->getNumArguments() != incoming.size()) { - visiting.erase(block); - return block->getTerminator()->emitError() - << "invalid branch operand count while structuring terminal " - "control flow"; - } - - for (auto [arg, value] : llvm::zip(block->getArguments(), incoming)) - mapping.map(arg, value); - - for (Operation &op : block->without_terminator()) - builder.clone(op, mapping); - - FailureOr> result = buildClonedTerminalTerminator( - block->getTerminator(), builder, mapping, visiting); - visiting.erase(block); - return result; -} - -static bool hasNonTreeCondBranch(Region &body) { - // A branch without a join cannot be consumed by the move-based path. If any - // such branch exists, clone the complete terminal tree atomically instead. - for (Block &block : body) { - auto condBr = dyn_cast(block.getTerminator()); - if (!condBr) - continue; - if (failed(findNearestCommonBlock(condBr.getTrueDest(), - condBr.getFalseDest(), condBr.getLoc(), - /*emitDiagnostic=*/false))) - return true; - } - return false; -} - -static LogicalResult structureTerminalReturnBody(Operation *funcOp, - Region &body) { - // This path is used when branch arms do not reconverge but both end in - // compatible returns. The cloned scf.if produces the return operands, after - // which every original non-entry block can be removed together. - Block &entryBlock = body.front(); - Operation *entryTerm = entryBlock.getTerminator(); - SmallPtrSet visited; - Operation *sampleReturn = findReturnOnPath(&entryBlock, visited); - if (!sampleReturn) { - return funcOp->emitError() - << "unsupported non-tree control flow: no terminal return found"; - } - - OpBuilder builder(entryTerm); - IRMapping mapping; - SmallPtrSet visiting; - FailureOr> returnOperands = - buildClonedTerminalTerminator(entryTerm, builder, mapping, visiting); - if (failed(returnOperands)) - return failure(); - - createReturnLike(builder, entryTerm->getLoc(), sampleReturn, *returnOperands); - - SmallVector eraseBlocks; - for (Block &block : llvm::drop_begin(body.getBlocks())) - eraseBlocks.push_back(&block); - - entryTerm->erase(); - for (Block *block : eraseBlocks) { - for (Operation &op : *block) - op.dropAllReferences(); - } - for (Block *block : llvm::reverse(eraseBlocks)) - block->erase(); - - return success(); -} - -static FailureOr buildTerminalValueIf(cf::CondBranchOp condBr, - OpBuilder &builder) { - SmallPtrSet thenVisiting; - FailureOr> thenTypes = - collectReturnPathTypes(condBr.getTrueDest(), thenVisiting); - SmallPtrSet elseVisiting; - FailureOr> elseTypes = - collectReturnPathTypes(condBr.getFalseDest(), elseVisiting); - if (failed(thenTypes) || failed(elseTypes)) - return failure(); - if (!haveSameTypes(TypeRange{*thenTypes}, TypeRange{*elseTypes})) { - condBr.emitError("terminal branch return types do not match"); - return failure(); - } - - auto ifOp = builder.create(condBr.getLoc(), *thenTypes, - condBr.getCondition(), - /*withElseRegion=*/true); - - { - OpBuilder::InsertionGuard branchGuard(builder); - Operation *autoYield = - thenTypes->empty() ? ifOp.thenBlock()->getTerminator() : nullptr; - if (autoYield) - builder.setInsertionPoint(autoYield); - else - builder.setInsertionPointToStart(ifOp.thenBlock()); - FailureOr thenReturn = buildReturnPath( - condBr.getTrueDest(), condBr.getTrueDestOperands(), builder); - if (failed(thenReturn)) - return failure(); - if (!haveSameTypes(TypeRange{ValueRange{thenReturn->operands}}, - TypeRange{*thenTypes})) { - condBr.emitError("then terminal branch returns incompatible values"); - return failure(); - } - if (!autoYield) - builder.create(condBr.getLoc(), thenReturn->operands); - } - - { - OpBuilder::InsertionGuard branchGuard(builder); - Operation *autoYield = - thenTypes->empty() ? ifOp.elseBlock()->getTerminator() : nullptr; - if (autoYield) - builder.setInsertionPoint(autoYield); - else - builder.setInsertionPointToStart(ifOp.elseBlock()); - FailureOr elseReturn = buildReturnPath( - condBr.getFalseDest(), condBr.getFalseDestOperands(), builder); - if (failed(elseReturn)) - return failure(); - if (!haveSameTypes(TypeRange{ValueRange{elseReturn->operands}}, - TypeRange{*thenTypes})) { - condBr.emitError("else terminal branch returns incompatible values"); - return failure(); - } - if (!autoYield) - builder.create(condBr.getLoc(), elseReturn->operands); - } - - return ifOp; -} - -static FailureOr -buildReturnPath(Block *block, ValueRange incoming, OpBuilder &builder) { - Operation *term = block->getTerminator(); - if (failed(replaceBlockArguments(block, incoming, term->getLoc()))) - return failure(); - moveBlockBodyBefore(block, builder); - - if (auto br = dyn_cast(term)) - return buildReturnPath(br.getDest(), br.getDestOperands(), builder); - - if (auto condBr = dyn_cast(term)) { - FailureOr nestedJoin = findNearestCommonBlock( - condBr.getTrueDest(), condBr.getFalseDest(), condBr.getLoc(), - /*emitDiagnostic=*/false); - if (succeeded(nestedJoin)) { - FailureOr nestedIf = - buildStructuredIf(condBr, *nestedJoin, builder); - if (failed(nestedIf)) - return failure(); - - SmallVector nestedResults((*nestedIf)->getResults().begin(), - (*nestedIf)->getResults().end()); - return buildReturnPath(*nestedJoin, nestedResults, builder); - } - - FailureOr terminalIf = buildTerminalValueIf(condBr, builder); - if (failed(terminalIf)) { - condBr.emitError() << "unsupported non-tree control flow: branch arms do " - "not both terminate with compatible returns"; - return failure(); - } - - ReturnPathResult result; - result.operands.assign((*terminalIf)->getResults().begin(), - (*terminalIf)->getResults().end()); - return result; - } - - if (isSupportedReturn(term)) { - ReturnPathResult result; - result.operands.assign(term->getOperands().begin(), - term->getOperands().end()); - return result; - } - - return term->emitError() - << "unsupported terminator while structuring terminal control flow"; -} - -//===----------------------------------------------------------------------===// -// Top-level CFG consumption -//===----------------------------------------------------------------------===// - -/// Consumes one branch arm until `stopBlock`. Nested conditionals are converted -/// recursively, and their results become the incoming values of the next join. -static FailureOr> buildRegionPath(Block *block, - ValueRange incoming, - Block *stopBlock, - OpBuilder &builder) { - if (block == stopBlock) - return SmallVector(incoming.begin(), incoming.end()); - - Operation *term = block->getTerminator(); - if (failed(replaceBlockArguments(block, incoming, term->getLoc()))) - return failure(); - moveBlockBodyBefore(block, builder); - - if (auto br = dyn_cast(term)) { - SmallVector operands(br.getDestOperands().begin(), - br.getDestOperands().end()); - if (br.getDest() == stopBlock) - return operands; - return buildRegionPath(br.getDest(), operands, stopBlock, builder); - } - - if (auto condBr = dyn_cast(term)) { - FailureOr nestedJoin = findNearestCommonBlock( - condBr.getTrueDest(), condBr.getFalseDest(), condBr.getLoc()); - if (failed(nestedJoin)) - return failure(); - - FailureOr nestedIf = - buildStructuredIf(condBr, *nestedJoin, builder); - if (failed(nestedIf)) - return failure(); - - SmallVector nestedResults((*nestedIf)->getResults().begin(), - (*nestedIf)->getResults().end()); - if (*nestedJoin == stopBlock) - return nestedResults; - return buildRegionPath(*nestedJoin, nestedResults, stopBlock, builder); - } - - if (isSupportedReturn(term)) { - return term->emitError() - << "unsupported early return while structuring control flow"; - } - - return term->emitError() - << "unsupported terminator while structuring control flow"; -} - -static LogicalResult appendStructuredBlock(Block *block, ValueRange incoming, - OpBuilder &builder, - Operation *anchorTerminator); - -static LogicalResult appendStructuredTerminator(Operation *term, - OpBuilder &builder, - Operation *anchorTerminator) { - // Walk the entry CFG in execution order. A converging conditional is emitted - // as scf.if and traversal resumes at its join; a terminal conditional emits - // the final return and finishes the function body. - if (auto br = dyn_cast(term)) { - return appendStructuredBlock(br.getDest(), br.getDestOperands(), builder, - anchorTerminator); - } - - if (auto condBr = dyn_cast(term)) { - FailureOr joinBlock = findNearestCommonBlock( - condBr.getTrueDest(), condBr.getFalseDest(), condBr.getLoc(), - /*emitDiagnostic=*/false); - if (failed(joinBlock)) { - SmallPtrSet visited; - Operation *sampleReturn = findReturnOnPath(condBr.getTrueDest(), visited); - if (!sampleReturn) { - visited.clear(); - sampleReturn = findReturnOnPath(condBr.getFalseDest(), visited); - } - if (!sampleReturn) { - return condBr.emitError() - << "unsupported non-tree control flow: branch arms do not " - "reach a common convergence block"; - } - - FailureOr terminalIf = buildTerminalValueIf(condBr, builder); - if (failed(terminalIf)) - return failure(); - - SmallVector returnOperands((*terminalIf)->getResults().begin(), - (*terminalIf)->getResults().end()); - createReturnLike(builder, condBr.getLoc(), sampleReturn, returnOperands); - return success(); - } - - FailureOr ifOp = buildStructuredIf(condBr, *joinBlock, builder); - if (failed(ifOp)) - return failure(); - - return appendStructuredBlock(*joinBlock, (*ifOp)->getResults(), builder, - anchorTerminator); - } - - if (isSupportedReturn(term)) { - term->moveBefore(anchorTerminator); - return success(); - } - - return term->emitError() - << "unsupported entry terminator while structuring control flow"; -} - -static LogicalResult appendStructuredBlock(Block *block, ValueRange incoming, - OpBuilder &builder, - Operation *anchorTerminator) { - Operation *term = block->getTerminator(); - if (failed(replaceBlockArguments(block, incoming, term->getLoc()))) - return failure(); - - moveBlockBodyBefore(block, builder); - return appendStructuredTerminator(term, builder, anchorTerminator); -} - -static LogicalResult validateSupportedCfg(Region &body) { - // Validate every reachable block before move-based construction starts, so - // unsupported terminators cannot leave a partially consumed function body. - for (Block &block : body) { - Operation *term = block.getTerminator(); - if (!isa(term) && !isSupportedReturn(term)) - return term->emitError() - << "unsupported terminator in multi-block function"; - } - return success(); -} - -/// Collects blocks reachable from the function entry through supported CFG -/// successors. Restrict successor traversal to the current region because a -/// branch-like operation must not make a nested or enclosing block reachable. -static void collectReachableBlocks(Block *block, - SmallPtrSetImpl &reachable) { - if (!reachable.insert(block).second) - return; - - for (Block *successor : getCfgSuccessors(block)) { - if (successor->getParent() == block->getParent()) - collectReachableBlocks(successor, reachable); - } -} - -/// Removes blocks that cannot be reached from the entry block before CFG -/// validation and structuring. Frontend lowering may leave detached return -/// blocks behind; they are not part of the function's executable CFG and must -/// not affect join discovery, cycle checks or the entry-terminator decision. -static void eraseUnreachableBlocks(Region &body) { - if (body.empty() || body.hasOneBlock()) - return; - - SmallPtrSet reachable; - collectReachableBlocks(&body.front(), reachable); - - SmallVector eraseBlocks; - for (Block &block : body) { - if (!reachable.contains(&block)) - eraseBlocks.push_back(&block); - } - - // Break successor and operand references before erasing. This also handles - // unreachable blocks that refer to one another, including unreachable - // cycles. - for (Block *block : eraseBlocks) { - for (Operation &op : *block) - op.dropAllReferences(); - } - for (Block *block : llvm::reverse(eraseBlocks)) - block->erase(); -} - -/// Rejects backedges before any destructive block movement occurs. General -/// restructuring of cyclic CFGs is not supported. -static LogicalResult rejectCyclicCfg(Block *block, - SmallPtrSetImpl &visiting, - SmallPtrSetImpl &visited) { - if (visited.contains(block)) - return success(); - if (!visiting.insert(block).second) - return block->getTerminator()->emitError() - << "unsupported cyclic control flow in multi-block function"; - - for (Block *successor : getCfgSuccessors(block)) { - if (successor->getParent() == block->getParent() && - failed(rejectCyclicCfg(successor, visiting, visited))) - return failure(); - } - - visiting.erase(block); - visited.insert(block); - return success(); -} - -static LogicalResult structureFunctionBody(Operation *funcOp, Region &body) { - // Validation is deliberately completed before the move-based path starts. - // From that point onward the function is rewritten as one SCF entry block and - // the consumed CFG blocks are erased only after construction succeeds. - if (body.empty() || body.hasOneBlock()) - return success(); - - eraseUnreachableBlocks(body); - if (body.hasOneBlock()) - return success(); - - if (failed(validateSupportedCfg(body))) - return failure(); - - SmallPtrSet visiting; - SmallPtrSet visited; - if (failed(rejectCyclicCfg(&body.front(), visiting, visited))) - return failure(); - - if (hasNonTreeCondBranch(body)) - return structureTerminalReturnBody(funcOp, body); - - Block &entryBlock = body.front(); - Operation *entryTerm = entryBlock.getTerminator(); - if (isSupportedReturn(entryTerm)) { - return funcOp->emitError() - << "multi-block function entry cannot terminate with return"; - } - - SmallVector eraseBlocks; - for (Block &block : llvm::drop_begin(body.getBlocks())) - eraseBlocks.push_back(&block); - - OpBuilder builder(entryTerm); - if (failed(appendStructuredTerminator(entryTerm, builder, entryTerm))) - return failure(); - - entryTerm->erase(); - for (Block *block : eraseBlocks) { - for (Operation &op : *block) - op.dropAllReferences(); - } - for (Block *block : llvm::reverse(eraseBlocks)) - block->erase(); - - return success(); -} - -} // namespace - -namespace mlir::triton::controlflow { - -LogicalResult structureCFG(ModuleOp module) { - // Collect functions first because structureFunctionBody mutates their nested - // regions. This keeps the module walk independent of those mutations. - SmallVector functions; - module.walk([&](Operation *op) { - if (isa(op)) - functions.push_back(op); - }); - - // Handle both Triton callables and ordinary func.func wrappers; declarations - // and functions erased by an enclosing transformation are skipped. - for (Operation *op : functions) { - if (!op->getParentOp()) - continue; - - if (auto funcOp = dyn_cast(op)) { - if (!funcOp.isDeclaration() && - failed(structureFunctionBody(funcOp, funcOp.getBody()))) - return failure(); - continue; - } - - if (auto funcOp = dyn_cast(op)) { - if (!funcOp.isDeclaration() && - failed(structureFunctionBody(funcOp, funcOp.getBody()))) - return failure(); - continue; - } - - auto mapOp = cast(op); - if (failed(structureFunctionBody(mapOp, mapOp.getRegion()))) - return failure(); - } - - return success(); -} - -} // namespace mlir::triton::controlflow diff --git a/third_party/ascend/lib/TritonControlFlowOpt/CMakeLists.txt b/third_party/ascend/lib/TritonControlFlowOpt/CMakeLists.txt index 190680c22a..1c76f4b913 100644 --- a/third_party/ascend/lib/TritonControlFlowOpt/CMakeLists.txt +++ b/third_party/ascend/lib/TritonControlFlowOpt/CMakeLists.txt @@ -1,9 +1,4 @@ add_triton_library(TritonControlFlowOpt - BlockPtrDecompose.cpp - CFGStructuring.cpp - ControlFlowAnalysis.cpp - ControlFlowRewrite.cpp - TensorPtrDecompose.cpp TritonControlFlowOptPass.cpp DEPENDS @@ -16,8 +11,6 @@ add_triton_library(TritonControlFlowOpt MLIRIR MLIRPass MLIRSCFDialect - MLIRSideEffectInterfaces MLIRSupport - MLIRTritonNPUUtils TritonIR ) diff --git a/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowAnalysis.cpp b/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowAnalysis.cpp deleted file mode 100644 index 67ad7c7c31..0000000000 --- a/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowAnalysis.cpp +++ /dev/null @@ -1,509 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#include "TritonControlFlowOpt/ControlFlowAnalysis.h" - -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/BuiltinOps.h" - -#include "llvm/ADT/STLExtras.h" - -#include - -using namespace mlir; - -namespace mlir::triton::controlflow { - -namespace { - -/// Control-flow kinds whose operand/result correspondence is understood by -/// both the analyzer and the mechanical rewrite. -static bool isSupportedControlFlow(Operation *op) { - return isa(op); -} - -static void setResultIdentity(AnalyzedValue &value, Value result, - ArrayRef componentIndices) { - // A transferred component is represented by the SCF result outside the op; - // invariant components retain their incoming symbolic identities. - for (unsigned index : componentIndices) - value.components[index].identity = - ComponentIdentity::fromValue(result, index); -} - -static ControlFlowSlotAnalysis -makeSlotAnalysis(unsigned oldIndex, unsigned componentCount, - ArrayRef componentIndices, - SmallVector componentTypes) { - // Preserve the full classification for future Recomputed support while also - // storing a compact ordered list used by the current signature rewrite. - ControlFlowSlotAnalysis slot; - slot.oldIndex = oldIndex; - slot.componentKinds.assign(componentCount, ComponentTransferKind::Invariant); - for (unsigned index : componentIndices) - slot.componentKinds[index] = ComponentTransferKind::Transferred; - slot.componentIndices.append(componentIndices.begin(), - componentIndices.end()); - slot.componentTypes = std::move(componentTypes); - return slot; -} - -} // namespace - -//===----------------------------------------------------------------------===// -// Shared value analysis and nested traversal -//===----------------------------------------------------------------------===// - -const ControlFlowOpAnalysis * -ControlFlowRewritePlan::lookup(Operation *op) const { - auto it = operations.find(op); - return it == operations.end() ? nullptr : &it->second; -} - -const AnalyzedValue * -ControlFlowAnalysisContext::lookupValue(Value value) const { - auto it = analyzedValues.find(value); - return it == analyzedValues.end() ? nullptr : &it->second; -} - -const ControlFlowOpAnalysis * -ControlFlowAnalysisContext::lookup(Operation *op) const { - auto it = analyzedOps.find(op); - return it == analyzedOps.end() ? nullptr : &it->second; -} - -FailureOr ControlFlowAnalysisContext::analyzeValue(Value value) { - if (const AnalyzedValue *known = lookupValue(value)) - return *known; - - // A control-flow result is meaningful only after all incoming states have - // been merged. Analyze its owner first instead of letting a pointer policy - // treat the opaque result as a new base. - if (auto result = dyn_cast(value)) { - Operation *owner = result.getOwner(); - if (isSupportedControlFlow(owner)) { - if (failed(analyzeControlFlowOp(owner))) - return failure(); - if (const AnalyzedValue *known = lookupValue(value)) - return *known; - return failure(); - } - } - - FailureOr result = policy.analyzeValue(value, *this); - if (failed(result)) - return failure(); - analyzedValues.try_emplace(value, *result); - return *result; -} - -void ControlFlowAnalysisContext::bindRegionArgument( - Value argument, const AnalyzedValue &initial, - ArrayRef componentIndices) { - // Only candidate loop components acquire a new identity at region entry. - // Policy-owned invariants and non-carried components remain traceable to the - // initial descriptor and can therefore be checked at the backedge. - AnalyzedValue argumentState = initial; - for (unsigned index : componentIndices) - argumentState.components[index].identity = - ComponentIdentity::fromValue(argument, index); - analyzedValues[argument] = std::move(argumentState); -} - -FailureOr> ControlFlowAnalysisContext::getTransferredTypes( - const AnalyzedValue &lhs, const AnalyzedValue &rhs, - ArrayRef componentIndices) const { - // Type joining is policy-owned because tensor offsets may widen i32 to i64, - // whereas block-pointer descriptor fields currently require exact equality. - SmallVector types; - types.reserve(componentIndices.size()); - for (unsigned index : componentIndices) { - if (index >= lhs.components.size() || index >= rhs.components.size()) - return failure(); - FailureOr type = policy.joinComponentTypes( - lhs.components[index].type, rhs.components[index].type); - if (failed(type)) - return failure(); - types.push_back(*type); - } - return types; -} - -LogicalResult -ControlFlowAnalysisContext::analyzeNestedOperations(Block *block, - bool &hasNestedRewrite) { - // Walk in program order so an inner result requested by a later address - // expression is already cached. Nested SCF is analyzed recursively and its - // rewrite requirement is propagated to every enclosing supported SCF op. - for (Operation &operation : block->without_terminator()) { - if (isSupportedControlFlow(&operation)) { - FailureOr nested = - analyzeControlFlowOp(&operation); - if (failed(nested)) - return failure(); - hasNestedRewrite |= nested->needsRewrite(); - continue; - } - - // SCF may be wrapped in an ordinary region-owning operation. Such regions - // do not change pointer schemas, but ControlFlowRewrite does not yet clone - // arbitrary region operations recursively. Reject an affected nested SCF - // instead of reporting success and leaving it opaque in the rewritten IR. - for (Region ®ion : operation.getRegions()) { - for (Block &nestedBlock : region) { - bool regionNeedsRewrite = false; - if (failed(analyzeNestedOperations(&nestedBlock, regionNeedsRewrite)) || - regionNeedsRewrite) - return failure(); - } - } - } - return success(); -} - -//===----------------------------------------------------------------------===// -// scf.for schema analysis -//===----------------------------------------------------------------------===// - -FailureOr -ControlFlowAnalysisContext::analyzeFor(Operation *operation) { - auto forOp = cast(operation); - auto yieldOp = cast(forOp.getBody()->getTerminator()); - if (forOp.getInitArgs().size() != forOp.getRegionIterArgs().size() || - yieldOp.getNumOperands() != forOp.getRegionIterArgs().size()) - return failure(); - - SmallVector> initialStates( - forOp.getInitArgs().size()); - - // Bind an abstract component state to each pointer region argument. The body - // can then be analyzed recursively without constructing a replacement loop. - for (auto [index, iterArg] : llvm::enumerate(forOp.getRegionIterArgs())) { - if (!policy.isDecompositionTarget(iterArg)) - continue; - FailureOr initial = analyzeValue(forOp.getInitArgs()[index]); - if (failed(initial)) - return failure(); - FailureOr> candidates = - policy.getLoopCandidateComponents(*initial); - if (failed(candidates)) - return failure(); - initialStates[index] = *initial; - bindRegionArgument(iterArg, *initial, *candidates); - } - - ControlFlowOpAnalysis result; - if (failed(analyzeNestedOperations(forOp.getBody(), result.hasNestedRewrite))) - return failure(); - - // Compare the abstract init, region-argument and yielded states only after - // the complete body (including nested SCF) has been analyzed. - for (auto [index, iterArg] : llvm::enumerate(forOp.getRegionIterArgs())) { - if (!initialStates[index]) - continue; - FailureOr next = analyzeValue(yieldOp.getOperand(index)); - const AnalyzedValue *argument = lookupValue(iterArg); - if (failed(next) || !argument) - return failure(); - FailureOr> transferred = - policy.getLoopTransferredComponents(*initialStates[index], *argument, - *next); - if (failed(transferred)) - return failure(); - - AnalyzedValue resultState = *initialStates[index]; - if (!transferred->empty()) { - // The ordered component/type list is the complete contract consumed by - // rewriteForOp; rewrite does not rediscover its signature on the fly. - FailureOr> types = - getTransferredTypes(*initialStates[index], *next, *transferred); - if (failed(types)) - return failure(); - for (auto [component, type] : llvm::zip(*transferred, *types)) - resultState.components[component].type = type; - result.slots.push_back(makeSlotAnalysis(index, - resultState.components.size(), - *transferred, std::move(*types))); - setResultIdentity(resultState, forOp.getResult(index), *transferred); - } - analyzedValues[forOp.getResult(index)] = std::move(resultState); - } - - return result; -} - -//===----------------------------------------------------------------------===// -// scf.while schema analysis -//===----------------------------------------------------------------------===// - -FailureOr -ControlFlowAnalysisContext::analyzeWhile(Operation *operation) { - auto whileOp = cast(operation); - scf::ConditionOp conditionOp = whileOp.getConditionOp(); - scf::YieldOp yieldOp = whileOp.getYieldOp(); - if (whileOp.getInits().size() != whileOp.getBeforeArguments().size() || - conditionOp.getArgs().size() != whileOp.getAfterArguments().size() || - yieldOp.getNumOperands() != whileOp.getBeforeArguments().size()) - return failure(); - - SmallVector> initialStates( - whileOp.getInits().size()); - SmallVector> candidateIndices( - whileOp.getInits().size()); - - // Bind the initial descriptor to the before-region arguments first. The - // condition and backedge are analyzed as two consecutive state transfers. - for (auto [index, beforeArg] : - llvm::enumerate(whileOp.getBeforeArguments())) { - if (!policy.isDecompositionTarget(beforeArg)) - continue; - FailureOr initial = analyzeValue(whileOp.getInits()[index]); - if (failed(initial)) - return failure(); - FailureOr> candidates = - policy.getLoopCandidateComponents(*initial); - if (failed(candidates)) - return failure(); - initialStates[index] = *initial; - candidateIndices[index] = *candidates; - bindRegionArgument(beforeArg, *initial, *candidates); - } - - ControlFlowOpAnalysis result; - if (failed(analyzeNestedOperations(whileOp.getBeforeBody(), - result.hasNestedRewrite))) - return failure(); - - // The after-region argument receives the value forwarded by scf.condition. - // Bind it before visiting the after region so address expressions there can - // recursively resolve the same abstract schema. - SmallVector> conditionStates( - conditionOp.getArgs().size()); - for (auto [index, afterArg] : llvm::enumerate(whileOp.getAfterArguments())) { - if (!initialStates[index]) - continue; - FailureOr condition = - analyzeValue(conditionOp.getArgs()[index]); - if (failed(condition)) - return failure(); - conditionStates[index] = *condition; - bindRegionArgument(afterArg, *condition, candidateIndices[index]); - } - - if (failed(analyzeNestedOperations(whileOp.getAfterBody(), - result.hasNestedRewrite))) - return failure(); - - for (auto [index, beforeArg] : - llvm::enumerate(whileOp.getBeforeArguments())) { - if (!initialStates[index]) - continue; - FailureOr next = analyzeValue(yieldOp.getOperand(index)); - const AnalyzedValue *argument = lookupValue(beforeArg); - if (failed(next) || !argument || !conditionStates[index]) - return failure(); - - // The policy merge is monotone over {Invariant, Transferred}: once either - // region changes a candidate component it must be present in the loop - // signature. Re-running the union would not remove a transferred bit, so - // this is the fixed point for the current two-state domain. - FailureOr> fromCondition = - policy.getLoopTransferredComponents(*initialStates[index], *argument, - *conditionStates[index]); - const AnalyzedValue *afterArgument = - lookupValue(whileOp.getAfterArguments()[index]); - if (!afterArgument || failed(fromCondition)) - return failure(); - FailureOr> fromBackedge = - policy.getLoopTransferredComponents(*conditionStates[index], - *afterArgument, *next); - if (failed(fromBackedge)) - return failure(); - - SmallVector transferred = *fromCondition; - for (unsigned component : *fromBackedge) { - if (!llvm::is_contained(transferred, component)) - transferred.push_back(component); - } - llvm::sort(transferred); - - AnalyzedValue resultState = *initialStates[index]; - if (!transferred.empty()) { - FailureOr> types = getTransferredTypes( - *initialStates[index], *conditionStates[index], transferred); - if (failed(types)) - return failure(); - for (auto [component, type] : llvm::zip(transferred, *types)) { - FailureOr joined = - policy.joinComponentTypes(type, next->components[component].type); - if (failed(joined)) - return failure(); - type = *joined; - resultState.components[component].type = type; - } - result.slots.push_back(makeSlotAnalysis(index, - resultState.components.size(), - transferred, std::move(*types))); - setResultIdentity(resultState, whileOp.getResult(index), transferred); - } - analyzedValues[whileOp.getResult(index)] = std::move(resultState); - } - - return result; -} - -//===----------------------------------------------------------------------===// -// scf.if schema analysis -//===----------------------------------------------------------------------===// - -FailureOr -ControlFlowAnalysisContext::analyzeIf(Operation *operation) { - auto ifOp = cast(operation); - ControlFlowOpAnalysis result; - - if (failed( - analyzeNestedOperations(ifOp.thenBlock(), result.hasNestedRewrite))) - return failure(); - if (ifOp.elseBlock() && failed(analyzeNestedOperations( - ifOp.elseBlock(), result.hasNestedRewrite))) - return failure(); - - if (ifOp.getNumResults() == 0) - // A result-less if may still need rebuilding solely to rewrite nested SCF. - return result; - if (!ifOp.elseBlock()) - return failure(); - - scf::YieldOp thenYield = ifOp.thenYield(); - scf::YieldOp elseYield = ifOp.elseYield(); - if (thenYield.getNumOperands() != ifOp.getNumResults() || - elseYield.getNumOperands() != ifOp.getNumResults()) - return failure(); - - // Each result position is independent. Only policy-matching pointer results - // are expanded; all other result positions keep their original type/order. - for (auto [index, opResult] : llvm::enumerate(ifOp.getResults())) { - if (!policy.isDecompositionTarget(opResult)) - continue; - FailureOr thenState = - analyzeValue(thenYield.getOperand(index)); - FailureOr elseState = - analyzeValue(elseYield.getOperand(index)); - if (failed(thenState) || failed(elseState)) - return failure(); - FailureOr> transferred = - policy.getIfTransferredComponents(*thenState, *elseState); - if (failed(transferred)) - return failure(); - - AnalyzedValue resultState = *thenState; - if (!transferred->empty()) { - FailureOr> types = - getTransferredTypes(*thenState, *elseState, *transferred); - if (failed(types)) - return failure(); - for (auto [component, type] : llvm::zip(*transferred, *types)) - resultState.components[component].type = type; - result.slots.push_back(makeSlotAnalysis(index, - resultState.components.size(), - *transferred, std::move(*types))); - setResultIdentity(resultState, opResult, *transferred); - } - analyzedValues[opResult] = std::move(resultState); - } - - return result; -} - -//===----------------------------------------------------------------------===// -// Stage-wide caching, plan freezing and entry-point discovery -//===----------------------------------------------------------------------===// - -FailureOr -ControlFlowAnalysisContext::analyzeControlFlowOp(Operation *op) { - if (const ControlFlowOpAnalysis *known = lookup(op)) - return *known; - // The in-progress set prevents accidental cyclic re-entry when value - // analysis asks to analyze the control-flow op that owns that value. - if (!isSupportedControlFlow(op) || !operationsBeingAnalyzed.insert(op).second) - return failure(); - - FailureOr result = failure(); - if (isa(op)) - result = analyzeFor(op); - else if (isa(op)) - result = analyzeWhile(op); - else if (isa(op)) - result = analyzeIf(op); - - operationsBeingAnalyzed.erase(op); - if (failed(result)) - return failure(); - analyzedOps.try_emplace(op, *result); - return *result; -} - -ControlFlowRewritePlan ControlFlowAnalysisContext::takeRewritePlan() && { - // analyzedValues intentionally dies with this context. Rewrite only needs - // the position/type decisions and must not retain Value handles that may be - // invalidated as earlier roots are replaced. - return ControlFlowRewritePlan{std::move(analyzedOps)}; -} - -SmallVector collectOutermostControlFlowOps(ModuleOp module) { - // Rewriting an outer root recursively replaces all affected descendants. - // Returning nested ops as independent roots would leave stale pointers after - // the parent is erased, so filter them here instead of relying on walk order. - SmallVector roots; - module.walk([&](Operation *operation) { - if (!isSupportedControlFlow(operation)) - return; - for (Operation *parent = operation->getParentOp(); parent; - parent = parent->getParentOp()) { - if (isSupportedControlFlow(parent)) - return; - } - roots.push_back(operation); - }); - return roots; -} - -FailureOr -analyzeControlFlow(ModuleOp module, const ControlFlowAnalysisPolicy &policy) { - SmallVector roots = collectOutermostControlFlowOps(module); - ControlFlowAnalysisContext context(policy); - - // One context covers the complete decomposition stage. Consequently a - // later root can reuse the merged state of a preceding sibling control-flow - // op, and no IR is mutated until every root has proved analyzable. - for (Operation *root : roots) { - if (failed(context.analyzeControlFlowOp(root))) { - root->emitError( - "failed to analyze pointer components across control flow"); - return failure(); - } - } - - return std::move(context).takeRewritePlan(); -} - -} // namespace mlir::triton::controlflow diff --git a/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowRewrite.cpp b/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowRewrite.cpp deleted file mode 100644 index 0b92e8d72c..0000000000 --- a/third_party/ascend/lib/TritonControlFlowOpt/ControlFlowRewrite.cpp +++ /dev/null @@ -1,1051 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#include "TritonControlFlowOpt/ControlFlowRewrite.h" -#include "Utils/Utils.h" - -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Interfaces/SideEffectInterfaces.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" - -#include - -using namespace mlir; -using mlir::triton::controlflow::ControlFlowOpAnalysis; -using mlir::triton::controlflow::ControlFlowRewriteContext; -using mlir::triton::controlflow::ControlFlowRewritePlan; -using mlir::triton::controlflow::ControlFlowRewritePolicy; -using mlir::triton::controlflow::ControlFlowSlotAnalysis; -using mlir::triton::controlflow::DecomposedValue; - -namespace mlir::triton::controlflow { - -Value ControlFlowRewriteContext::remap(Value value) const { - if (Value mapped = valueMapping.lookupOrNull(value)) - return mapped; - return value; -} - -const DecomposedValue *ControlFlowRewriteContext::lookup(Value value) const { - auto it = decomposedValues.find(value); - return it == decomposedValues.end() ? nullptr : &it->second; -} - -} // namespace mlir::triton::controlflow - -namespace { - -// Keep the mechanical if/for/while rewrite in one translation unit. These -// handlers are mutually recursive through rewriteBodyOps(), share one -// short-lived RewriteEnv, and must agree on signature expansion, nested-op -// ordering and failure cleanup. Splitting them by op kind would expose those -// private invariants through additional internal headers without creating an -// independently reusable component. -// -//===----------------------------------------------------------------------===// -// Per-rewrite state and generic component manipulation -//===----------------------------------------------------------------------===// - -/// Carries everything needed to translate values from one original -/// control-flow path into its replacement path. -/// -/// valueMapping answers which new SSA value replaces an old SSA value. -/// decomposedValues remembers the policy-specific pieces of an old pointer. -/// policy defines what those pieces mean and how to rebuild the pointer, -/// while plan says which loop/if operands must be expanded into such pieces. -/// -/// For example, after a block-pointer loop argument is expanded into scalar -/// offsets, a region-local environment can contain: -/// valueMapping: %old_ptr_arg -> %rebuilt_ptr -/// decomposedValues: %old_ptr_arg -> { -/// components = [%shape0, %stride0, %new_offset0], -/// invariants = [%base], attributes = [order] -/// } -/// Operations cloned into that region use the mapping, while pointer -/// decomposition uses the stored components. -struct RewriteEnv { - /// Starts a rewrite path with no old-to-new mappings or decomposed values. - /// The referenced policy and analysis plan are shared by child environments; - /// only the two mutable state tables above are copied for each nested region. - /// - /// For example, the top-level environment may enter a rewritten scf.for. - /// The loop body copies it, then adds mappings from the old body arguments to - /// the new body arguments without changing the state of sibling regions: - /// RewriteEnv env(blockPtrPolicy, rewritePlan); - /// RewriteEnv bodyEnv = env; - /// bodyEnv.valueMapping.map(oldBodyArg, newBodyArg); - RewriteEnv(const ControlFlowRewritePolicy &policy, - const ControlFlowRewritePlan &plan) - : policy(policy), plan(plan) {} - - /// Gives a decomposition policy read-only access to this path's two state - /// tables. The context can resolve an old SSA value to its replacement and - /// retrieve a previously stored DecomposedValue; it does not own or copy - /// either table and therefore must not outlive this environment. - /// - /// For example, while the block-pointer policy decomposes this operation, - /// it needs both the replacement delta and the saved state of the input: - /// %next = tt.advance %old_ptr, [%old_delta] - /// - /// valueMapping: %old_delta -> %new_delta - /// decomposedValues: %old_ptr -> ptrInfo - /// - /// auto context = bodyEnv.getRewriteContext(); - /// context.remap(oldDelta); // The Value for %new_delta. - /// context.lookup(oldPtr); // A pointer to ptrInfo. - ControlFlowRewriteContext getRewriteContext() const { - return ControlFlowRewriteContext(valueMapping, decomposedValues); - } - - /// Translates one SSA value referenced by the original IR into the value that - /// must be used in the replacement IR. It queries valueMapping; a value - /// defined outside the rewritten area is already valid and is returned - /// unchanged when the table has no entry. - /// - /// This prevents a newly built operation from referring back to a block - /// argument or result owned by the old region. For example: - /// valueMapping: %old_iter_arg -> %new_iter_arg - /// - /// // Original: scf.yield %old_iter_arg, %outer_value - /// newYieldOperands = { - /// bodyEnv.remap(oldIterArg), // The Value for %new_iter_arg. - /// bodyEnv.remap(outerValue) // The unchanged %outer_value. - /// }; - /// The returned values are then used as operands of the replacement yield. - Value remap(Value value) const { return getRewriteContext().remap(value); } - - /// Asks the active policy to express one high-level SSA value as the runtime - /// components that may cross an expanded scf.for, scf.while, or scf.if - /// boundary. value is the original pointer-like SSA value. The policy - /// receives the current rewrite context so it can reuse a known decomposition - /// and remap operands to the replacement IR. It may use builder and loc - /// to insert scalar address arithmetic at the correct rewrite position. - /// - /// For the block-pointer policy, decomposing an advance conceptually changes - /// only its offset components; base, shape, strides, and order are preserved: - /// %next = tt.advance %ptr, [%delta0, %delta1] - /// - /// decomposeValue(%next) -> DecomposedValue { - /// components = [%shape0, %shape1, %stride0, %stride1, - /// %offset0 + %delta0, %offset1 + %delta1], - /// invariants = [%base], attributes = [order] - /// } - /// The caller can put selected components into a new control-flow signature - /// or pass the whole descriptor to policy.recompose(). Unsupported values - /// or inconsistent component layouts return failure. - FailureOr decomposeValue(Value value, OpBuilder &builder, - Location loc) const { - return policy.decompose(value, getRewriteContext(), builder, loc); - } - - /// Records both ways in which later rewriting must understand an original - /// value. oldValue is the key from the original IR, info is its - /// component descriptor, and rebuiltValue is the pointer-like SSA value - /// created in the replacement IR. Pointer-aware code reads info from - /// decomposedValues; ordinary cloned users read rebuiltValue through - /// valueMapping. - /// - /// For example, after rebuilding an expanded loop argument: - /// %rebuilt_ptr = tt.make_tensor_ptr %base, ... %new_offset ... - /// bodyEnv.recordDecomposition(oldPtrArg, ptrInfo, rebuiltPtr); - /// - /// // An ordinary cloned tt.load receives %rebuilt_ptr. - /// bodyEnv.remap(oldPtrArg); - /// - /// // A later tt.advance/yield reuses the flattened ptrInfo directly. - /// bodyEnv.getRewriteContext().lookup(oldPtrArg); - void recordDecomposition(Value oldValue, const DecomposedValue &info, - Value rebuiltValue) { - decomposedValues[oldValue] = info; - valueMapping.map(oldValue, rebuiltValue); - } - - // Maps values from the original region to values in the replacement region. - IRMapping valueMapping; - // Concrete component state keyed by original values. Keeping this alongside - // the mapping lets pointer producers be flattened across nested rewrites. - DenseMap decomposedValues; - const ControlFlowRewritePolicy &policy; - const ControlFlowRewritePlan &plan; -}; - -// RewriteEnv is copied when entering a newly built region. The copy inherits -// mappings visible at the region boundary and records additional mappings only -// for that recursive rewrite. Nothing is stored on the IR or shared between -// decomposition policies. - -struct LoopPointerInfo { - // Original iter-argument/result position before signature expansion. - unsigned oldIndex = 0; - // Concrete descriptor used as the reconstruction template. - DecomposedValue initInfo; - // Ordered schema decided by ControlFlowSlotAnalysis. - SmallVector componentIndices; - SmallVector componentTypes; - // Positions occupied by those components in the replacement operation. - SmallVector newIndices; -}; - -struct IfPointerInfo { - unsigned oldIndex = 0; - SmallVector componentIndices; - SmallVector componentTypes; - std::optional thenInfo; -}; - -// Copies the values selected by indices into a new owning vector while -// preserving their input order. Indices must be unique and in bounds; this -// function reports failure instead of deduplicating or accessing invalid input. -// `sourceValues` is only borrowed while this function executes. -// -// Example: sourceValues = [shape, stride, offset] and indices = [2, 0, 1] -// produce [offset, shape, stride]. Indices [2, 0, 2] produce failure. -static FailureOr> gatherValues(ValueRange sourceValues, - ArrayRef indices) { - SmallVector values; - values.reserve(indices.size()); - llvm::SmallDenseSet seenIndices; - for (unsigned index : indices) { - if (index >= sourceValues.size() || !seenIndices.insert(index).second) - return failure(); - values.push_back(sourceValues[index]); - } - return values; -} - -// Returns a copy of decomposition with selected component values replaced. -// componentIndices and replacements are paired by position. The replacement -// fails if the ranges have different sizes, an index is out of bounds, or a -// replacement changes the component type; the input object remains unchanged. -// -// Example: -// decomposition.components = [shape, stride, originalOffset] -// componentIndices = [2], replacements = [nextOffset] -// result.components = [shape, stride, nextOffset] -static FailureOr -withReplacedComponents(DecomposedValue decomposition, - ArrayRef componentIndices, - ArrayRef replacements) { - if (componentIndices.size() != replacements.size()) - return failure(); - for (auto [componentIndex, replacement] : - llvm::zip(componentIndices, replacements)) { - if (componentIndex >= decomposition.components.size() || - decomposition.components[componentIndex].getType() != - replacement.getType()) - return failure(); - decomposition.components[componentIndex] = replacement; - } - return decomposition; -} - -static LogicalResult castPlannedComponents(DecomposedValue &value, - ArrayRef componentIndices, - ArrayRef componentTypes, - OpBuilder &builder, Location loc) { - if (componentIndices.size() != componentTypes.size()) - return failure(); - for (auto [index, type] : llvm::zip(componentIndices, componentTypes)) { - if (index >= value.components.size()) - return failure(); - FailureOr component = - castIntegerLike(builder, loc, value.components[index], type); - if (failed(component)) - return failure(); - value.components[index] = *component; - } - return success(); -} - -//===----------------------------------------------------------------------===// -// Shared recursive body rewrite -//===----------------------------------------------------------------------===// - -// Returns the first pointer descriptor whose oldIndex matches the requested -// position in the original control-flow signature. The range is an lvalue -// reference so the returned pointer refers to caller-owned storage, and the -// return type preserves whether that range exposes mutable or const elements. -// A missing index returns nullptr; duplicate indices keep first-match behavior. -// -// Example: [{oldIndex = 1}, {oldIndex = 3}] queried with 3 returns a pointer to -// the second element, while a query for 2 returns nullptr. -template -static auto findPointerInfoByOldIndex(InfoRange &pointerInfos, - unsigned oldIndex) - -> decltype(pointerInfos.data()) { - for (auto &info : pointerInfos) { - if (info.oldIndex == oldIndex) - return &info; - } - return nullptr; -} - -// A replacement loop carries selected scalar or tensor descriptor components -// instead of the original pointer iter-argument. Operations cloned from the -// original body still expect one pointer-typed block argument, so this function -// reconstructs that pointer at the replacement region entry. -// `pointerInfo.newIndices` selects the current component values from -// `newRegionArguments`, while -// `pointerInfo.componentIndices` identifies the descriptor fields that those -// values replace. The untouched fields come from `pointerInfo.initInfo`. -// -// On success, this function updates both parts of `regionEnv`: remapping the -// old region argument produces the rebuilt pointer, and looking up its -// decomposition returns the descriptor containing the current components. -// Operations cloned later in the same region can therefore use the complete -// pointer or its flattened state without referring to the old loop block. -// -// Example: -// oldRegionArgument = %old_ptr -// newRegionArguments = [%ordinary, %current_offset0, %current_offset1] -// pointerInfo.newIndices = [1, 2] -// pointerInfo.componentIndices = [4, 5] -// pointerInfo.initInfo.components = -// [shape0, shape1, stride0, stride1, initial_offset0, initial_offset1] -// -// The rebuilt descriptor keeps shape and stride, replaces the final two -// components with the current offsets, and records `%old_ptr -> %rebuilt_ptr` -// in `regionEnv`. Invalid indices, incompatible component types, or a policy -// that cannot recompose the descriptor return failure without recording a -// partial binding; the enclosing loop rewrite owns cleanup of inserted IR. -static LogicalResult bindLoopCarriedPointer(Value oldRegionArgument, - const LoopPointerInfo &pointerInfo, - ValueRange newRegionArguments, - OpBuilder &builder, Location loc, - RewriteEnv ®ionEnv) { - FailureOr> carriedComponentValues = - gatherValues(newRegionArguments, pointerInfo.newIndices); - if (failed(carriedComponentValues)) - return failure(); - - FailureOr argumentInfo = - withReplacedComponents(pointerInfo.initInfo, pointerInfo.componentIndices, - *carriedComponentValues); - if (failed(argumentInfo)) - return failure(); - - Value rebuiltPointer = - regionEnv.policy.recompose(*argumentInfo, builder, loc); - if (!rebuiltPointer) - return failure(); - - regionEnv.recordDecomposition(oldRegionArgument, *argumentInfo, - rebuiltPointer); - return success(); -} - -// Binds every original loop region argument to its replacement-region state. -// A pointer slot delegates to bindLoopCarriedPointer because one original -// pointer may occupy several component positions in the new signature. An -// ordinary slot keeps one SSA value and is mapped through oldToNewStart. -// -// Example: -// oldRegionArguments = [%x, %old_ptr, %y] -// newRegionArguments = [%new_x, %offset0, %offset1, %new_y] -// pointerInfos = [{oldIndex = 1, newIndices = [1, 2]}] -// oldToNewStart = [0, 1, 3] -// -// The resulting environment maps %x to %new_x and %y to %new_y. For -// %old_ptr, it gathers %offset0 and %offset1, rebuilds the complete pointer, -// and records both its SSA mapping and current decomposition. -// -// The function deliberately visits every argument after a binding failure. -// The surrounding builder callback may still need all available mappings to -// construct structurally valid temporary IR before its new loop is erased. -// It therefore accumulates failure and reports it only after the full range. -static LogicalResult bindLoopRegionArguments( - ValueRange oldRegionArguments, ValueRange newRegionArguments, - ArrayRef pointerInfos, ArrayRef oldToNewStart, - OpBuilder &builder, Location loc, RewriteEnv ®ionEnv) { - bool allArgumentsBound = true; - for (auto [oldIndex, oldRegionArgument] : - llvm::enumerate(oldRegionArguments)) { - const LoopPointerInfo *pointerInfo = - findPointerInfoByOldIndex(pointerInfos, oldIndex); - if (pointerInfo) { - if (failed(bindLoopCarriedPointer(oldRegionArgument, *pointerInfo, - newRegionArguments, builder, loc, - regionEnv))) - allArgumentsBound = false; - continue; - } - - if (oldIndex >= oldToNewStart.size() || - oldToNewStart[oldIndex] >= newRegionArguments.size()) { - allArgumentsBound = false; - continue; - } - regionEnv.valueMapping.map(oldRegionArgument, - newRegionArguments[oldToNewStart[oldIndex]]); - } - return allArgumentsBound ? success() : failure(); -} - -// Rewrites one original loop terminator operand list to the expanded signature -// used by the replacement loop. Ordinary operands remain one value and are -// remapped through `regionEnv`. A pointer operand is decomposed, normalized to -// the component types frozen by analysis, and expanded into the components -// selected by its `LoopPointerInfo`. -// -// Example: -// oldOperands = [%next_ptr, %sum] -// pointerInfo = { -// oldIndex = 0, componentIndices = [4, 5], newIndices = [0, 1] -// } -// currentRegionArguments = [%current_offset0, %current_offset1, %sum_arg] -// -// A valid `%next_ptr` decomposition produces -// `[%next_offset0, %next_offset1, %mapped_sum]`. If pointer decomposition or -// component normalization fails, the output instead uses -// `[%current_offset0, %current_offset1, %mapped_sum]`. The fallback keeps the -// temporary scf.yield/scf.condition structurally complete until the enclosing -// failed loop rewrite erases it. -// -// The output vector is separate from the LogicalResult intentionally. The -// function visits every old operand and fills all available fallback positions -// even after a pointer failure, then reports whether every pointer succeeded. -static LogicalResult rewriteLoopTerminatorOperands( - ValueRange oldOperands, ValueRange currentRegionArguments, - ArrayRef pointerInfos, OpBuilder &builder, Location loc, - RewriteEnv ®ionEnv, SmallVectorImpl &newOperands) { - bool allOperandsValid = true; - newOperands.clear(); - newOperands.reserve(currentRegionArguments.size()); - - auto appendFallbackComponents = [&](const LoopPointerInfo &pointerInfo) { - for (unsigned newIndex : pointerInfo.newIndices) - newOperands.push_back(currentRegionArguments[newIndex]); - }; - - for (auto [oldIndex, oldOperand] : llvm::enumerate(oldOperands)) { - const LoopPointerInfo *pointerInfo = - findPointerInfoByOldIndex(pointerInfos, oldIndex); - if (!pointerInfo) { - newOperands.push_back(regionEnv.remap(oldOperand)); - continue; - } - - FailureOr nextInfo = - regionEnv.decomposeValue(oldOperand, builder, loc); - if (failed(nextInfo) || failed(castPlannedComponents( - *nextInfo, pointerInfo->componentIndices, - pointerInfo->componentTypes, builder, loc))) { - allOperandsValid = false; - appendFallbackComponents(*pointerInfo); - continue; - } - - FailureOr> carriedComponentValues = - gatherValues(nextInfo->components, pointerInfo->componentIndices); - if (failed(carriedComponentValues)) { - allOperandsValid = false; - appendFallbackComponents(*pointerInfo); - continue; - } - newOperands.append(carriedComponentValues->begin(), - carriedComponentValues->end()); - } - - return allOperandsValid ? success() : failure(); -} - -// Rebuilds and records every result produced by a replacement loop. Ordinary -// results still occupy one position and are mapped through `oldToNewStart`. -// A pointer result occupies the positions listed in its `LoopPointerInfo`; the -// function gathers those components, writes them into the initial descriptor, -// recomposes the complete pointer, and records both its SSA mapping and current -// decomposition in `env`. -// -// Example: -// oldResults = [%old_sum, %old_ptr, %old_flag] -// newResults = [%new_sum, %offset0, %offset1, %new_flag] -// pointerInfo = { -// oldIndex = 1, componentIndices = [4, 5], newIndices = [1, 2] -// } -// oldToNewStart = [0, 1, 3] -// -// The function maps `%old_sum -> %new_sum` and `%old_flag -> %new_flag`. It -// inserts `%offset0` and `%offset1` into the pointer descriptor, rebuilds -// `%old_ptr`, and records `%old_ptr -> %rebuilt_ptr` plus that decomposition. -// -// The caller must set the builder insertion point after the replacement loop, -// so any rebuilt pointer dominates later operations. On failure this function -// returns immediately; the caller still owns and erases the replacement loop. -static LogicalResult -rebuildAndMapLoopResults(ValueRange oldResults, ValueRange newResults, - ArrayRef pointerInfos, - ArrayRef oldToNewStart, OpBuilder &builder, - RewriteEnv &env) { - for (auto [oldIndex, oldResult] : llvm::enumerate(oldResults)) { - const LoopPointerInfo *pointerInfo = - findPointerInfoByOldIndex(pointerInfos, oldIndex); - if (!pointerInfo) { - env.valueMapping.map(oldResult, newResults[oldToNewStart[oldIndex]]); - continue; - } - - FailureOr> resultComponentValues = - gatherValues(newResults, pointerInfo->newIndices); - if (failed(resultComponentValues)) - return failure(); - - FailureOr resultInfo = withReplacedComponents( - pointerInfo->initInfo, pointerInfo->componentIndices, - *resultComponentValues); - if (failed(resultInfo)) - return failure(); - - Value rebuiltPointer = - env.policy.recompose(*resultInfo, builder, oldResult.getLoc()); - if (!rebuiltPointer) - return failure(); - - env.recordDecomposition(oldResult, *resultInfo, rebuiltPointer); - } - - return success(); -} - -static LogicalResult rewriteControlFlowOp(Operation *op, OpBuilder &builder, - RewriteEnv &env); - -static LogicalResult materializePointerResult(Operation &originalOp, - Operation *clonedOp, - OpBuilder &builder, - RewriteEnv &env) { - // Each policy decides which pointer-producing operations need their exact - // components recorded immediately after cloning. - if (!env.policy.shouldDecomposeOperation(&originalOp)) - return success(); - - OpBuilder::InsertionGuard guard(builder); - builder.setInsertionPointAfter(clonedOp); - - bool decomposedAllResults = clonedOp->getNumResults() != 0; - for (auto [oldResult, clonedResult] : - llvm::zip(originalOp.getResults(), clonedOp->getResults())) { - if (!env.policy.isDecompositionTarget(oldResult)) { - decomposedAllResults = false; - continue; - } - - FailureOr info = - env.decomposeValue(clonedResult, builder, oldResult.getLoc()); - if (failed(info)) - return failure(); - - Value rebuilt = env.policy.recompose(*info, builder, oldResult.getLoc()); - if (!rebuilt) - return failure(); - env.recordDecomposition(oldResult, *info, rebuilt); - } - - // The replacement is recorded in the SSA mapping, so a side-effect-free - // clone whose every result was decomposed is redundant once it has no users. - if (decomposedAllResults && clonedOp->use_empty() && - isMemoryEffectFree(clonedOp)) - clonedOp->erase(); - - return success(); -} - -static LogicalResult rewriteBodyOps(Block *oldBlock, OpBuilder &builder, - RewriteEnv &env) { - // Process operations in program order. Nested control flow is rewritten - // recursively with the same policy; ordinary operations are cloned through - // the current SSA mapping. - for (Operation &originalOp : oldBlock->without_terminator()) { - if (isa(originalOp)) { - const ControlFlowOpAnalysis *analysis = env.plan.lookup(&originalOp); - if (!analysis) - return failure(); - if (analysis->needsRewrite()) { - if (failed(rewriteControlFlowOp(&originalOp, builder, env))) - return failure(); - continue; - } - } - Operation *clonedOp = builder.clone(originalOp, env.valueMapping); - if (failed(materializePointerResult(originalOp, clonedOp, builder, env))) - return failure(); - } - return success(); -} - -//===----------------------------------------------------------------------===// -// scf.for rewrite -//===----------------------------------------------------------------------===// - -static LogicalResult rewriteForOp(scf::ForOp forOp, OpBuilder &builder, - RewriteEnv &env) { - const ControlFlowOpAnalysis *analysis = env.plan.lookup(forOp); - if (!analysis || !analysis->needsRewrite()) - return failure(); - auto yieldOp = cast(forOp.getBody()->getTerminator()); - SmallVector pointerInfos; - - // The read-only analysis has already fixed every expanded slot and type. - // Materialization here recovers only the concrete values for that schema. - for (const ControlFlowSlotAnalysis &slot : analysis->slots) { - unsigned idx = slot.oldIndex; - if (idx >= forOp.getInitArgs().size() || idx >= yieldOp.getNumOperands() || - !env.policy.matches(forOp.getRegionIterArgs()[idx].getType())) - return failure(); - - FailureOr initInfo = - env.decomposeValue(forOp.getInitArgs()[idx], builder, forOp.getLoc()); - if (failed(initInfo) || failed(castPlannedComponents( - *initInfo, slot.componentIndices, - slot.componentTypes, builder, forOp.getLoc()))) - return failure(); - pointerInfos.push_back(LoopPointerInfo{ - idx, *initInfo, slot.componentIndices, slot.componentTypes, {}}); - } - - SmallVector newInitArgs; - SmallVector oldToNewStart(forOp.getInitArgs().size(), 0); - // Expand each owned pointer init into its runtime components. Non-pointer and - // other-policy slots retain one position in the new signature. - for (auto [idx, initArg] : llvm::enumerate(forOp.getInitArgs())) { - oldToNewStart[idx] = newInitArgs.size(); - if (LoopPointerInfo *info = findPointerInfoByOldIndex(pointerInfos, idx)) { - FailureOr> initComponents = - gatherValues(info->initInfo.components, info->componentIndices); - if (failed(initComponents)) - return failure(); - for (Value component : *initComponents) { - info->newIndices.push_back(newInitArgs.size()); - newInitArgs.push_back(component); - } - continue; - } - newInitArgs.push_back(env.remap(initArg)); - } - - bool bodyOk = true; - auto newForOp = builder.create( - forOp.getLoc(), env.remap(forOp.getLowerBound()), - env.remap(forOp.getUpperBound()), env.remap(forOp.getStep()), newInitArgs, - [&](OpBuilder &bodyBuilder, Location loc, Value iv, - ValueRange newRegionArgs) { - RewriteEnv bodyEnv = env; - bodyEnv.valueMapping.map(forOp.getInductionVar(), iv); - - // Bind pointer and ordinary iter-arguments before cloning body users. - if (failed(bindLoopRegionArguments( - forOp.getRegionIterArgs(), newRegionArgs, pointerInfos, - oldToNewStart, bodyBuilder, loc, bodyEnv))) - bodyOk = false; - - if (failed(rewriteBodyOps(forOp.getBody(), bodyBuilder, bodyEnv))) - bodyOk = false; - - SmallVector newYieldOperands; - if (failed(rewriteLoopTerminatorOperands( - yieldOp.getOperands(), newRegionArgs, pointerInfos, bodyBuilder, - yieldOp.getLoc(), bodyEnv, newYieldOperands))) - bodyOk = false; - - bodyBuilder.create(yieldOp.getLoc(), newYieldOperands); - }); - newForOp->setAttrs(forOp->getAttrs()); - - if (!bodyOk) { - newForOp.erase(); - return failure(); - } - - builder.setInsertionPointAfter(newForOp); - if (failed(rebuildAndMapLoopResults(forOp.getResults(), newForOp.getResults(), - pointerInfos, oldToNewStart, builder, - env))) { - newForOp.erase(); - return failure(); - } - - return success(); -} - -//===----------------------------------------------------------------------===// -// scf.while rewrite -//===----------------------------------------------------------------------===// - -static LogicalResult rewriteWhileOp(scf::WhileOp whileOp, OpBuilder &builder, - RewriteEnv &env) { - const ControlFlowOpAnalysis *analysis = env.plan.lookup(whileOp); - if (!analysis || !analysis->needsRewrite()) - return failure(); - scf::ConditionOp conditionOp = whileOp.getConditionOp(); - scf::YieldOp yieldOp = whileOp.getYieldOp(); - SmallVector pointerInfos; - - // The before arguments, condition forwarded values, after arguments and - // yield operands all consume the same precomputed positional schema. - for (const ControlFlowSlotAnalysis &slot : analysis->slots) { - unsigned idx = slot.oldIndex; - if (idx >= whileOp.getBeforeArguments().size() || - !env.policy.matches(whileOp.getBeforeArguments()[idx].getType()) || - idx >= whileOp.getInits().size() || - idx >= conditionOp.getArgs().size() || idx >= yieldOp.getNumOperands()) - return failure(); - - FailureOr initInfo = - env.decomposeValue(whileOp.getInits()[idx], builder, whileOp.getLoc()); - if (failed(initInfo) || - failed(castPlannedComponents(*initInfo, slot.componentIndices, - slot.componentTypes, builder, - whileOp.getLoc()))) - return failure(); - pointerInfos.push_back(LoopPointerInfo{ - idx, *initInfo, slot.componentIndices, slot.componentTypes, {}}); - } - - // Expand inits and result types in lockstep. oldToNewStart keeps untouched - // positions addressable even when earlier pointer slots expand by rank. - SmallVector newInits; - SmallVector newResultTypes; - SmallVector oldToNewStart(whileOp.getInits().size(), 0); - for (auto [idx, initArg] : llvm::enumerate(whileOp.getInits())) { - oldToNewStart[idx] = newInits.size(); - if (LoopPointerInfo *info = findPointerInfoByOldIndex(pointerInfos, idx)) { - FailureOr> initComponents = - gatherValues(info->initInfo.components, info->componentIndices); - if (failed(initComponents)) - return failure(); - for (Value component : *initComponents) { - info->newIndices.push_back(newInits.size()); - newInits.push_back(component); - newResultTypes.push_back(component.getType()); - } - continue; - } - newInits.push_back(env.remap(initArg)); - newResultTypes.push_back(whileOp.getResult(idx).getType()); - } - - bool bodyOk = true; - auto newWhileOp = builder.create( - whileOp.getLoc(), newResultTypes, newInits, - [&](OpBuilder &bodyBuilder, Location loc, ValueRange newRegionArgs) { - RewriteEnv beforeEnv = env; - // Bind the before-region arguments, then rewrite the body and the - // values forwarded by scf.condition. - if (failed(bindLoopRegionArguments( - whileOp.getBeforeArguments(), newRegionArgs, pointerInfos, - oldToNewStart, bodyBuilder, loc, beforeEnv))) - bodyOk = false; - - if (failed(rewriteBodyOps(whileOp.getBeforeBody(), bodyBuilder, - beforeEnv))) - bodyOk = false; - - SmallVector newConditionArgs; - if (failed(rewriteLoopTerminatorOperands( - conditionOp.getArgs(), newRegionArgs, pointerInfos, bodyBuilder, - conditionOp.getLoc(), beforeEnv, newConditionArgs))) - bodyOk = false; - - bodyBuilder.create( - conditionOp.getLoc(), beforeEnv.remap(conditionOp.getCondition()), - newConditionArgs); - }, - [&](OpBuilder &bodyBuilder, Location loc, ValueRange newRegionArgs) { - RewriteEnv afterEnv = env; - // Bind the after-region arguments before rewriting the body and its - // backedge yield. - if (failed(bindLoopRegionArguments( - whileOp.getAfterArguments(), newRegionArgs, pointerInfos, - oldToNewStart, bodyBuilder, loc, afterEnv))) - bodyOk = false; - - if (failed( - rewriteBodyOps(whileOp.getAfterBody(), bodyBuilder, afterEnv))) - bodyOk = false; - - SmallVector newYieldOperands; - if (failed(rewriteLoopTerminatorOperands( - yieldOp.getOperands(), newRegionArgs, pointerInfos, bodyBuilder, - yieldOp.getLoc(), afterEnv, newYieldOperands))) - bodyOk = false; - - bodyBuilder.create(yieldOp.getLoc(), newYieldOperands); - }); - newWhileOp->setAttrs(whileOp->getAttrs()); - - if (!bodyOk) { - newWhileOp.erase(); - return failure(); - } - - builder.setInsertionPointAfter(newWhileOp); - if (failed(rebuildAndMapLoopResults(whileOp.getResults(), - newWhileOp.getResults(), pointerInfos, - oldToNewStart, builder, env))) { - newWhileOp.erase(); - return failure(); - } - - return success(); -} - -//===----------------------------------------------------------------------===// -// scf.if component planning and rewrite -//===----------------------------------------------------------------------===// - -static LogicalResult rewriteIfOp(scf::IfOp ifOp, OpBuilder &builder, - RewriteEnv &env) { - const ControlFlowOpAnalysis *analysis = env.plan.lookup(ifOp); - if (!analysis || !analysis->needsRewrite() || - (!ifOp.elseBlock() && analysis->rewritesOwnSignature())) - return failure(); - - bool hasElse = static_cast(ifOp.elseBlock()); - scf::YieldOp thenYield = ifOp.thenYield(); - scf::YieldOp elseYield = hasElse ? ifOp.elseYield() : scf::YieldOp(); - SmallVector pointerInfos; - - for (const ControlFlowSlotAnalysis &slot : analysis->slots) { - if (slot.oldIndex >= ifOp.getNumResults() || - !env.policy.matches(ifOp.getResult(slot.oldIndex).getType()) || - slot.componentIndices.size() != slot.componentTypes.size()) - return failure(); - pointerInfos.push_back(IfPointerInfo{slot.oldIndex, slot.componentIndices, - slot.componentTypes, std::nullopt}); - } - - // Expand only result positions selected by analysis. An if with no pointer - // result may still be rebuilt because one of its nested operations changes. - SmallVector newResultTypes; - for (auto [idx, result] : llvm::enumerate(ifOp.getResults())) { - if (const IfPointerInfo *info = - findPointerInfoByOldIndex(pointerInfos, idx)) { - newResultTypes.append(info->componentTypes.begin(), - info->componentTypes.end()); - continue; - } - newResultTypes.push_back(result.getType()); - } - - bool bodyOk = true; - auto buildBranch = [&](OpBuilder &branchBuilder, - bool isThen) -> LogicalResult { - // Each branch gets an isolated environment because values defined in one - // branch must never be visible while cloning the other branch. - RewriteEnv branchEnv = env; - Block *oldBlock = isThen ? ifOp.thenBlock() : ifOp.elseBlock(); - scf::YieldOp oldYield = isThen ? thenYield : elseYield; - if (failed(rewriteBodyOps(oldBlock, branchBuilder, branchEnv))) - return failure(); - - SmallVector newYieldOperands; - for (auto [idx, oldOperand] : llvm::enumerate(oldYield.getOperands())) { - if (IfPointerInfo *info = findPointerInfoByOldIndex(pointerInfos, idx)) { - FailureOr branchInfo = branchEnv.decomposeValue( - oldOperand, branchBuilder, oldYield.getLoc()); - if (failed(branchInfo) || - failed(castPlannedComponents(*branchInfo, info->componentIndices, - info->componentTypes, branchBuilder, - oldYield.getLoc()))) - return failure(); - if (isThen) - info->thenInfo = *branchInfo; - FailureOr> values = - gatherValues(branchInfo->components, info->componentIndices); - if (failed(values)) - return failure(); - newYieldOperands.append(values->begin(), values->end()); - continue; - } - newYieldOperands.push_back(branchEnv.remap(oldOperand)); - } - branchBuilder.create(oldYield.getLoc(), newYieldOperands); - return success(); - }; - - // Create the shell first, then clone each old branch into the corresponding - // new region with independent mappings. - auto newIfOp = builder.create( - ifOp.getLoc(), newResultTypes, env.remap(ifOp.getCondition()), hasElse); - newIfOp->setAttrs(ifOp->getAttrs()); - - // The then region always exists, including for a result-less one-arm if. - // Rewriting it is still required when it contains affected nested SCF. - { - OpBuilder::InsertionGuard guard(builder); - if (newResultTypes.empty()) { - newIfOp.thenBlock()->getTerminator()->erase(); - builder.setInsertionPointToEnd(newIfOp.thenBlock()); - } else { - builder.setInsertionPointToStart(newIfOp.thenBlock()); - } - if (failed(buildBranch(builder, /*isThen=*/true))) - bodyOk = false; - } - // An else region exists only for the two-arm form. In particular, do not - // access elseBlock() merely because the then region contains nested work. - if (hasElse) { - OpBuilder::InsertionGuard guard(builder); - if (newResultTypes.empty()) { - newIfOp.elseBlock()->getTerminator()->erase(); - builder.setInsertionPointToEnd(newIfOp.elseBlock()); - } else { - builder.setInsertionPointToStart(newIfOp.elseBlock()); - } - if (failed(buildBranch(builder, /*isThen=*/false))) - bodyOk = false; - } - - if (!bodyOk) { - newIfOp.erase(); - return failure(); - } - - // Reassemble the pointer immediately after the replacement if. Downstream - // operations therefore keep their original operand types; decomposition is - // limited to the control-flow boundary itself. - builder.setInsertionPointAfter(newIfOp); - unsigned newResultIndex = 0; - for (auto [idx, oldResult] : llvm::enumerate(ifOp.getResults())) { - if (const IfPointerInfo *info = - findPointerInfoByOldIndex(pointerInfos, idx)) { - SmallVector componentValues; - for (unsigned i = 0; i < info->componentIndices.size(); ++i) - componentValues.push_back(newIfOp.getResult(newResultIndex++)); - FailureOr resultInfo = withReplacedComponents( - *info->thenInfo, info->componentIndices, componentValues); - if (failed(resultInfo)) { - newIfOp.erase(); - return failure(); - } - Value rebuilt = - env.policy.recompose(*resultInfo, builder, oldResult.getLoc()); - if (!rebuilt) { - newIfOp.erase(); - return failure(); - } - env.recordDecomposition(oldResult, *resultInfo, rebuilt); - continue; - } - env.valueMapping.map(oldResult, newIfOp.getResult(newResultIndex++)); - } - - return success(); -} - -static LogicalResult rewriteControlFlowOp(Operation *op, OpBuilder &builder, - RewriteEnv &env) { - // Keep the operation dispatch next to the shared recursive implementation: - // all supported region operations must obey the same mapping and cleanup - // rules. Pointer-specific semantics enter only through env.policy. - if (auto forOp = dyn_cast(op)) - return rewriteForOp(forOp, builder, env); - if (auto whileOp = dyn_cast(op)) - return rewriteWhileOp(whileOp, builder, env); - if (auto ifOp = dyn_cast(op)) - return rewriteIfOp(ifOp, builder, env); - // TODO: Add the frontend-produced scope.scope operation here. Scope support - // belongs in this shared plumbing rather than in both decompositions. - return failure(); -} - -static FailureOr> -collectReplacements(Operation *op, const RewriteEnv &env) { - SmallVector replacements; - replacements.reserve(op->getNumResults()); - for (Value result : op->getResults()) { - // Unlike RewriteEnv::remap(), replacement collection must not fall back to - // the original result. Such a fallback would hide an unhandled result slot - // and ask replaceOp to replace a value with itself. - Value replacement = env.valueMapping.lookupOrNull(result); - if (!replacement) - return failure(); - replacements.push_back(replacement); - } - return replacements; -} - -static LogicalResult -tryDecoupleControlFlowOp(Operation *op, IRRewriter &rewriter, - const ControlFlowRewritePolicy &policy, - const ControlFlowRewritePlan &plan) { - // Build a replacement beside the original operation. The original operation - // itself remains until every result has a valid mapped value, after which the - // standard rewriter performs the externally visible replacement. - // TODO: Track and erase policy materializations created outside the new SCF - // operation if an unexpected rewrite-time validation fails. Read-only - // analysis makes that path exceptional, but failure should still be atomic. - RewriteEnv env(policy, plan); - rewriter.setInsertionPoint(op); - if (failed(rewriteControlFlowOp(op, rewriter, env))) - return failure(); - - FailureOr> replacements = collectReplacements(op, env); - if (failed(replacements)) - return failure(); - rewriter.replaceOp(op, *replacements); - return success(); -} - -} // namespace - -namespace mlir::triton::controlflow { - -LogicalResult -applyControlFlowRewritePlan(ModuleOp module, - const ControlFlowRewritePolicy &policy, - const ControlFlowRewritePlan &plan) { - IRRewriter rewriter(module.getContext()); - // Analysis and application are consecutive and no IR mutation occurs in - // between, so rediscover the roots from the module instead of duplicating - // traversal state in the immutable operation plan. - for (Operation *root : collectOutermostControlFlowOps(module)) { - const ControlFlowOpAnalysis *rootAnalysis = plan.lookup(root); - if (!rootAnalysis) { - root->emitError("missing frozen control-flow rewrite decision"); - return failure(); - } - if (!rootAnalysis->needsRewrite()) - continue; - if (failed(tryDecoupleControlFlowOp(root, rewriter, policy, plan))) { - root->emitError("failed to apply analyzed pointer decomposition"); - return failure(); - } - } - return success(); -} - -LogicalResult rewriteControlFlow(ModuleOp module, - const ControlFlowRewritePolicy &policy) { - FailureOr plan = analyzeControlFlow(module, policy); - if (failed(plan)) - return failure(); - return applyControlFlowRewritePlan(module, policy, *plan); -} - -} // namespace mlir::triton::controlflow diff --git a/third_party/ascend/lib/TritonControlFlowOpt/TensorPtrDecompose.cpp b/third_party/ascend/lib/TritonControlFlowOpt/TensorPtrDecompose.cpp deleted file mode 100644 index c55a80186e..0000000000 --- a/third_party/ascend/lib/TritonControlFlowOpt/TensorPtrDecompose.cpp +++ /dev/null @@ -1,414 +0,0 @@ -/* - * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ - -#include "TritonControlFlowOpt/TensorPtrDecompose.h" - -#include "TritonControlFlowOpt/ControlFlowRewrite.h" -#include "Utils/Utils.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -using namespace mlir; -using namespace mlir::triton; -using namespace mlir::triton::controlflow; - -namespace { - -static constexpr unsigned kOffsetsComponent = 0; -static constexpr unsigned kBaseInvariant = 0; -static constexpr unsigned kBaseIsScalarAttribute = 0; - -/// Identifies tensor-of-pointers handled by this stage: -/// `tensor<...x!tt.ptr>`. A block pointer has the different scalar type -/// `!tt.ptr>` and is handled by BlockPtrDecompose. -static bool isTensorPointerType(Type type) { - auto tensorType = dyn_cast(type); - return tensorType && isa(tensorType.getElementType()); -} - -/// Tensor-pointer state used only by this policy: -/// -/// control-flow components = [complete_offsets] -/// rewrite-only invariants = [common_base] -/// attributes = [base_is_scalar] -/// -/// Only `components` expand an SCF signature. The common base is deliberately -/// kept out of iter-args and results; it must be identical at every incoming -/// edge and is used to rebuild the tensor-of-pointers inside and after the -/// rewritten control-flow operation. -static RankedTensorType getDefaultOffsetsType(Type pointerType) { - auto pointerTensor = cast(pointerType); - return RankedTensorType::get(pointerTensor.getShape(), - IntegerType::get(pointerType.getContext(), 32), - pointerTensor.getEncoding()); -} - -/// Validates the policy-owned component layout, including the offsets shape -/// and the representation selected for the invariant base. -static bool hasValidSchema(Type originalType, Type offsetsType, - ArrayRef invariants, - ArrayRef attributes) { - auto pointerTensor = dyn_cast(originalType); - auto offsetsTensor = dyn_cast(offsetsType); - if (!pointerTensor || !offsetsTensor || invariants.size() != 1 || - attributes.size() != 1 || - !isa(pointerTensor.getElementType()) || - !isa(offsetsTensor.getElementType()) || - pointerTensor.getShape() != offsetsTensor.getShape() || - pointerTensor.getEncoding() != offsetsTensor.getEncoding()) - return false; - - auto baseIsScalar = dyn_cast(attributes[kBaseIsScalarAttribute]); - if (!baseIsScalar) - return false; - Type expectedBaseType = - baseIsScalar.getValue() ? pointerTensor.getElementType() : originalType; - return invariants[kBaseInvariant].getType() == expectedBaseType; -} - -/// Validates the concrete Value-based state created while rewriting IR. -static bool hasValidLayout(const DecomposedValue &value) { - return value.components.size() == 1 && - hasValidSchema(value.originalType, - value.components[kOffsetsComponent].getType(), - value.invariants, value.attributes); -} - -/// Validates the read-only counterpart without materializing component Values. -static bool hasValidLayout(const AnalyzedValue &value) { - return value.components.size() == 1 && - hasValidSchema(value.originalType, - value.components[kOffsetsComponent].type, - value.invariants, value.attributes); -} - -template -static bool haveSameTensorPtrBaseSchema(const StateT &lhs, const StateT &rhs) { - return hasValidLayout(lhs) && hasValidLayout(rhs) && - lhs.originalType == rhs.originalType && - lhs.invariants == rhs.invariants && lhs.attributes == rhs.attributes; -} - -/// Whether the invariant base must be broadcast before rebuilding addptr. -static bool hasScalarBase(const DecomposedValue &value) { - return cast(value.attributes[kBaseIsScalarAttribute]).getValue(); -} - -/// Creates the initial complete offsets for an opaque tensor-of-pointers. -/// i32 is the current default. A later addptr or control-flow merge may promote -/// it to i64 through getWiderOffsetsType. -static Value createZeroOffsets(OpBuilder &builder, Location loc, - Type pointerType) { - RankedTensorType offsetsType = getDefaultOffsetsType(pointerType); - auto elementType = cast(offsetsType.getElementType()); - auto zero = DenseElementsAttr::get(offsetsType, - builder.getIntegerAttr(elementType, 0)); - return builder.create(loc, zero); -} - -/// Computes the common type used by an offset addition or SCF merge. -/// -/// Complete offsets are always integer tensors with the pointer tensor's shape -/// and encoding. When element widths differ, use the wider signed width. -static FailureOr getWiderOffsetsType(Type lhs, Type rhs) { - auto lhsTensor = dyn_cast(lhs); - auto rhsTensor = dyn_cast(rhs); - if (!lhsTensor || !rhsTensor || - lhsTensor.getShape() != rhsTensor.getShape() || - lhsTensor.getEncoding() != rhsTensor.getEncoding()) - return failure(); - auto lhsElementInt = dyn_cast(lhsTensor.getElementType()); - auto rhsElementInt = dyn_cast(rhsTensor.getElementType()); - if (!lhsElementInt || !rhsElementInt) - return failure(); - if (lhsElementInt.getWidth() == rhsElementInt.getWidth()) { - if (lhsElementInt != rhsElementInt) - return failure(); - return lhs; - } - return lhsElementInt.getWidth() >= rhsElementInt.getWidth() ? lhs : rhs; -} - -/// Adds two offset values after promoting both operands to their common type. -/// This is the concrete IR counterpart of joinComponentTypes used in the -/// read-only analysis. -static Value createOffsetsAdd(OpBuilder &builder, Location loc, Value lhs, - Value rhs) { - if (!lhs || !rhs) - return nullptr; - FailureOr type = getWiderOffsetsType(lhs.getType(), rhs.getType()); - if (failed(type)) - return nullptr; - FailureOr convertedLhs = castIntegerLike(builder, loc, lhs, *type); - FailureOr convertedRhs = castIntegerLike(builder, loc, rhs, *type); - if (failed(convertedLhs) || failed(convertedRhs)) - return nullptr; - return builder.create(loc, *convertedLhs, *convertedRhs); -} - -/// Tensor-pointer semantics plugged into the shared control-flow machinery. -/// -/// The analysis methods compute a symbolic schema without changing IR. The -/// rewrite methods later materialize the exact Values for that already-chosen -/// schema. Keeping both implementations here makes their layouts directly -/// comparable and avoids hiding pointer semantics in ControlFlowRewrite. -class TensorPtrDecomposePolicy final : public ControlFlowRewritePolicy { -public: - /// Selects only tensor-of-pointers owned by this decomposition stage. - bool matches(Type type) const override { - // A tensor pointer here means tensor<...x!tt.ptr<...>>. Scalar block - // pointers have already been handled by BlockPtrDecompose. - return isTensorPointerType(type); - } - - //===--------------------------------------------------------------------===// - // Read-only schema analysis - //===--------------------------------------------------------------------===// - - /// Recovers `{common_base, complete_offsets, base_is_scalar}` without - /// creating constants, additions, or pointer operations. - FailureOr - analyzeValue(Value value, - ControlFlowAnalysisContext &context) const override { - // Region arguments and previously merged control-flow results are already - // bound in the stage-scoped cache. Reusing them is what lets analysis cross - // sibling and nested SCF without inspecting rewritten IR. - if (const AnalyzedValue *known = context.lookupValue(value)) { - if (!matches(known->originalType)) - return failure(); - return *known; - } - - // addptr preserves the upstream common base and adds another contribution - // to complete_offsets. The result Value is used only as a symbolic identity - // showing that this component changed. - if (auto addPtr = value.getDefiningOp()) { - FailureOr result = context.analyzeValue(addPtr.getPtr()); - if (failed(result) || !hasValidLayout(*result)) - return failure(); - FailureOr offsetsType = - getWiderOffsetsType(result->components[kOffsetsComponent].type, - addPtr.getOffset().getType()); - if (failed(offsetsType)) - return failure(); - result->originalType = value.getType(); - result->components[kOffsetsComponent] = { - *offsetsType, ComponentIdentity::fromValue(value, kOffsetsComponent)}; - return *result; - } - - // Splatting one scalar pointer establishes the canonical common-base form: - // every lane starts at the scalar base and therefore has zero offset. - if (auto splat = value.getDefiningOp()) { - if (!isa(splat.getSrc().getType())) - return failure(); - Type offsetsType = getDefaultOffsetsType(value.getType()); - return AnalyzedValue{value.getType(), - {{offsetsType, ComponentIdentity::zero()}}, - {splat.getSrc()}, - {BoolAttr::get(value.getContext(), true)}}; - } - - // An otherwise opaque tensor-of-pointers is treated as an already-vector - // base with zero additional offsets. This preserves current behavior until - // common-base analysis is shared with TritonToUnstructure. - if (!matches(value.getType())) - return failure(); - Type offsetsType = getDefaultOffsetsType(value.getType()); - return AnalyzedValue{value.getType(), - {{offsetsType, ComponentIdentity::zero()}}, - {value}, - {BoolAttr::get(value.getContext(), false)}}; - } - - /// Tensor pointers have exactly one loop-transfer candidate: the complete - /// per-lane offsets tensor at component index 0. - FailureOr> - getLoopCandidateComponents(const AnalyzedValue &value) const override { - if (!hasValidLayout(value)) - return failure(); - return SmallVector{kOffsetsComponent}; - } - - /// Classifies the loop offsets as transferred only when the backedge changes - /// their symbolic identity. The original pointer type, common base, and base - /// representation must remain invariant for reconstruction to be valid. - FailureOr> - getLoopTransferredComponents(const AnalyzedValue &initial, - const AnalyzedValue ®ionArgument, - const AnalyzedValue &next) const override { - if (!haveSameTensorPtrBaseSchema(initial, regionArgument) || - !haveSameTensorPtrBaseSchema(initial, next) || - failed(joinComponentTypes(initial.components[kOffsetsComponent].type, - next.components[kOffsetsComponent].type))) - return failure(); - // Yielding the region argument unchanged requires no new SCF iter-arg. - if (regionArgument.components[kOffsetsComponent].identity == - next.components[kOffsetsComponent].identity) - return SmallVector{}; - return SmallVector{kOffsetsComponent}; - } - - /// Merges the two `scf.if` pointer states. Different complete-offset - /// identities become an if result; the base and its scalar/tensor form must - /// agree so one pointer can be rebuilt after the if. - FailureOr> - getIfTransferredComponents(const AnalyzedValue &thenValue, - const AnalyzedValue &elseValue) const override { - if (!haveSameTensorPtrBaseSchema(thenValue, elseValue) || - failed( - joinComponentTypes(thenValue.components[kOffsetsComponent].type, - elseValue.components[kOffsetsComponent].type))) - return failure(); - // Identical symbolic offsets are available outside the if as an invariant. - if (thenValue.components[kOffsetsComponent].identity == - elseValue.components[kOffsetsComponent].identity) - return SmallVector{}; - return SmallVector{kOffsetsComponent}; - } - - /// Chooses the offsets type carried by the replacement control-flow op. - FailureOr joinComponentTypes(Type lhs, Type rhs) const override { - return getWiderOffsetsType(lhs, rhs); - } - - //===--------------------------------------------------------------------===// - // Concrete rewrite materialization - //===--------------------------------------------------------------------===// - - /// addptr results must be decomposed immediately after cloning so later users - /// can obtain their accumulated offsets from the rewrite context. - bool shouldDecomposeOperation(Operation *op) const override { - return isa(op); - } - - /// Materializes the same decomposition described by analyzeValue. This may - /// create zero constants and integer additions, so it is called only after - /// the complete control-flow subtree has passed read-only analysis. - FailureOr decompose(Value value, - const ControlFlowRewriteContext &context, - OpBuilder &builder, - Location loc) const override { - // Rewritten region arguments and nested results have exact component Values - // recorded by ControlFlowRewrite; never reconstruct them from old IR. - if (const DecomposedValue *known = context.lookup(value)) { - if (!matches(known->originalType)) - return failure(); - return *known; - } - - // Ordinary SSA inputs may already have been cloned into the replacement - // region, so pointer producer matching must use the remapped Value. - value = context.remap(value); - // Recursively flatten an addptr chain into one complete offsets tensor. - if (auto addPtr = value.getDefiningOp()) { - FailureOr result = - decompose(addPtr.getPtr(), context, builder, loc); - if (failed(result) || !hasValidLayout(*result)) - return failure(); - - // Insert the accumulated addition beside the addptr being decomposed; - // ControlFlowRewrite will use this Value in the expanded SCF signature. - OpBuilder::InsertionGuard guard(builder); - builder.setInsertionPoint(addPtr); - Value offsets = createOffsetsAdd(builder, addPtr.getLoc(), - result->components[kOffsetsComponent], - context.remap(addPtr.getOffset())); - if (!offsets) - return failure(); - result->originalType = value.getType(); - result->components[kOffsetsComponent] = offsets; - return *result; - } - - // A scalar pointer splat materializes zero offsets while retaining the - // scalar source as the common-base invariant. - if (auto splat = value.getDefiningOp()) { - if (!isa(splat.getSrc().getType())) - return failure(); - OpBuilder::InsertionGuard guard(builder); - builder.setInsertionPoint(splat); - Value offsets = - createZeroOffsets(builder, splat.getLoc(), value.getType()); - if (!offsets) - return failure(); - return DecomposedValue{value.getType(), - {offsets}, - {splat.getSrc()}, - {builder.getBoolAttr(true)}}; - } - - // Fallback for an opaque tensor base: use the entire tensor-of-pointers as - // the invariant base and represent only subsequent displacement in offsets. - if (!matches(value.getType())) - return failure(); - - OpBuilder::InsertionGuard guard(builder); - if (Operation *definingOp = value.getDefiningOp()) - builder.setInsertionPointAfter(definingOp); - else if (auto blockArg = dyn_cast(value)) - builder.setInsertionPointToStart(blockArg.getOwner()); - // Place the zero where it dominates every later reconstruction using it. - Value offsets = createZeroOffsets(builder, loc, value.getType()); - if (!offsets) - return failure(); - return DecomposedValue{ - value.getType(), {offsets}, {value}, {builder.getBoolAttr(false)}}; - } - - /// Rebuilds the original tensor-of-pointers from the invariant base and the - /// complete offsets selected/carried by the rewritten control flow. - Value recompose(const DecomposedValue &value, OpBuilder &builder, - Location loc) const override { - if (!hasValidLayout(value)) - return nullptr; - Value base = value.invariants[kBaseInvariant]; - // addptr requires matching tensor lanes; broadcast a scalar common base - // only when the decomposition recorded `base_is_scalar = true`. - if (hasScalarBase(value)) - base = builder.create(loc, value.originalType, base); - return builder.create( - loc, value.originalType, base, value.components[kOffsetsComponent]); - } -}; - -} // namespace - -namespace mlir::triton::controlflow { - -/// Runs tensor-pointer decomposition after CFG structuring/block-pointer -/// handling. The shared driver analyzes each outermost SCF root, then rewrites -/// only the complete-offset components selected by this policy. -LogicalResult runTensorPtrDecompose(ModuleOp module) { - // Carry only complete per-lane offsets through SCF. The common scalar base - // remains a rewrite invariant and is used to rebuild tensor-of-pointers at - // each region boundary. This decomposition is independent of - // BlockPtrDecompose. - // TODO: Replace this local extraction with TritonToUnstructure's common-base - // analysis. Different or lane-wise bases must become explicit diagnostics - // instead of pattern misses. - TensorPtrDecomposePolicy policy; - return rewriteControlFlow(module, policy); -} - -} // namespace mlir::triton::controlflow diff --git a/third_party/ascend/lib/TritonControlFlowOpt/TritonControlFlowOptPass.cpp b/third_party/ascend/lib/TritonControlFlowOpt/TritonControlFlowOptPass.cpp index bc6addb507..e2f74c70fa 100644 --- a/third_party/ascend/lib/TritonControlFlowOpt/TritonControlFlowOptPass.cpp +++ b/third_party/ascend/lib/TritonControlFlowOpt/TritonControlFlowOptPass.cpp @@ -22,50 +22,2139 @@ #include "TritonControlFlowOpt/TritonControlFlowOptPass.h" -#include "TritonControlFlowOpt/BlockPtrDecompose.h" -#include "TritonControlFlowOpt/CFGStructuring.h" -#include "TritonControlFlowOpt/TensorPtrDecompose.h" - #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/PatternMatch.h" #include "mlir/IR/Verifier.h" +#include "mlir/IR/Visitors.h" #include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" + +#include + +#define DEBUG_TYPE "triton-control-flow-opt" + +using namespace mlir; +using namespace triton; + +namespace { + +static bool isSupportedReturn(Operation *op) { + return isa( + op); +} + +static SmallVector getCfgSuccessors(Block *block) { + Operation *term = block->getTerminator(); + if (auto br = dyn_cast(term)) + return {br.getDest()}; + if (auto condBr = dyn_cast(term)) + return {condBr.getTrueDest(), condBr.getFalseDest()}; + return {}; +} + +static DenseMap computeDistances(Block *start) { + DenseMap distances; + SmallVector worklist; + + distances[start] = 0; + worklist.push_back(start); + + for (unsigned i = 0; i < worklist.size(); ++i) { + Block *block = worklist[i]; + unsigned nextDistance = distances[block] + 1; + for (Block *successor : getCfgSuccessors(block)) { + if (successor->getParent() != start->getParent()) + continue; + if (distances.count(successor)) + continue; + distances[successor] = nextDistance; + worklist.push_back(successor); + } + } + + return distances; +} + +static FailureOr findNearestCommonBlock(Block *lhs, Block *rhs, + Location loc, + bool emitDiagnostic = true) { + DenseMap lhsDistances = computeDistances(lhs); + DenseMap rhsDistances = computeDistances(rhs); + + Block *best = nullptr; + unsigned bestMaxDistance = std::numeric_limits::max(); + unsigned bestTotalDistance = std::numeric_limits::max(); + + for (auto &entry : lhsDistances) { + Block *candidate = entry.first; + auto rhsIt = rhsDistances.find(candidate); + if (rhsIt == rhsDistances.end()) + continue; + + unsigned lhsDistance = entry.second; + unsigned rhsDistance = rhsIt->second; + unsigned maxDistance = std::max(lhsDistance, rhsDistance); + unsigned totalDistance = lhsDistance + rhsDistance; + if (maxDistance < bestMaxDistance || + (maxDistance == bestMaxDistance && totalDistance < bestTotalDistance)) { + best = candidate; + bestMaxDistance = maxDistance; + bestTotalDistance = totalDistance; + } + } + + if (!best && emitDiagnostic) { + emitError(loc) << "unsupported non-tree control flow: branch arms do not " + "reach a common convergence block"; + return failure(); + } + if (!best) + return failure(); + + return best; +} + +static LogicalResult replaceBlockArguments(Block *block, ValueRange incoming, + Location loc) { + if (block->getNumArguments() != incoming.size()) { + emitError(loc) << "invalid branch operand count while structuring " + "control flow: " + << incoming.size() << " operands for " + << block->getNumArguments() << " block arguments"; + return failure(); + } + + for (auto [arg, value] : llvm::zip(block->getArguments(), incoming)) + arg.replaceAllUsesWith(value); + return success(); +} + +static void moveBlockBodyBefore(Block *block, OpBuilder &builder) { + SmallVector movedOps = llvm::map_to_vector( + block->without_terminator(), [](Operation &op) { return &op; }); + for (Operation *op : movedOps) + op->moveBefore(builder.getInsertionBlock(), builder.getInsertionPoint()); +} + +struct ReturnPathResult { + SmallVector operands; +}; + +static FailureOr> buildRegionPath(Block *block, + ValueRange incoming, + Block *stopBlock, + OpBuilder &builder); + +static FailureOr +buildReturnPath(Block *block, ValueRange incoming, OpBuilder &builder); + +static FailureOr buildTerminalValueIf(cf::CondBranchOp condBr, + OpBuilder &builder); + +static FailureOr buildStructuredIf(cf::CondBranchOp condBr, + Block *joinBlock, + OpBuilder &builder) { + SmallVector resultTypes; + resultTypes.reserve(joinBlock->getNumArguments()); + for (BlockArgument arg : joinBlock->getArguments()) + resultTypes.push_back(arg.getType()); + + auto ifOp = builder.create(condBr.getLoc(), resultTypes, + condBr.getCondition(), + /*withElseRegion=*/true); + + { + OpBuilder::InsertionGuard guard(builder); + Operation *autoYield = + resultTypes.empty() ? ifOp.thenBlock()->getTerminator() : nullptr; + if (autoYield) + builder.setInsertionPoint(autoYield); + else + builder.setInsertionPointToStart(ifOp.thenBlock()); + FailureOr> thenYield = buildRegionPath( + condBr.getTrueDest(), condBr.getTrueDestOperands(), joinBlock, builder); + if (failed(thenYield)) + return failure(); + if (thenYield->size() != resultTypes.size()) { + condBr.emitError("then branch yields ") + << thenYield->size() << " values, expected " << resultTypes.size(); + return failure(); + } + if (!autoYield) + builder.create(condBr.getLoc(), *thenYield); + } + + { + OpBuilder::InsertionGuard guard(builder); + Operation *autoYield = + resultTypes.empty() ? ifOp.elseBlock()->getTerminator() : nullptr; + if (autoYield) + builder.setInsertionPoint(autoYield); + else + builder.setInsertionPointToStart(ifOp.elseBlock()); + FailureOr> elseYield = + buildRegionPath(condBr.getFalseDest(), condBr.getFalseDestOperands(), + joinBlock, builder); + if (failed(elseYield)) + return failure(); + if (elseYield->size() != resultTypes.size()) { + condBr.emitError("else branch yields ") + << elseYield->size() << " values, expected " << resultTypes.size(); + return failure(); + } + if (!autoYield) + builder.create(condBr.getLoc(), *elseYield); + } + + return ifOp; +} + +static bool haveSameTypes(ValueRange lhs, ValueRange rhs) { + if (lhs.size() != rhs.size()) + return false; + for (auto [lhsValue, rhsValue] : llvm::zip(lhs, rhs)) { + if (lhsValue.getType() != rhsValue.getType()) + return false; + } + return true; +} + +static bool haveSameTypes(ValueRange values, ArrayRef types) { + if (values.size() != types.size()) + return false; + for (auto [value, type] : llvm::zip(values, types)) { + if (value.getType() != type) + return false; + } + return true; +} + +static bool haveSameTypes(ArrayRef lhs, ArrayRef rhs) { + if (lhs.size() != rhs.size()) + return false; + for (auto [lhsType, rhsType] : llvm::zip(lhs, rhs)) { + if (lhsType != rhsType) + return false; + } + return true; +} + +static Operation *createReturnLike(OpBuilder &builder, Location loc, + Operation *sampleReturn, + ValueRange operands) { + OperationState state(loc, sampleReturn->getName()); + state.addOperands(operands); + state.addAttributes(sampleReturn->getAttrs()); + return builder.create(state); +} + +static FailureOr> +collectReturnPathTypes(Block *block, SmallPtrSetImpl &visiting) { + if (!visiting.insert(block).second) + return block->getTerminator()->emitError() + << "unsupported cyclic terminal control flow"; + + Operation *term = block->getTerminator(); + if (auto br = dyn_cast(term)) { + FailureOr> result = + collectReturnPathTypes(br.getDest(), visiting); + visiting.erase(block); + return result; + } + + if (auto condBr = dyn_cast(term)) { + FailureOr nestedJoin = findNearestCommonBlock( + condBr.getTrueDest(), condBr.getFalseDest(), condBr.getLoc(), + /*emitDiagnostic=*/false); + if (succeeded(nestedJoin)) { + FailureOr> result = + collectReturnPathTypes(*nestedJoin, visiting); + visiting.erase(block); + return result; + } + + FailureOr> thenTypes = + collectReturnPathTypes(condBr.getTrueDest(), visiting); + FailureOr> elseTypes = + collectReturnPathTypes(condBr.getFalseDest(), visiting); + visiting.erase(block); + if (failed(thenTypes) || failed(elseTypes)) + return failure(); + if (!haveSameTypes(*thenTypes, *elseTypes)) { + condBr.emitError("terminal branch return types do not match"); + return failure(); + } + return *thenTypes; + } + + if (isSupportedReturn(term)) { + SmallVector types; + for (Value operand : term->getOperands()) + types.push_back(operand.getType()); + visiting.erase(block); + return types; + } + + visiting.erase(block); + return term->emitError() + << "unsupported terminator while analyzing terminal control flow"; +} + +static Operation *findReturnOnPath(Block *block, + SmallPtrSetImpl &visited) { + if (!visited.insert(block).second) + return nullptr; + + Operation *term = block->getTerminator(); + if (isSupportedReturn(term)) + return term; + for (Block *successor : getCfgSuccessors(block)) { + if (successor->getParent() != block->getParent()) + continue; + if (Operation *returnOp = findReturnOnPath(successor, visited)) + return returnOp; + } + return nullptr; +} + +static SmallVector mapValues(ValueRange values, IRMapping &mapping) { + SmallVector mapped; + mapped.reserve(values.size()); + for (Value value : values) + mapped.push_back(mapping.lookupOrDefault(value)); + return mapped; +} + +static FailureOr> +buildClonedTerminalPath(Block *block, ValueRange incoming, OpBuilder &builder, + IRMapping mapping, SmallPtrSetImpl &visiting); + +static FailureOr> +buildClonedTerminalTerminator(Operation *term, OpBuilder &builder, + IRMapping mapping, + SmallPtrSetImpl &visiting) { + if (auto br = dyn_cast(term)) { + SmallVector incoming = mapValues(br.getDestOperands(), mapping); + return buildClonedTerminalPath(br.getDest(), incoming, builder, mapping, + visiting); + } + + if (auto condBr = dyn_cast(term)) { + SmallPtrSet thenVisiting; + FailureOr> thenTypes = + collectReturnPathTypes(condBr.getTrueDest(), thenVisiting); + SmallPtrSet elseVisiting; + FailureOr> elseTypes = + collectReturnPathTypes(condBr.getFalseDest(), elseVisiting); + if (failed(thenTypes) || failed(elseTypes)) + return failure(); + if (!haveSameTypes(*thenTypes, *elseTypes)) { + condBr.emitError("terminal branch return types do not match"); + return failure(); + } + + auto ifOp = builder.create( + condBr.getLoc(), *thenTypes, + mapping.lookupOrDefault(condBr.getCondition()), + /*withElseRegion=*/true); + + { + OpBuilder::InsertionGuard guard(builder); + Operation *autoYield = + thenTypes->empty() ? ifOp.thenBlock()->getTerminator() : nullptr; + if (autoYield) + builder.setInsertionPoint(autoYield); + else + builder.setInsertionPointToStart(ifOp.thenBlock()); + SmallVector incoming = + mapValues(condBr.getTrueDestOperands(), mapping); + FailureOr> thenReturn = buildClonedTerminalPath( + condBr.getTrueDest(), incoming, builder, mapping, visiting); + if (failed(thenReturn)) + return failure(); + if (!haveSameTypes(*thenReturn, *thenTypes)) { + condBr.emitError("then terminal branch returns incompatible values"); + return failure(); + } + if (!autoYield) + builder.create(condBr.getLoc(), *thenReturn); + } + + { + OpBuilder::InsertionGuard guard(builder); + Operation *autoYield = + thenTypes->empty() ? ifOp.elseBlock()->getTerminator() : nullptr; + if (autoYield) + builder.setInsertionPoint(autoYield); + else + builder.setInsertionPointToStart(ifOp.elseBlock()); + SmallVector incoming = + mapValues(condBr.getFalseDestOperands(), mapping); + FailureOr> elseReturn = buildClonedTerminalPath( + condBr.getFalseDest(), incoming, builder, mapping, visiting); + if (failed(elseReturn)) + return failure(); + if (!haveSameTypes(*elseReturn, *thenTypes)) { + condBr.emitError("else terminal branch returns incompatible values"); + return failure(); + } + if (!autoYield) + builder.create(condBr.getLoc(), *elseReturn); + } + + return SmallVector(ifOp->getResults().begin(), + ifOp->getResults().end()); + } + + if (isSupportedReturn(term)) + return mapValues(term->getOperands(), mapping); + + return term->emitError() + << "unsupported terminator while structuring terminal control flow"; +} + +static FailureOr> +buildClonedTerminalPath(Block *block, ValueRange incoming, OpBuilder &builder, + IRMapping mapping, SmallPtrSetImpl &visiting) { + if (!visiting.insert(block).second) + return block->getTerminator()->emitError() + << "unsupported cyclic terminal control flow"; + + if (block->getNumArguments() != incoming.size()) { + visiting.erase(block); + return block->getTerminator()->emitError() + << "invalid branch operand count while structuring terminal " + "control flow"; + } + + for (auto [arg, value] : llvm::zip(block->getArguments(), incoming)) + mapping.map(arg, value); + + for (Operation &op : block->without_terminator()) + builder.clone(op, mapping); + + FailureOr> result = buildClonedTerminalTerminator( + block->getTerminator(), builder, mapping, visiting); + visiting.erase(block); + return result; +} + +static bool hasNonTreeCondBranch(Region &body) { + for (Block &block : body) { + auto condBr = dyn_cast(block.getTerminator()); + if (!condBr) + continue; + if (failed(findNearestCommonBlock(condBr.getTrueDest(), + condBr.getFalseDest(), condBr.getLoc(), + /*emitDiagnostic=*/false))) + return true; + } + return false; +} + +static LogicalResult structureTerminalReturnBody(Operation *funcOp, + Region &body) { + Block &entryBlock = body.front(); + Operation *entryTerm = entryBlock.getTerminator(); + SmallPtrSet visited; + Operation *sampleReturn = findReturnOnPath(&entryBlock, visited); + if (!sampleReturn) { + return funcOp->emitError() + << "unsupported non-tree control flow: no terminal return found"; + } + + OpBuilder builder(entryTerm); + IRMapping mapping; + SmallPtrSet visiting; + FailureOr> returnOperands = + buildClonedTerminalTerminator(entryTerm, builder, mapping, visiting); + if (failed(returnOperands)) + return failure(); + + createReturnLike(builder, entryTerm->getLoc(), sampleReturn, *returnOperands); + + SmallVector eraseBlocks; + for (Block &block : llvm::drop_begin(body.getBlocks())) + eraseBlocks.push_back(&block); + + entryTerm->erase(); + for (Block *block : eraseBlocks) { + for (Operation &op : *block) + op.dropAllReferences(); + } + for (Block *block : llvm::reverse(eraseBlocks)) + block->erase(); + + return success(); +} + +static FailureOr buildTerminalValueIf(cf::CondBranchOp condBr, + OpBuilder &builder) { + SmallPtrSet thenVisiting; + FailureOr> thenTypes = + collectReturnPathTypes(condBr.getTrueDest(), thenVisiting); + SmallPtrSet elseVisiting; + FailureOr> elseTypes = + collectReturnPathTypes(condBr.getFalseDest(), elseVisiting); + if (failed(thenTypes) || failed(elseTypes)) + return failure(); + if (!haveSameTypes(*thenTypes, *elseTypes)) { + condBr.emitError("terminal branch return types do not match"); + return failure(); + } + + auto ifOp = builder.create(condBr.getLoc(), *thenTypes, + condBr.getCondition(), + /*withElseRegion=*/true); + + { + OpBuilder::InsertionGuard branchGuard(builder); + Operation *autoYield = + thenTypes->empty() ? ifOp.thenBlock()->getTerminator() : nullptr; + if (autoYield) + builder.setInsertionPoint(autoYield); + else + builder.setInsertionPointToStart(ifOp.thenBlock()); + FailureOr thenReturn = buildReturnPath( + condBr.getTrueDest(), condBr.getTrueDestOperands(), builder); + if (failed(thenReturn)) + return failure(); + if (!haveSameTypes(thenReturn->operands, *thenTypes)) { + condBr.emitError("then terminal branch returns incompatible values"); + return failure(); + } + if (!autoYield) + builder.create(condBr.getLoc(), thenReturn->operands); + } + + { + OpBuilder::InsertionGuard branchGuard(builder); + Operation *autoYield = + thenTypes->empty() ? ifOp.elseBlock()->getTerminator() : nullptr; + if (autoYield) + builder.setInsertionPoint(autoYield); + else + builder.setInsertionPointToStart(ifOp.elseBlock()); + FailureOr elseReturn = buildReturnPath( + condBr.getFalseDest(), condBr.getFalseDestOperands(), builder); + if (failed(elseReturn)) + return failure(); + if (!haveSameTypes(elseReturn->operands, *thenTypes)) { + condBr.emitError("else terminal branch returns incompatible values"); + return failure(); + } + if (!autoYield) + builder.create(condBr.getLoc(), elseReturn->operands); + } + + return ifOp; +} + +static FailureOr +buildReturnPath(Block *block, ValueRange incoming, OpBuilder &builder) { + Operation *term = block->getTerminator(); + if (failed(replaceBlockArguments(block, incoming, term->getLoc()))) + return failure(); + moveBlockBodyBefore(block, builder); + + if (auto br = dyn_cast(term)) + return buildReturnPath(br.getDest(), br.getDestOperands(), builder); + + if (auto condBr = dyn_cast(term)) { + FailureOr nestedJoin = findNearestCommonBlock( + condBr.getTrueDest(), condBr.getFalseDest(), condBr.getLoc(), + /*emitDiagnostic=*/false); + if (succeeded(nestedJoin)) { + FailureOr nestedIf = + buildStructuredIf(condBr, *nestedJoin, builder); + if (failed(nestedIf)) + return failure(); + + SmallVector nestedResults((*nestedIf)->getResults().begin(), + (*nestedIf)->getResults().end()); + return buildReturnPath(*nestedJoin, nestedResults, builder); + } + + FailureOr terminalIf = buildTerminalValueIf(condBr, builder); + if (failed(terminalIf)) { + condBr.emitError() << "unsupported non-tree control flow: branch arms do " + "not both terminate with compatible returns"; + return failure(); + } + + ReturnPathResult result; + result.operands.assign((*terminalIf)->getResults().begin(), + (*terminalIf)->getResults().end()); + return result; + } + + if (isSupportedReturn(term)) { + ReturnPathResult result; + result.operands.assign(term->getOperands().begin(), + term->getOperands().end()); + return result; + } + + return term->emitError() + << "unsupported terminator while structuring terminal control flow"; +} + +static FailureOr> buildRegionPath(Block *block, + ValueRange incoming, + Block *stopBlock, + OpBuilder &builder) { + if (block == stopBlock) + return SmallVector(incoming.begin(), incoming.end()); + + Operation *term = block->getTerminator(); + if (failed(replaceBlockArguments(block, incoming, term->getLoc()))) + return failure(); + moveBlockBodyBefore(block, builder); + + if (auto br = dyn_cast(term)) { + SmallVector operands(br.getDestOperands().begin(), + br.getDestOperands().end()); + if (br.getDest() == stopBlock) + return operands; + return buildRegionPath(br.getDest(), operands, stopBlock, builder); + } + + if (auto condBr = dyn_cast(term)) { + FailureOr nestedJoin = findNearestCommonBlock( + condBr.getTrueDest(), condBr.getFalseDest(), condBr.getLoc()); + if (failed(nestedJoin)) + return failure(); + + FailureOr nestedIf = + buildStructuredIf(condBr, *nestedJoin, builder); + if (failed(nestedIf)) + return failure(); + + SmallVector nestedResults((*nestedIf)->getResults().begin(), + (*nestedIf)->getResults().end()); + if (*nestedJoin == stopBlock) + return nestedResults; + return buildRegionPath(*nestedJoin, nestedResults, stopBlock, builder); + } + + if (isSupportedReturn(term)) { + return term->emitError() + << "unsupported early return while structuring control flow"; + } + + return term->emitError() + << "unsupported terminator while structuring control flow"; +} + +static LogicalResult appendStructuredBlock(Block *block, ValueRange incoming, + OpBuilder &builder, + Operation *anchorTerminator); + +static LogicalResult appendStructuredTerminator(Operation *term, + OpBuilder &builder, + Operation *anchorTerminator) { + if (auto br = dyn_cast(term)) { + return appendStructuredBlock(br.getDest(), br.getDestOperands(), builder, + anchorTerminator); + } + + if (auto condBr = dyn_cast(term)) { + FailureOr joinBlock = findNearestCommonBlock( + condBr.getTrueDest(), condBr.getFalseDest(), condBr.getLoc(), + /*emitDiagnostic=*/false); + if (failed(joinBlock)) { + SmallPtrSet visited; + Operation *sampleReturn = findReturnOnPath(condBr.getTrueDest(), visited); + if (!sampleReturn) { + visited.clear(); + sampleReturn = findReturnOnPath(condBr.getFalseDest(), visited); + } + if (!sampleReturn) { + return condBr.emitError() + << "unsupported non-tree control flow: branch arms do not " + "reach a common convergence block"; + } + + FailureOr terminalIf = buildTerminalValueIf(condBr, builder); + if (failed(terminalIf)) + return failure(); + + SmallVector returnOperands((*terminalIf)->getResults().begin(), + (*terminalIf)->getResults().end()); + createReturnLike(builder, condBr.getLoc(), sampleReturn, returnOperands); + return success(); + } + + FailureOr ifOp = buildStructuredIf(condBr, *joinBlock, builder); + if (failed(ifOp)) + return failure(); + + return appendStructuredBlock(*joinBlock, (*ifOp)->getResults(), builder, + anchorTerminator); + } + + if (isSupportedReturn(term)) { + term->moveBefore(anchorTerminator); + return success(); + } + + return term->emitError() + << "unsupported entry terminator while structuring control flow"; +} + +static LogicalResult appendStructuredBlock(Block *block, ValueRange incoming, + OpBuilder &builder, + Operation *anchorTerminator) { + Operation *term = block->getTerminator(); + if (failed(replaceBlockArguments(block, incoming, term->getLoc()))) + return failure(); + + moveBlockBodyBefore(block, builder); + return appendStructuredTerminator(term, builder, anchorTerminator); +} + +static LogicalResult validateSupportedCfg(Region &body) { + for (Block &block : body) { + Operation *term = block.getTerminator(); + if (!isa(term) && !isSupportedReturn(term)) + return term->emitError() + << "unsupported terminator in multi-block function"; + } + return success(); +} + +static LogicalResult rejectCyclicCfg(Block *block, + SmallPtrSetImpl &visiting, + SmallPtrSetImpl &visited) { + if (visited.contains(block)) + return success(); + if (!visiting.insert(block).second) + return block->getTerminator()->emitError() + << "unsupported cyclic control flow in multi-block function"; + + for (Block *successor : getCfgSuccessors(block)) { + if (successor->getParent() == block->getParent() && + failed(rejectCyclicCfg(successor, visiting, visited))) + return failure(); + } + + visiting.erase(block); + visited.insert(block); + return success(); +} + +static LogicalResult structureFunctionBody(Operation *funcOp, Region &body) { + if (body.empty() || body.hasOneBlock()) + return success(); + + if (failed(validateSupportedCfg(body))) + return failure(); + + SmallPtrSet visiting; + SmallPtrSet visited; + if (failed(rejectCyclicCfg(&body.front(), visiting, visited))) + return failure(); + + if (hasNonTreeCondBranch(body)) + return structureTerminalReturnBody(funcOp, body); + + Block &entryBlock = body.front(); + Operation *entryTerm = entryBlock.getTerminator(); + if (isSupportedReturn(entryTerm)) { + return funcOp->emitError() + << "multi-block function entry cannot terminate with return"; + } + + SmallVector eraseBlocks; + for (Block &block : llvm::drop_begin(body.getBlocks())) + eraseBlocks.push_back(&block); + + OpBuilder builder(entryTerm); + if (failed(appendStructuredTerminator(entryTerm, builder, entryTerm))) + return failure(); + + entryTerm->erase(); + for (Block *block : eraseBlocks) { + for (Operation &op : *block) + op.dropAllReferences(); + } + for (Block *block : llvm::reverse(eraseBlocks)) + block->erase(); + + return success(); +} + +enum class PtrKind { Tensor, Block }; + +struct TensorPtrInfo { + Type resultType; + Value base; + Value offset; + bool scalarBase = false; +}; + +struct BlockPtrInfo { + Type resultType; + Value base; + SmallVector shape; + SmallVector strides; + SmallVector offsets; + DenseI32ArrayAttr order; +}; + +struct CFPtrInfo { + PtrKind kind; + TensorPtrInfo tensor; + BlockPtrInfo block; +}; + +struct RewriteEnv { + IRMapping valueMapping; + DenseMap pointerComponents; +}; + +struct LoopPointerInfo { + unsigned oldIndex = 0; + CFPtrInfo initInfo; + SmallVector newIndices; + SmallVector ivDeltas; +}; + +enum class IfComponentKind { + TensorOffset, + BlockShape, + BlockStride, + BlockOffset +}; + +struct IfComponent { + IfComponentKind kind; + unsigned dim = 0; + Type type; +}; + +struct IfPointerInfo { + unsigned oldIndex = 0; + CFPtrInfo thenInfo; + CFPtrInfo elseInfo; + SmallVector components; +}; + +static bool isTensorPointerType(Type type) { + auto tensorType = dyn_cast(type); + return tensorType && isa(tensorType.getElementType()); +} + +static bool isBlockPointerType(Type type) { + auto ptrType = dyn_cast(type); + return ptrType && isa(ptrType.getPointeeType()); +} + +static bool isControlFlowPointerType(Type type) { + return isTensorPointerType(type) || isBlockPointerType(type); +} + +static Value createZeroLike(OpBuilder &builder, Location loc, Type type) { + if (auto tensorType = dyn_cast(type)) { + auto elementType = dyn_cast(tensorType.getElementType()); + if (!elementType) + return nullptr; + auto attr = DenseElementsAttr::get(tensorType, + builder.getIntegerAttr(elementType, 0)); + return builder.create(loc, attr); + } + + if (type.isIndex()) + return builder.create(loc, 0); + + if (auto intType = dyn_cast(type)) + return builder.create(loc, 0, intType.getWidth()); + + return nullptr; +} + +static Value createZeroOffset(OpBuilder &builder, Location loc, Type ptrType) { + Type i32 = builder.getI32Type(); + if (auto tensorType = dyn_cast(ptrType)) + return createZeroLike(builder, loc, + RankedTensorType::get(tensorType.getShape(), i32)); + return createZeroLike(builder, loc, i32); +} + +static Value castIntegerLike(OpBuilder &builder, Location loc, Value value, + Type targetType) { + if (value.getType() == targetType) + return value; + + Type sourceType = value.getType(); + if ((sourceType.isIndex() && isa(targetType)) || + (isa(sourceType) && targetType.isIndex())) + return builder.create(loc, targetType, value); + + auto sourceInt = dyn_cast(sourceType); + auto targetInt = dyn_cast(targetType); + if (sourceInt && targetInt) { + if (sourceInt.getWidth() < targetInt.getWidth()) + return builder.create(loc, targetType, value); + if (sourceInt.getWidth() > targetInt.getWidth()) + return builder.create(loc, targetType, value); + return nullptr; + } + + auto sourceTensor = dyn_cast(sourceType); + auto targetTensor = dyn_cast(targetType); + if (!sourceTensor || !targetTensor || + sourceTensor.getShape() != targetTensor.getShape()) + return nullptr; + + auto sourceElement = dyn_cast(sourceTensor.getElementType()); + auto targetElement = dyn_cast(targetTensor.getElementType()); + if (!sourceElement || !targetElement) + return nullptr; + + if (sourceElement.getWidth() == targetElement.getWidth()) + return value; + if (sourceElement.getWidth() < targetElement.getWidth()) + return builder.create(loc, targetType, value); + return builder.create(loc, targetType, value); +} + +static FailureOr getWiderIntegerLikeType(Type lhs, Type rhs) { + if (lhs == rhs) + return lhs; + + if (lhs.isIndex() && rhs.isIndex()) + return lhs; + if (lhs.isIndex() && isa(rhs)) + return lhs; + if (isa(lhs) && rhs.isIndex()) + return rhs; + + auto lhsInt = dyn_cast(lhs); + auto rhsInt = dyn_cast(rhs); + if (lhsInt && rhsInt) + return lhsInt.getWidth() >= rhsInt.getWidth() ? lhs : rhs; + + auto lhsTensor = dyn_cast(lhs); + auto rhsTensor = dyn_cast(rhs); + if (!lhsTensor || !rhsTensor || lhsTensor.getShape() != rhsTensor.getShape()) + return failure(); + + Type lhsElement = lhsTensor.getElementType(); + Type rhsElement = rhsTensor.getElementType(); + if (lhsElement == rhsElement) + return lhs; + if (lhsElement.isIndex() && isa(rhsElement)) + return lhs; + if (isa(lhsElement) && rhsElement.isIndex()) + return rhs; + + auto lhsElementInt = dyn_cast(lhsElement); + auto rhsElementInt = dyn_cast(rhsElement); + if (!lhsElementInt || !rhsElementInt) + return failure(); + return lhsElementInt.getWidth() >= rhsElementInt.getWidth() ? lhs : rhs; +} + +static Value createAdd(OpBuilder &builder, Location loc, Value lhs, Value rhs) { + if (!lhs || !rhs) + return nullptr; + if (lhs.getType() != rhs.getType()) { + rhs = castIntegerLike(builder, loc, rhs, lhs.getType()); + if (!rhs) + return nullptr; + } + return builder.create(loc, lhs, rhs); +} + +static Value createAddWithWiderType(OpBuilder &builder, Location loc, Value lhs, + Value rhs) { + if (!lhs || !rhs) + return nullptr; + + FailureOr targetType = + getWiderIntegerLikeType(lhs.getType(), rhs.getType()); + if (failed(targetType)) + return nullptr; + + lhs = castIntegerLike(builder, loc, lhs, *targetType); + rhs = castIntegerLike(builder, loc, rhs, *targetType); + if (!lhs || !rhs) + return nullptr; + return builder.create(loc, lhs, rhs); +} + +static Value createMul(OpBuilder &builder, Location loc, Value lhs, Value rhs) { + if (!lhs || !rhs) + return nullptr; + if (lhs.getType() != rhs.getType()) { + rhs = castIntegerLike(builder, loc, rhs, lhs.getType()); + if (!rhs) + return nullptr; + } + return builder.create(loc, lhs, rhs); +} + +static Value remapValue(Value value, const RewriteEnv &env) { + if (Value mapped = env.valueMapping.lookupOrNull(value)) + return mapped; + return value; +} + +static FailureOr analyzePtr(Value value, const RewriteEnv &env, + OpBuilder &builder, Location loc); + +static FailureOr analyzeTensorPtr(Value value, + const RewriteEnv &env, + OpBuilder &builder, + Location loc) { + if (auto it = env.pointerComponents.find(value); + it != env.pointerComponents.end() && it->second.kind == PtrKind::Tensor) + return it->second.tensor; + + value = remapValue(value, env); + + if (auto addPtrOp = value.getDefiningOp()) { + FailureOr baseInfo = + analyzePtr(addPtrOp.getPtr(), env, builder, loc); + if (failed(baseInfo) || (*baseInfo).kind != PtrKind::Tensor) + return failure(); + + TensorPtrInfo tensor = (*baseInfo).tensor; + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPoint(addPtrOp); + Value offset = addPtrOp.getOffset(); + Value newOffset = createAddWithWiderType(builder, addPtrOp.getLoc(), + tensor.offset, offset); + if (!newOffset) + return failure(); + tensor.offset = newOffset; + tensor.resultType = value.getType(); + return tensor; + } + + if (auto splatOp = value.getDefiningOp()) { + if (!isa(splatOp.getSrc().getType())) + return failure(); + TensorPtrInfo parts; + parts.resultType = value.getType(); + parts.base = splatOp.getSrc(); + parts.scalarBase = true; + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPoint(splatOp); + parts.offset = createZeroOffset(builder, splatOp.getLoc(), value.getType()); + if (!parts.offset) + return failure(); + return parts; + } + + if (!isTensorPointerType(value.getType())) + return failure(); + + TensorPtrInfo parts; + parts.resultType = value.getType(); + parts.base = value; + parts.scalarBase = false; + OpBuilder::InsertionGuard guard(builder); + if (Operation *defOp = value.getDefiningOp()) + builder.setInsertionPointAfter(defOp); + else if (auto blockArg = dyn_cast(value)) + builder.setInsertionPointToStart(blockArg.getOwner()); + parts.offset = createZeroOffset(builder, loc, value.getType()); + if (!parts.offset) + return failure(); + return parts; +} + +static FailureOr analyzeBlockPtr(Value value, + const RewriteEnv &env, + OpBuilder &builder, + Location loc) { + if (auto it = env.pointerComponents.find(value); + it != env.pointerComponents.end() && it->second.kind == PtrKind::Block) + return it->second.block; + + value = remapValue(value, env); + + if (auto makePtrOp = value.getDefiningOp()) { + BlockPtrInfo parts; + parts.resultType = value.getType(); + parts.base = makePtrOp.getBase(); + parts.shape.assign(makePtrOp.getShape().begin(), + makePtrOp.getShape().end()); + parts.strides.assign(makePtrOp.getStrides().begin(), + makePtrOp.getStrides().end()); + parts.offsets.assign(makePtrOp.getOffsets().begin(), + makePtrOp.getOffsets().end()); + parts.order = makePtrOp.getOrderAttr(); + return parts; + } + + if (auto advanceOp = value.getDefiningOp()) { + FailureOr baseInfo = + analyzePtr(advanceOp.getPtr(), env, builder, loc); + if (failed(baseInfo) || (*baseInfo).kind != PtrKind::Block) + return failure(); + BlockPtrInfo block = (*baseInfo).block; + if (block.offsets.size() != advanceOp.getOffsets().size()) + return failure(); + + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPoint(advanceOp); + for (auto [idx, delta] : llvm::enumerate(advanceOp.getOffsets())) { + Value newOffset = + createAdd(builder, advanceOp.getLoc(), block.offsets[idx], delta); + if (!newOffset) + return failure(); + block.offsets[idx] = newOffset; + } + block.resultType = value.getType(); + return block; + } + + return failure(); +} + +static FailureOr analyzePtr(Value value, const RewriteEnv &env, + OpBuilder &builder, Location loc) { + if (auto it = env.pointerComponents.find(value); + it != env.pointerComponents.end()) + return it->second; + + value = remapValue(value, env); + + if (isBlockPointerType(value.getType())) { + FailureOr block = analyzeBlockPtr(value, env, builder, loc); + if (failed(block)) + return failure(); + CFPtrInfo info{PtrKind::Block}; + info.block = *block; + return info; + } + + if (isTensorPointerType(value.getType())) { + FailureOr tensor = + analyzeTensorPtr(value, env, builder, loc); + if (failed(tensor)) + return failure(); + CFPtrInfo info{PtrKind::Tensor}; + info.tensor = *tensor; + return info; + } + + return failure(); +} + +static Value rebuildTensorPtr(OpBuilder &builder, Location loc, + const TensorPtrInfo &parts, Value base, + Value offset) { + Value ptrBase = base; + if (parts.scalarBase && isTensorPointerType(parts.resultType)) + ptrBase = builder.create(loc, parts.resultType, base); + return builder.create(loc, parts.resultType, ptrBase, + offset); +} + +static Value rebuildBlockPtr(OpBuilder &builder, Location loc, + const BlockPtrInfo &parts, Value base, + ArrayRef shape, ArrayRef strides, + ArrayRef offsets) { + return builder.create( + loc, parts.resultType, base, ValueRange(shape), ValueRange(strides), + ValueRange(offsets), parts.order); +} + +static Value rebuildPtr(OpBuilder &builder, Location loc, + const CFPtrInfo &info) { + if (info.kind == PtrKind::Tensor) + return rebuildTensorPtr(builder, loc, info.tensor, info.tensor.base, + info.tensor.offset); + return rebuildBlockPtr(builder, loc, info.block, info.block.base, + info.block.shape, info.block.strides, + info.block.offsets); +} + +static void recordPointer(Value oldPtr, const CFPtrInfo &info, Value rebuiltPtr, + RewriteEnv &env) { + env.pointerComponents[oldPtr] = info; + env.valueMapping.map(oldPtr, rebuiltPtr); +} + +static SmallVector getLoopComponentValues(const CFPtrInfo &info) { + if (info.kind == PtrKind::Tensor) + return {info.tensor.offset}; + return info.block.offsets; +} + +static SmallVector getLoopComponentTypes(const CFPtrInfo &info) { + SmallVector types; + for (Value value : getLoopComponentValues(info)) + types.push_back(value.getType()); + return types; +} + +static CFPtrInfo withLoopComponentValues(CFPtrInfo info, + ArrayRef values) { + if (info.kind == PtrKind::Tensor) { + if (values.size() != 1) + return info; + info.tensor.offset = values[0]; + return info; + } + + if (values.size() == info.block.offsets.size()) + info.block.offsets.assign(values.begin(), values.end()); + return info; +} + +static bool areLoopCompatible(const CFPtrInfo &initInfo, + const CFPtrInfo &nextInfo) { + if (initInfo.kind != nextInfo.kind) + return false; + if (initInfo.kind == PtrKind::Tensor) + return initInfo.tensor.resultType == nextInfo.tensor.resultType && + initInfo.tensor.base == nextInfo.tensor.base && + initInfo.tensor.scalarBase == nextInfo.tensor.scalarBase && + haveSameTypes(getLoopComponentValues(initInfo), + getLoopComponentValues(nextInfo)); + + return initInfo.block.resultType == nextInfo.block.resultType && + initInfo.block.base == nextInfo.block.base && + initInfo.block.order == nextInfo.block.order && + initInfo.block.shape == nextInfo.block.shape && + initInfo.block.strides == nextInfo.block.strides && + initInfo.block.offsets.size() == nextInfo.block.offsets.size() && + haveSameTypes(getLoopComponentValues(initInfo), + getLoopComponentValues(nextInfo)); +} + +static bool isScalarIntegerLike(Type type) { + return type.isIndex() || isa(type); +} + +static bool isConstantIndex(Value value, int64_t expected) { + auto constOp = value.getDefiningOp(); + return constOp && constOp.value() == expected; +} + +static bool isRangeFromZeroByOne(scf::ForOp forOp) { + return isConstantIndex(forOp.getLowerBound(), 0) && + isConstantIndex(forOp.getStep(), 1); +} + +static bool isDefinedOutside(Operation *scope, Value value) { + if (Operation *defOp = value.getDefiningOp()) + return !scope->isAncestor(defOp); + + auto blockArg = dyn_cast(value); + if (!blockArg) + return false; + + Operation *ownerOp = blockArg.getOwner()->getParentOp(); + return ownerOp != scope && (!ownerOp || !scope->isAncestor(ownerOp)); +} + +static FailureOr> +matchSimpleForIvDeltas(scf::ForOp forOp, const LoopPointerInfo &info, + Value yieldOperand) { + if (!isRangeFromZeroByOne(forOp) || info.initInfo.kind != PtrKind::Block) + return failure(); + + auto advanceOp = yieldOperand.getDefiningOp(); + if (!advanceOp || + advanceOp.getPtr() != forOp.getRegionIterArgs()[info.oldIndex]) + return failure(); + + if (advanceOp.getOffsets().size() != info.initInfo.block.offsets.size()) + return failure(); + + SmallVector deltas; + for (auto [initOffset, delta] : + llvm::zip(info.initInfo.block.offsets, advanceOp.getOffsets())) { + if (!isScalarIntegerLike(initOffset.getType()) || + !isScalarIntegerLike(delta.getType()) || + !isDefinedOutside(forOp, delta)) + return failure(); + deltas.push_back(delta); + } + return deltas; +} + +static FailureOr +withForIvClosedFormComponents(const LoopPointerInfo &info, Value iv, + OpBuilder &builder, Location loc, + const RewriteEnv &env) { + if (info.ivDeltas.empty()) + return failure(); + + SmallVector initComponents = getLoopComponentValues(info.initInfo); + if (initComponents.size() != info.ivDeltas.size()) + return failure(); + + SmallVector components; + components.reserve(initComponents.size()); + for (auto [initComponent, delta] : llvm::zip(initComponents, info.ivDeltas)) { + Type componentType = initComponent.getType(); + if (!isScalarIntegerLike(componentType)) + return failure(); + + Value typedIv = castIntegerLike(builder, loc, iv, componentType); + Value typedDelta = + castIntegerLike(builder, loc, remapValue(delta, env), componentType); + if (!typedIv || !typedDelta) + return failure(); + + Value scaledDelta = createMul(builder, loc, typedIv, typedDelta); + Value component = createAdd(builder, loc, initComponent, scaledDelta); + if (!scaledDelta || !component) + return failure(); + components.push_back(component); + } + + return withLoopComponentValues(info.initInfo, components); +} + +static LoopPointerInfo *findLoopInfo(SmallVectorImpl &infos, + unsigned oldIndex) { + for (LoopPointerInfo &info : infos) { + if (info.oldIndex == oldIndex) + return &info; + } + return nullptr; +} + +static const LoopPointerInfo *findLoopInfo(ArrayRef infos, + unsigned oldIndex) { + for (const LoopPointerInfo &info : infos) { + if (info.oldIndex == oldIndex) + return &info; + } + return nullptr; +} + +static SmallVector collectForComponents(const LoopPointerInfo &info, + scf::ForOp forOp, + bool useResults) { + SmallVector values; + for (unsigned newIndex : info.newIndices) + values.push_back(useResults ? forOp.getResult(newIndex) + : forOp.getRegionIterArgs()[newIndex]); + return values; +} + +static SmallVector collectWhileComponents(const LoopPointerInfo &info, + scf::WhileOp whileOp, + bool useResults, + bool useAfterArgs) { + SmallVector values; + for (unsigned newIndex : info.newIndices) { + if (useResults) + values.push_back(whileOp.getResult(newIndex)); + else if (useAfterArgs) + values.push_back(whileOp.getAfterArguments()[newIndex]); + else + values.push_back(whileOp.getBeforeArguments()[newIndex]); + } + return values; +} + +static LogicalResult rewriteControlFlowOp(Operation *op, OpBuilder &builder, + RewriteEnv &env); + +static LogicalResult materializePointerResult(Operation &bodyOp, + Operation *clonedOp, + OpBuilder &builder, + RewriteEnv &env) { + if (!isa(bodyOp)) + return success(); + + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointAfter(clonedOp); + + for (auto [oldResult, clonedResult] : + llvm::zip(bodyOp.getResults(), clonedOp->getResults())) { + if (!isControlFlowPointerType(oldResult.getType())) + continue; + + FailureOr info = + analyzePtr(clonedResult, env, builder, oldResult.getLoc()); + if (failed(info)) + continue; + + Value rebuilt = rebuildPtr(builder, oldResult.getLoc(), *info); + if (!rebuilt) + return failure(); + recordPointer(oldResult, *info, rebuilt, env); + } + + return success(); +} + +static LogicalResult rewriteBodyOps(Block *oldBlock, OpBuilder &builder, + RewriteEnv &env) { + for (Operation &bodyOp : oldBlock->without_terminator()) { + if (isa(bodyOp) && + succeeded(rewriteControlFlowOp(&bodyOp, builder, env))) + continue; + Operation *clonedOp = builder.clone(bodyOp, env.valueMapping); + if (failed(materializePointerResult(bodyOp, clonedOp, builder, env))) + return failure(); + } + return success(); +} + +static LogicalResult rewriteForOp(scf::ForOp forOp, OpBuilder &builder, + RewriteEnv &env) { + auto yieldOp = cast(forOp.getBody()->getTerminator()); + SmallVector pointerInfos; + + OpBuilder analysisBuilder(forOp.getContext()); + analysisBuilder.setInsertionPoint(forOp); + + for (auto [idx, iterArg] : llvm::enumerate(forOp.getRegionIterArgs())) { + if (!isControlFlowPointerType(iterArg.getType())) + continue; + if (idx >= forOp.getInitArgs().size() || idx >= yieldOp.getNumOperands()) + return failure(); + + FailureOr initInfo = analyzePtr(forOp.getInitArgs()[idx], env, + analysisBuilder, forOp.getLoc()); + if (failed(initInfo)) + continue; + pointerInfos.push_back( + LoopPointerInfo{static_cast(idx), *initInfo, {}}); + } + + if (pointerInfos.empty()) + return failure(); + + for (LoopPointerInfo &info : pointerInfos) { + FailureOr> deltas = + matchSimpleForIvDeltas(forOp, info, yieldOp.getOperand(info.oldIndex)); + if (succeeded(deltas)) + info.ivDeltas = *deltas; + } + + SmallVector newInitArgs; + SmallVector oldToNewStart(forOp.getInitArgs().size(), 0); + for (auto [idx, initArg] : llvm::enumerate(forOp.getInitArgs())) { + oldToNewStart[idx] = newInitArgs.size(); + if (LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + for (Value component : getLoopComponentValues(info->initInfo)) { + info->newIndices.push_back(newInitArgs.size()); + newInitArgs.push_back(component); + } + continue; + } + newInitArgs.push_back(remapValue(initArg, env)); + } + + bool bodyOk = true; + auto newForOp = builder.create( + forOp.getLoc(), remapValue(forOp.getLowerBound(), env), + remapValue(forOp.getUpperBound(), env), remapValue(forOp.getStep(), env), + newInitArgs, + [&](OpBuilder &bodyBuilder, Location loc, Value iv, ValueRange args) { + RewriteEnv bodyEnv = env; + bodyEnv.valueMapping.map(forOp.getInductionVar(), iv); + + for (auto [idx, oldArg] : llvm::enumerate(forOp.getRegionIterArgs())) { + if (const LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + SmallVector values; + for (unsigned newIndex : info->newIndices) + values.push_back(args[newIndex]); + CFPtrInfo argInfo = withLoopComponentValues(info->initInfo, values); + FailureOr closedFormInfo = withForIvClosedFormComponents( + *info, iv, bodyBuilder, loc, bodyEnv); + if (succeeded(closedFormInfo)) + argInfo = *closedFormInfo; + Value rebuilt = rebuildPtr(bodyBuilder, loc, argInfo); + if (!rebuilt) { + bodyOk = false; + continue; + } + recordPointer(oldArg, argInfo, rebuilt, bodyEnv); + continue; + } + bodyEnv.valueMapping.map(oldArg, args[oldToNewStart[idx]]); + } + + if (failed(rewriteBodyOps(forOp.getBody(), bodyBuilder, bodyEnv))) + bodyOk = false; + + SmallVector newYieldOperands; + for (auto [idx, oldOperand] : llvm::enumerate(yieldOp.getOperands())) { + if (const LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + FailureOr nextInfo = + analyzePtr(oldOperand, bodyEnv, bodyBuilder, yieldOp.getLoc()); + if (failed(nextInfo) || + !areLoopCompatible(info->initInfo, *nextInfo)) { + bodyOk = false; + for (unsigned newIndex : info->newIndices) + newYieldOperands.push_back(args[newIndex]); + continue; + } + for (Value component : getLoopComponentValues(*nextInfo)) + newYieldOperands.push_back(component); + continue; + } + newYieldOperands.push_back(remapValue(oldOperand, bodyEnv)); + } + + bodyBuilder.create(yieldOp.getLoc(), newYieldOperands); + }); + newForOp->setAttrs(forOp->getAttrs()); + + if (!bodyOk) { + newForOp.erase(); + return failure(); + } + + builder.setInsertionPointAfter(newForOp); + for (auto [idx, oldResult] : llvm::enumerate(forOp.getResults())) { + if (const LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + CFPtrInfo resultInfo = withLoopComponentValues( + info->initInfo, + collectForComponents(*info, newForOp, /*useResults=*/true)); + Value rebuilt = rebuildPtr(builder, oldResult.getLoc(), resultInfo); + if (!rebuilt) { + newForOp.erase(); + return failure(); + } + recordPointer(oldResult, resultInfo, rebuilt, env); + continue; + } + env.valueMapping.map(oldResult, newForOp.getResult(oldToNewStart[idx])); + } + + return success(); +} + +static LogicalResult rewriteWhileOp(scf::WhileOp whileOp, OpBuilder &builder, + RewriteEnv &env) { + scf::ConditionOp conditionOp = whileOp.getConditionOp(); + scf::YieldOp yieldOp = whileOp.getYieldOp(); + SmallVector pointerInfos; + + OpBuilder analysisBuilder(whileOp.getContext()); + analysisBuilder.setInsertionPoint(whileOp); + + for (auto [idx, beforeArg] : llvm::enumerate(whileOp.getBeforeArguments())) { + if (!isControlFlowPointerType(beforeArg.getType())) + continue; + if (idx >= whileOp.getInits().size() || + idx >= conditionOp.getArgs().size() || idx >= yieldOp.getNumOperands()) + return failure(); + + FailureOr initInfo = analyzePtr( + whileOp.getInits()[idx], env, analysisBuilder, whileOp.getLoc()); + if (failed(initInfo)) + continue; + pointerInfos.push_back( + LoopPointerInfo{static_cast(idx), *initInfo, {}}); + } + + if (pointerInfos.empty()) + return failure(); + + SmallVector newInits; + SmallVector newResultTypes; + SmallVector oldToNewStart(whileOp.getInits().size(), 0); + for (auto [idx, initArg] : llvm::enumerate(whileOp.getInits())) { + oldToNewStart[idx] = newInits.size(); + if (LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + for (Value component : getLoopComponentValues(info->initInfo)) { + info->newIndices.push_back(newInits.size()); + newInits.push_back(component); + newResultTypes.push_back(component.getType()); + } + continue; + } + newInits.push_back(remapValue(initArg, env)); + newResultTypes.push_back(whileOp.getResult(idx).getType()); + } + + bool bodyOk = true; + auto newWhileOp = builder.create( + whileOp.getLoc(), newResultTypes, newInits, + [&](OpBuilder &bodyBuilder, Location loc, ValueRange args) { + RewriteEnv beforeEnv = env; + for (auto [idx, oldArg] : + llvm::enumerate(whileOp.getBeforeArguments())) { + if (const LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + SmallVector values; + for (unsigned newIndex : info->newIndices) + values.push_back(args[newIndex]); + CFPtrInfo argInfo = withLoopComponentValues(info->initInfo, values); + Value rebuilt = rebuildPtr(bodyBuilder, loc, argInfo); + if (!rebuilt) { + bodyOk = false; + continue; + } + recordPointer(oldArg, argInfo, rebuilt, beforeEnv); + continue; + } + beforeEnv.valueMapping.map(oldArg, args[oldToNewStart[idx]]); + } + + if (failed(rewriteBodyOps(whileOp.getBeforeBody(), bodyBuilder, + beforeEnv))) + bodyOk = false; + + SmallVector newConditionArgs; + for (auto [idx, oldArg] : llvm::enumerate(conditionOp.getArgs())) { + if (const LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + FailureOr conditionInfo = analyzePtr( + oldArg, beforeEnv, bodyBuilder, conditionOp.getLoc()); + if (failed(conditionInfo) || + !areLoopCompatible(info->initInfo, *conditionInfo)) { + bodyOk = false; + for (unsigned newIndex : info->newIndices) + newConditionArgs.push_back(args[newIndex]); + continue; + } + for (Value component : getLoopComponentValues(*conditionInfo)) + newConditionArgs.push_back(component); + continue; + } + newConditionArgs.push_back(remapValue(oldArg, beforeEnv)); + } + + bodyBuilder.create( + conditionOp.getLoc(), + remapValue(conditionOp.getCondition(), beforeEnv), + newConditionArgs); + }, + [&](OpBuilder &bodyBuilder, Location loc, ValueRange args) { + RewriteEnv afterEnv = env; + for (auto [idx, oldArg] : + llvm::enumerate(whileOp.getAfterArguments())) { + if (const LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + SmallVector values; + for (unsigned newIndex : info->newIndices) + values.push_back(args[newIndex]); + CFPtrInfo argInfo = withLoopComponentValues(info->initInfo, values); + Value rebuilt = rebuildPtr(bodyBuilder, loc, argInfo); + if (!rebuilt) { + bodyOk = false; + continue; + } + recordPointer(oldArg, argInfo, rebuilt, afterEnv); + continue; + } + afterEnv.valueMapping.map(oldArg, args[oldToNewStart[idx]]); + } + + if (failed( + rewriteBodyOps(whileOp.getAfterBody(), bodyBuilder, afterEnv))) + bodyOk = false; + + SmallVector newYieldOperands; + for (auto [idx, oldOperand] : llvm::enumerate(yieldOp.getOperands())) { + if (const LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + FailureOr nextInfo = + analyzePtr(oldOperand, afterEnv, bodyBuilder, yieldOp.getLoc()); + if (failed(nextInfo) || + !areLoopCompatible(info->initInfo, *nextInfo)) { + bodyOk = false; + for (unsigned newIndex : info->newIndices) + newYieldOperands.push_back(args[newIndex]); + continue; + } + for (Value component : getLoopComponentValues(*nextInfo)) + newYieldOperands.push_back(component); + continue; + } + newYieldOperands.push_back(remapValue(oldOperand, afterEnv)); + } + + bodyBuilder.create(yieldOp.getLoc(), newYieldOperands); + }); + newWhileOp->setAttrs(whileOp->getAttrs()); + + if (!bodyOk) { + newWhileOp.erase(); + return failure(); + } + + builder.setInsertionPointAfter(newWhileOp); + for (auto [idx, oldResult] : llvm::enumerate(whileOp.getResults())) { + if (const LoopPointerInfo *info = findLoopInfo(pointerInfos, idx)) { + CFPtrInfo resultInfo = withLoopComponentValues( + info->initInfo, + collectWhileComponents(*info, newWhileOp, /*useResults=*/true, + /*useAfterArgs=*/false)); + Value rebuilt = rebuildPtr(builder, oldResult.getLoc(), resultInfo); + if (!rebuilt) { + newWhileOp.erase(); + return failure(); + } + recordPointer(oldResult, resultInfo, rebuilt, env); + continue; + } + env.valueMapping.map(oldResult, newWhileOp.getResult(oldToNewStart[idx])); + } + + return success(); +} + +static LogicalResult +addIfTensorComponents(const TensorPtrInfo &thenParts, + const TensorPtrInfo &elseParts, + SmallVectorImpl &components) { + if (thenParts.base != elseParts.base) + return failure(); + if (thenParts.offset.getType() != elseParts.offset.getType()) + return failure(); + components.push_back( + {IfComponentKind::TensorOffset, 0, thenParts.offset.getType()}); + return success(); +} + +static LogicalResult +addIfBlockComponents(const BlockPtrInfo &thenParts, + const BlockPtrInfo &elseParts, + SmallVectorImpl &components) { + if (thenParts.resultType != elseParts.resultType || + thenParts.order != elseParts.order || + thenParts.shape.size() != elseParts.shape.size() || + thenParts.strides.size() != elseParts.strides.size() || + thenParts.offsets.size() != elseParts.offsets.size()) + return failure(); + + if (thenParts.base != elseParts.base) + return failure(); + + for (auto [idx, values] : + llvm::enumerate(llvm::zip(thenParts.shape, elseParts.shape))) { + Value thenValue = std::get<0>(values); + Value elseValue = std::get<1>(values); + if (thenValue == elseValue) + continue; + if (thenValue.getType() != elseValue.getType()) + return failure(); + components.push_back({IfComponentKind::BlockShape, + static_cast(idx), thenValue.getType()}); + } + + for (auto [idx, values] : + llvm::enumerate(llvm::zip(thenParts.strides, elseParts.strides))) { + Value thenValue = std::get<0>(values); + Value elseValue = std::get<1>(values); + if (thenValue == elseValue) + continue; + if (thenValue.getType() != elseValue.getType()) + return failure(); + components.push_back({IfComponentKind::BlockStride, + static_cast(idx), thenValue.getType()}); + } + + for (auto [idx, values] : + llvm::enumerate(llvm::zip(thenParts.offsets, elseParts.offsets))) { + Value thenValue = std::get<0>(values); + Value elseValue = std::get<1>(values); + if (thenValue == elseValue) + continue; + if (thenValue.getType() != elseValue.getType()) + return failure(); + components.push_back({IfComponentKind::BlockOffset, + static_cast(idx), thenValue.getType()}); + } + return success(); +} + +static FailureOr analyzePtrForIfPlanning(Value value, + const RewriteEnv &env, + OpBuilder &builder, + Location loc); + +static FailureOr +analyzeNestedIfResultForPlanning(scf::IfOp ifOp, unsigned resultIndex, + const RewriteEnv &env, OpBuilder &builder, + Location loc) { + if (!ifOp.elseBlock() || resultIndex >= ifOp.getNumResults()) + return failure(); + + scf::YieldOp thenYield = ifOp.thenYield(); + scf::YieldOp elseYield = ifOp.elseYield(); + if (resultIndex >= thenYield.getNumOperands() || + resultIndex >= elseYield.getNumOperands()) + return failure(); + + FailureOr thenInfo = analyzePtrForIfPlanning( + thenYield.getOperand(resultIndex), env, builder, loc); + FailureOr elseInfo = analyzePtrForIfPlanning( + elseYield.getOperand(resultIndex), env, builder, loc); + if (failed(thenInfo) || failed(elseInfo) || thenInfo->kind != elseInfo->kind) + return failure(); + + SmallVector components; + if (thenInfo->kind == PtrKind::Tensor) { + if (thenInfo->tensor.resultType != elseInfo->tensor.resultType || + thenInfo->tensor.scalarBase != elseInfo->tensor.scalarBase) + return failure(); + if (failed(addIfTensorComponents(thenInfo->tensor, elseInfo->tensor, + components))) + return failure(); + } else if (failed(addIfBlockComponents(thenInfo->block, elseInfo->block, + components))) { + return failure(); + } + + return *thenInfo; +} + +static FailureOr analyzePtrForIfPlanning(Value value, + const RewriteEnv &env, + OpBuilder &builder, + Location loc) { + if (auto it = env.pointerComponents.find(value); + it != env.pointerComponents.end()) + return it->second; + + Value mapped = remapValue(value, env); + if (auto result = dyn_cast(mapped)) { + if (auto nestedIf = dyn_cast(result.getOwner())) { + FailureOr nestedInfo = analyzeNestedIfResultForPlanning( + nestedIf, result.getResultNumber(), env, builder, loc); + if (succeeded(nestedInfo)) + return nestedInfo; + } + } + + return analyzePtr(value, env, builder, loc); +} + +static Value getComponentValue(const CFPtrInfo &info, + const IfComponent &component) { + if (info.kind == PtrKind::Tensor) { + if (component.kind == IfComponentKind::TensorOffset) + return info.tensor.offset; + return nullptr; + } + + switch (component.kind) { + case IfComponentKind::BlockShape: + return info.block.shape[component.dim]; + case IfComponentKind::BlockStride: + return info.block.strides[component.dim]; + case IfComponentKind::BlockOffset: + return info.block.offsets[component.dim]; + default: + return nullptr; + } +} + +static const IfPointerInfo *findIfInfo(ArrayRef infos, + unsigned oldIndex) { + for (const IfPointerInfo &info : infos) { + if (info.oldIndex == oldIndex) + return &info; + } + return nullptr; +} + +static FailureOr makeIfResultInfo(const IfPointerInfo &info, + ArrayRef componentValues) { + unsigned componentIndex = 0; + CFPtrInfo resultInfo = info.thenInfo; + if (info.thenInfo.kind == PtrKind::Tensor) { + for (const IfComponent &component : info.components) { + Value value = componentValues[componentIndex++]; + switch (component.kind) { + case IfComponentKind::TensorOffset: + resultInfo.tensor.offset = value; + break; + default: + return failure(); + } + } + return resultInfo; + } + + for (const IfComponent &component : info.components) { + Value value = componentValues[componentIndex++]; + switch (component.kind) { + case IfComponentKind::BlockShape: + resultInfo.block.shape[component.dim] = value; + break; + case IfComponentKind::BlockStride: + resultInfo.block.strides[component.dim] = value; + break; + case IfComponentKind::BlockOffset: + resultInfo.block.offsets[component.dim] = value; + break; + default: + return failure(); + } + } + return resultInfo; +} + +static LogicalResult rewriteIfOp(scf::IfOp ifOp, OpBuilder &builder, + RewriteEnv &env) { + if (!ifOp.elseBlock() || ifOp->getNumResults() == 0) + return failure(); + + scf::YieldOp thenYield = ifOp.thenYield(); + scf::YieldOp elseYield = ifOp.elseYield(); + SmallVector pointerInfos; + + OpBuilder analysisBuilder(ifOp.getContext()); + analysisBuilder.setInsertionPoint(ifOp); + + for (auto [idx, result] : llvm::enumerate(ifOp.getResults())) { + if (!isControlFlowPointerType(result.getType())) + continue; + if (thenYield.getOperand(idx) == elseYield.getOperand(idx)) + continue; + + FailureOr thenInfo = analyzePtrForIfPlanning( + thenYield.getOperand(idx), env, analysisBuilder, thenYield.getLoc()); + FailureOr elseInfo = analyzePtrForIfPlanning( + elseYield.getOperand(idx), env, analysisBuilder, elseYield.getLoc()); + if (failed(thenInfo) || failed(elseInfo) || + (*thenInfo).kind != (*elseInfo).kind) + continue; + + IfPointerInfo info; + info.oldIndex = idx; + info.thenInfo = *thenInfo; + info.elseInfo = *elseInfo; + if (info.thenInfo.kind == PtrKind::Tensor) { + if (info.thenInfo.tensor.resultType != info.elseInfo.tensor.resultType || + info.thenInfo.tensor.scalarBase != info.elseInfo.tensor.scalarBase) + continue; + if (failed(addIfTensorComponents(info.thenInfo.tensor, + info.elseInfo.tensor, info.components))) + continue; + } else { + if (failed(addIfBlockComponents(info.thenInfo.block, info.elseInfo.block, + info.components))) + continue; + } + pointerInfos.push_back(info); + } + + if (pointerInfos.empty()) + return failure(); + + SmallVector newResultTypes; + for (auto [idx, result] : llvm::enumerate(ifOp.getResults())) { + if (const IfPointerInfo *info = findIfInfo(pointerInfos, idx)) { + for (const IfComponent &component : info->components) + newResultTypes.push_back(component.type); + continue; + } + newResultTypes.push_back(result.getType()); + } + + bool bodyOk = true; + auto buildBranch = [&](OpBuilder &branchBuilder, Location loc, + bool isThen) -> LogicalResult { + RewriteEnv branchEnv = env; + Block *oldBlock = isThen ? ifOp.thenBlock() : ifOp.elseBlock(); + scf::YieldOp oldYield = isThen ? thenYield : elseYield; + if (failed(rewriteBodyOps(oldBlock, branchBuilder, branchEnv))) + return failure(); + + SmallVector newYieldOperands; + for (auto [idx, oldOperand] : llvm::enumerate(oldYield.getOperands())) { + if (const IfPointerInfo *info = findIfInfo(pointerInfos, idx)) { + FailureOr branchInfo = + analyzePtr(oldOperand, branchEnv, branchBuilder, oldYield.getLoc()); + if (failed(branchInfo) || branchInfo->kind != info->thenInfo.kind) + return failure(); + for (const IfComponent &component : info->components) { + Value value = getComponentValue(*branchInfo, component); + if (!value || value.getType() != component.type) + return failure(); + newYieldOperands.push_back(value); + } + continue; + } + newYieldOperands.push_back(remapValue(oldOperand, branchEnv)); + } + branchBuilder.create(oldYield.getLoc(), newYieldOperands); + return success(); + }; + + auto newIfOp = + builder.create(ifOp.getLoc(), newResultTypes, + remapValue(ifOp.getCondition(), env), true); + newIfOp->setAttrs(ifOp->getAttrs()); + + { + OpBuilder::InsertionGuard guard(builder); + if (newResultTypes.empty()) { + newIfOp.thenBlock()->getTerminator()->erase(); + builder.setInsertionPointToEnd(newIfOp.thenBlock()); + } else { + builder.setInsertionPointToStart(newIfOp.thenBlock()); + } + if (failed(buildBranch(builder, ifOp.getLoc(), /*isThen=*/true))) + bodyOk = false; + } + { + OpBuilder::InsertionGuard guard(builder); + if (newResultTypes.empty()) { + newIfOp.elseBlock()->getTerminator()->erase(); + builder.setInsertionPointToEnd(newIfOp.elseBlock()); + } else { + builder.setInsertionPointToStart(newIfOp.elseBlock()); + } + if (failed(buildBranch(builder, ifOp.getLoc(), /*isThen=*/false))) + bodyOk = false; + } + + if (!bodyOk) { + newIfOp.erase(); + return failure(); + } + + builder.setInsertionPointAfter(newIfOp); + unsigned newResultIndex = 0; + for (auto [idx, oldResult] : llvm::enumerate(ifOp.getResults())) { + if (const IfPointerInfo *info = findIfInfo(pointerInfos, idx)) { + SmallVector componentValues; + for (unsigned i = 0; i < info->components.size(); ++i) + componentValues.push_back(newIfOp.getResult(newResultIndex++)); + FailureOr resultInfo = + makeIfResultInfo(*info, componentValues); + if (failed(resultInfo)) { + newIfOp.erase(); + return failure(); + } + Value rebuilt = rebuildPtr(builder, oldResult.getLoc(), *resultInfo); + if (!rebuilt) { + newIfOp.erase(); + return failure(); + } + recordPointer(oldResult, *resultInfo, rebuilt, env); + continue; + } + env.valueMapping.map(oldResult, newIfOp.getResult(newResultIndex++)); + } + + return success(); +} + +static LogicalResult rewriteControlFlowOp(Operation *op, OpBuilder &builder, + RewriteEnv &env) { + if (auto forOp = dyn_cast(op)) + return rewriteForOp(forOp, builder, env); + if (auto whileOp = dyn_cast(op)) + return rewriteWhileOp(whileOp, builder, env); + if (auto ifOp = dyn_cast(op)) + return rewriteIfOp(ifOp, builder, env); + return failure(); +} + +static SmallVector collectReplacements(Operation *op, + const RewriteEnv &env) { + SmallVector replacements; + replacements.reserve(op->getNumResults()); + for (Value result : op->getResults()) + replacements.push_back(remapValue(result, env)); + return replacements; +} + +static LogicalResult tryDecoupleControlFlowOp(Operation *op, + IRRewriter &rewriter) { + RewriteEnv env; + rewriter.setInsertionPoint(op); + if (failed(rewriteControlFlowOp(op, rewriter, env))) + return failure(); + + SmallVector replacements = collectReplacements(op, env); + if (replacements.size() != op->getNumResults() || + llvm::any_of(replacements, [](Value value) { return !value; })) + return failure(); + rewriter.replaceOp(op, replacements); + return success(); +} + +} // namespace + namespace mlir::triton { void TritonControlFlowOptPass::getDependentDialects( DialectRegistry ®istry) const { - // CFG structuring creates SCF operations, while pointer decomposition - // materializes arith and Triton pointer operations. registry.insert(); } void TritonControlFlowOptPass::runOnOperation() { - ModuleOp module = getOperation(); - - // Apply the control-flow preprocessing pipeline in dependency order: - // 1. normalize supported cf graphs to scf; - // 2. decompose block-pointer descriptors across SCF boundaries; - // 3. replace common-base tensor pointers at SCF boundaries with offsets. - // TODO: Extend stage 3 to carry both the base and complete offsets, then add - // StructuredOffsetsDecompose to further split structured offsets into a - // base offset and per-dimension strides. - // Keep these calls explicit so each transformation has one owner, can be - // tested independently and can be implemented without changing pass - // registration. - if (failed(controlflow::structureCFG(module)) || - failed(controlflow::runBlockPtrDecompose(module)) || - failed(controlflow::runTensorPtrDecompose(module)) || - failed(verify(module))) + ModuleOp moduleOp = getOperation(); + SmallVector funcs; + moduleOp.walk([&](Operation *op) { + if (isa(op)) + funcs.push_back(op); + }); + + for (Operation *op : funcs) { + if (op->getParentOp() == nullptr) + continue; + if (auto funcOp = dyn_cast(op)) { + if (!funcOp.isDeclaration() && + failed(structureFunctionBody(funcOp, funcOp.getBody()))) { + signalPassFailure(); + return; + } + continue; + } + + if (auto funcOp = dyn_cast(op)) { + if (!funcOp.isDeclaration() && + failed(structureFunctionBody(funcOp, funcOp.getBody()))) { + signalPassFailure(); + return; + } + continue; + } + + if (auto mapOp = dyn_cast(op)) { + if (failed(structureFunctionBody(mapOp, mapOp.getRegion()))) { + signalPassFailure(); + return; + } + } + } + + SmallVector targets; + moduleOp.walk([&](Operation *op) { + if (isa(op)) + targets.push_back(op); + }); + + LLVM_DEBUG({ + llvm::dbgs() << "TritonControlFlowOpt collected " << targets.size() + << " structured control-flow targets for later decoupling\n"; + }); + + IRRewriter rewriter(moduleOp.getContext()); + for (Operation *op : targets) { + if (op->getParentOp() == nullptr) + continue; + + (void)tryDecoupleControlFlowOp(op, rewriter); + } + + if (failed(verify(moduleOp))) signalPassFailure(); } std::unique_ptr> createTritonControlFlowOptPass() { - // Keep construction in this translation unit; registration is generated - // from Passes.td and exposes only this public factory. return std::make_unique(); } diff --git a/third_party/ascend/lib/TritonToLinalg/BlockPtrAnalysis.cpp b/third_party/ascend/lib/TritonToLinalg/BlockPtrAnalysis.cpp index 3d21e90027..56b679d6a7 100644 --- a/third_party/ascend/lib/TritonToLinalg/BlockPtrAnalysis.cpp +++ b/third_party/ascend/lib/TritonToLinalg/BlockPtrAnalysis.cpp @@ -148,8 +148,7 @@ OpFoldResult BlockData::inferBlockOffset(const Location &loc, OpBuilder &builder) const { OpFoldResult retOffset = builder.getIndexAttr(0); for (auto ofr : offsets) { - retOffset = - addOpFoldResult(retOffset, ofr, loc, builder, builder.getIndexType()); + retOffset = addOpFoldResult(retOffset, ofr, loc, builder); } return retOffset; } @@ -193,9 +192,8 @@ void BlockData::addBlock(BlockData &lBlock, BlockData &rBlock, Location loc, // 2. otherwise, both lhs and rhs are scalar type with rank 0 // Except above, original `scalar` has been fused into `offset` under add. if (lBlock.isScalar() && rBlock.isScalar()) { - auto addScalar = - addOpFoldResult(lBlock.getScalarRef(), rBlock.getScalarRef(), loc, - rewriter, rewriter.getIndexType()); + auto addScalar = addOpFoldResult(lBlock.getScalarRef(), + rBlock.getScalarRef(), loc, rewriter); this->scalar = addScalar; } else if (lBlock.getRank() == 0) { // When both lhs and rhs are scalar type with rank 0, just try passing @@ -206,14 +204,12 @@ void BlockData::addBlock(BlockData &lBlock, BlockData &rBlock, Location loc, for (const auto &[lOffset, rOffset] : llvm::zip(lBlock.getOffsetsRef(), rBlock.getOffsetsRef())) { - this->offsets.push_back(addOpFoldResult(lOffset, rOffset, loc, rewriter, - rewriter.getIndexType())); + this->offsets.push_back(addOpFoldResult(lOffset, rOffset, loc, rewriter)); } for (const auto &[lStride, rStride] : llvm::zip(lBlock.getStridesRef(), rBlock.getStridesRef())) { - this->strides.push_back(addOpFoldResult(lStride, rStride, loc, rewriter, - rewriter.getIndexType())); + this->strides.push_back(addOpFoldResult(lStride, rStride, loc, rewriter)); } // Both sizes are same implicitly under `add` @@ -271,8 +267,8 @@ void BlockData::mulBlock(BlockData &lBlock, BlockData &rBlock, Location loc, << " rBlbock.scalar:" << rBlock.getScalar() << "\n"; }); - auto scalar = mulOpFoldResult(lBlock.getScalar(), rBlock.getScalar(), loc, - rewriter, rewriter.getIndexType()); + auto scalar = + mulOpFoldResult(lBlock.getScalar(), rBlock.getScalar(), loc, rewriter); this->scalar = scalar; } @@ -290,13 +286,11 @@ void BlockData::mulBlock(BlockData &lBlock, BlockData &rBlock, Location loc, // In mulBlock, `scalar` will be accumulated into `offset` and `stride` OpFoldResult rScalar = rb->getScalarRef(); for (const auto &lOffset : lb->getOffsetsRef()) { - this->offsets.push_back(mulOpFoldResult(lOffset, rScalar, loc, rewriter, - rewriter.getIndexType())); + this->offsets.push_back(mulOpFoldResult(lOffset, rScalar, loc, rewriter)); } for (const auto &lStride : lb->getStridesRef()) { - this->strides.push_back(mulOpFoldResult(lStride, rScalar, loc, rewriter, - rewriter.getIndexType())); + this->strides.push_back(mulOpFoldResult(lStride, rScalar, loc, rewriter)); } this->sizes = lb->getSizesRef(); @@ -1398,7 +1392,7 @@ accumulatePotentialOffsetOnBase(triton::MakeTensorPtrOp op, Value base, "base of MakeTensorPtrOp only comes from native ptr or AddPtrOp"); return addOpFoldResult(offset, baseRecast.getConstifiedMixedOffset(), - op.getLoc(), rewriter, rewriter.getIndexType()); + op.getLoc(), rewriter); } return offset; @@ -1568,8 +1562,7 @@ void BlockDataParser::rewriteMakeTensorPtrOp( SmallVector newOffsets; for (auto [offset, stride] : llvm::zip(data.getOffsetsRef(), data.getStridesRef())) - newOffsets.push_back(mulOpFoldResult(offset, stride, loc, rewriter, - rewriter.getIndexType())); + newOffsets.push_back(mulOpFoldResult(offset, stride, loc, rewriter)); // 1. Consider that current base ptr may comes from `triton::AddPtrOp`, // which have been converted to `memref::ReinterpretCastOp` with 1D @@ -1737,9 +1730,8 @@ void BlockDataParser::rewriteAdvanceOp( llvm::zip(incrementOffsets, blockData.getOffsetsRef(), blockData.getStridesRef())) { auto curDimOffset = - addOpFoldResult(mulOpFoldResult(increment, stride, loc, rewriter, - rewriter.getIndexType()), - originalOffset, loc, rewriter, rewriter.getIndexType()); + addOpFoldResult(mulOpFoldResult(increment, stride, loc, rewriter), + originalOffset, loc, rewriter); newOffsets.push_back(curDimOffset); } @@ -2544,8 +2536,7 @@ void BlockDataParser::rewriteAddPtrToUnstrucMemAcc( bLoc, bB.getIndexType(), scalarOffsetRaw); OpFoldResult baseOffset = bB.getIndexAttr(0); for (auto ofr : data.getOffsetsRef()) { - baseOffset = - addOpFoldResult(baseOffset, ofr, bLoc, bB, bB.getIndexType()); + baseOffset = addOpFoldResult(baseOffset, ofr, bLoc, bB); } Value baseVal = getValueOrCreateConstantIndexOp(bB, bLoc, baseOffset); Value combinedOffset = diff --git a/third_party/ascend/lib/TritonToLinalg/MaskAnalysis.cpp b/third_party/ascend/lib/TritonToLinalg/MaskAnalysis.cpp index a95f39ef63..d794b0b5d4 100644 --- a/third_party/ascend/lib/TritonToLinalg/MaskAnalysis.cpp +++ b/third_party/ascend/lib/TritonToLinalg/MaskAnalysis.cpp @@ -257,10 +257,8 @@ LogicalResult MaskState::addStateScalar(const MaskState &state, const OpFoldResult scalar, const Location &loc, OpBuilder &builder) { - start = addOpFoldResult(state.start, scalar, loc, builder, - builder.getIndexType()); - end = - addOpFoldResult(state.end, scalar, loc, builder, builder.getIndexType()); + start = addOpFoldResult(state.start, scalar, loc, builder); + end = addOpFoldResult(state.end, scalar, loc, builder); dims = state.dims; offsets = state.offsets; @@ -343,10 +341,8 @@ LogicalResult MaskState::minStates(const MaskState &lhsState, auto newOffset = maxOpFoldResult(lhsOffset, rhsOffset, loc, builder); auto lhsDim = lhsState.dims[i]; auto rhsDim = rhsState.dims[i]; - auto lhsEnd = addOpFoldResult(lhsOffset, lhsDim, loc, builder, - builder.getIndexType()); - auto rhsEnd = addOpFoldResult(rhsOffset, rhsDim, loc, builder, - builder.getIndexType()); + auto lhsEnd = addOpFoldResult(lhsOffset, lhsDim, loc, builder); + auto rhsEnd = addOpFoldResult(rhsOffset, rhsDim, loc, builder); auto newEnd = minOpFoldResult(lhsEnd, rhsEnd, loc, builder); auto newDim = subOpFoldResult(newEnd, newOffset, loc, builder); auto clampedNewDim = @@ -567,8 +563,8 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location &loc, } case arith::CmpIPredicate::sle: { // lhs <= rhs <=> lhs < rhs + 1 - auto rhsPlusOne = addOpFoldResult(rhsState.scalar, builder.getIndexAttr(1), - loc, builder, builder.getIndexType()); + auto rhsPlusOne = + addOpFoldResult(rhsState.scalar, builder.getIndexAttr(1), loc, builder); auto realBound = maxOpFoldResult(lhsState.start, rhsPlusOne, loc, builder); auto newEnd = minOpFoldResult(lhsState.end, realBound, loc, builder); auto newDim = subOpFoldResult(newEnd, lhsState.start, loc, builder); @@ -701,8 +697,8 @@ LogicalResult MaskState::parseSplat(triton::SplatOp splatOp, if (src.getType().isInteger(1) && !splatOp->use_empty() && llvm::all_of(splatOp->getUsers(), splatAsMask)) { for (auto s : dstShape) { - auto currentDim = mulOpFoldResult(builder.getIndexAttr(s), this->scalar, - loc, builder, builder.getIndexType()); + auto currentDim = + mulOpFoldResult(builder.getIndexAttr(s), this->scalar, loc, builder); this->dims.push_back(currentDim); this->offsets.push_back(builder.getIndexAttr(0)); } diff --git a/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp b/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp index 6b3d23ce4a..d3dadca0a3 100644 --- a/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp +++ b/third_party/ascend/lib/TritonToLinalg/TritonOpConverter.cpp @@ -294,8 +294,7 @@ SelectCanonicalizer::matchAndRewrite(arith::SelectOp op, } invertFalseDims.push_back(offVal); - trueDimOp = addOpFoldResult(offVal, dimVal, loc, rewriter, - rewriter.getIndexType()); + trueDimOp = addOpFoldResult(offVal, dimVal, loc, rewriter); invertTrueDims.push_back(trueDimOp); } } diff --git a/third_party/ascend/lib/TritonToStructured/MaskAnalysis.cpp b/third_party/ascend/lib/TritonToStructured/MaskAnalysis.cpp index 3760291b60..11be1a5b15 100644 --- a/third_party/ascend/lib/TritonToStructured/MaskAnalysis.cpp +++ b/third_party/ascend/lib/TritonToStructured/MaskAnalysis.cpp @@ -403,8 +403,7 @@ LogicalResult MaskState::addStateScalar(const MaskState &state, const OpFoldResult scalar, Location loc, OpBuilder &builder) { for (auto info : state.stateInfo) { - info.offset = addOpFoldResult(info.offset, scalar, loc, builder, - builder.getIndexType()); + info.offset = addOpFoldResult(info.offset, scalar, loc, builder); this->stateInfo.emplace_back(info); } return success(); diff --git a/third_party/ascend/lib/TritonToStructured/PtrAnalysis.cpp b/third_party/ascend/lib/TritonToStructured/PtrAnalysis.cpp index 6ec01f786d..3c26dc7497 100644 --- a/third_party/ascend/lib/TritonToStructured/PtrAnalysis.cpp +++ b/third_party/ascend/lib/TritonToStructured/PtrAnalysis.cpp @@ -208,8 +208,7 @@ void PtrState::normalizeState(const Location loc, OpBuilder &builder) { for (++it; it != this->stateInfo.end() && isZero(it->stride) && it->dimIndex == dimIndex; ++it) { - newShape = mulOpFoldResult(newShape, it->shape, loc, builder, - builder.getIndexType()); + newShape = mulOpFoldResult(newShape, it->shape, loc, builder); } newStateInfo.emplace_back(zeroAttr, newShape, dimIndex); } @@ -443,13 +442,13 @@ LogicalResult PtrState::mulState(const PtrState &lhsState, SmallVector newStateInfo; for (auto info : lhs->stateInfo) { - OpFoldResult newStride = mulOpFoldResult(info.stride, rhs->offset, loc, - builder, builder.getIndexType()); + OpFoldResult newStride = + mulOpFoldResult(info.stride, rhs->offset, loc, builder); newStateInfo.emplace_back(newStride, info.shape, info.dimIndex); } - auto newOffset = mulOpFoldResult(lhsState.offset, rhsState.offset, loc, - builder, builder.getIndexType()); + auto newOffset = + mulOpFoldResult(lhsState.offset, rhsState.offset, loc, builder); updatePtrState(newStateInfo, lhs->sizes, lhs->source, newOffset, loc, builder, lhs->shouldLinearize); @@ -561,8 +560,7 @@ LogicalResult PtrState::addState(PtrState &lhsState, PtrState &rhsState, return failure(); } - auto newStride = addOpFoldResult(lIt->stride, rIt->stride, loc, builder, - builder.getIndexType()); + auto newStride = addOpFoldResult(lIt->stride, rIt->stride, loc, builder); newStateInfo.emplace_back(newStride, newShape, lIt->dimIndex); if (isEqual(lIt->shape, newShape)) @@ -583,8 +581,8 @@ LogicalResult PtrState::addState(PtrState &lhsState, PtrState &rhsState, } auto newSource = source = lhsState.source ? lhsState.source : rhsState.source; - auto newOffset = addOpFoldResult(lhsState.offset, rhsState.offset, loc, - builder, builder.getIndexType()); + auto newOffset = + addOpFoldResult(lhsState.offset, rhsState.offset, loc, builder); auto newShouldLinearize = lhsState.shouldLinearize || rhsState.shouldLinearize; auto newSizes = lhsState.sizes; diff --git a/third_party/ascend/lib/Utils/Utils.cpp b/third_party/ascend/lib/Utils/Utils.cpp index 0e5bb1c283..dc967f81e3 100644 --- a/third_party/ascend/lib/Utils/Utils.cpp +++ b/third_party/ascend/lib/Utils/Utils.cpp @@ -47,7 +47,6 @@ #include "triton/Dialect/Triton/IR/Dialect.h" #include "triton/Dialect/Triton/IR/Types.h" -#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/SmallVectorExtras.h" @@ -85,87 +84,6 @@ static std::optional getConstantOfAttr(const OpFoldResult &arg) { return std::nullopt; } -static FailureOr materializeIndexOperand(const OpFoldResult &operand, - std::optional constant, - Location loc, - OpBuilder &builder) { - if (constant) - return createConstIndexValueOp(loc, builder, *constant); - auto value = dyn_cast(operand); - if (!value) - return failure(); - return castIntegerLike(builder, loc, value, builder.getIndexType()); -} - -static Type getOpFoldResultType(const OpFoldResult &operand) { - if (auto value = dyn_cast(operand)) - return value.getType(); - if (auto attribute = dyn_cast(operand)) { - if (auto typedAttribute = dyn_cast(attribute)) - return typedAttribute.getType(); - } - return Type(); -} - -static bool isSupportedScalarIntegerType(Type type) { - return type && (type.isIndex() || type.isSignlessInteger()); -} - -static FailureOr resolveIntegerResultType(const OpFoldResult &lhs, - const OpFoldResult &rhs, - Type requestedType) { - Type lhsType = getOpFoldResultType(lhs); - Type rhsType = getOpFoldResultType(rhs); - if (!isSupportedScalarIntegerType(lhsType) || - !isSupportedScalarIntegerType(rhsType) || - (requestedType && !isSupportedScalarIntegerType(requestedType))) - return failure(); - - if (requestedType) - return requestedType; - if (lhsType == rhsType) - return lhsType; - - // Index has no fixed IR-level width, so it does not participate in default - // widening against ordinary integers. - auto lhsInteger = dyn_cast(lhsType); - auto rhsInteger = dyn_cast(rhsType); - if (!lhsInteger || !rhsInteger || - lhsInteger.getWidth() == rhsInteger.getWidth()) - return failure(); - return lhsInteger.getWidth() > rhsInteger.getWidth() ? lhsType : rhsType; -} - -static FailureOr -castIntegerFoldResult(const OpFoldResult &operand, Type targetType, - Location loc, OpBuilder &builder) { - if (std::optional constant = getConstantOfAttr(operand)) - return OpFoldResult(builder.getIntegerAttr(targetType, *constant)); - - auto value = dyn_cast(operand); - if (!value) - return failure(); - FailureOr converted = castIntegerLike(builder, loc, value, targetType); - if (failed(converted)) - return failure(); - return OpFoldResult(*converted); -} - -static FailureOr -materializeIntegerOperand(const OpFoldResult &operand, - std::optional constant, Type targetType, - Location loc, OpBuilder &builder) { - if (constant) - return builder - .create( - loc, builder.getIntegerAttr(targetType, *constant)) - .getResult(); - auto value = dyn_cast(operand); - if (!value) - return failure(); - return castIntegerLike(builder, loc, value, targetType); -} - namespace ConverterUtils { std::optional @@ -1078,96 +996,37 @@ bool isTensorPtrType(Type type) { } // namespace triton -bool haveSameTypes(TypeRange lhs, TypeRange rhs) { - return llvm::equal(lhs, rhs); -} - -FailureOr castIntegerLike(OpBuilder &builder, Location loc, Value value, - Type targetType, - IntegerExtensionKind extension) { - if (!value || !targetType) - return failure(); - - Type sourceType = value.getType(); - if (sourceType == targetType) - return value; - - if ((sourceType.isIndex() && isa(targetType)) || - (isa(sourceType) && targetType.isIndex())) { - if (extension == IntegerExtensionKind::Unsigned) - return builder.create(loc, targetType, value) - .getResult(); - return builder.create(loc, targetType, value) - .getResult(); - } - - auto extendOrTruncate = [&](unsigned sourceWidth, - unsigned targetWidth) -> FailureOr { - if (sourceWidth < targetWidth) { - if (extension == IntegerExtensionKind::Unsigned) - return builder.create(loc, targetType, value) - .getResult(); - return builder.create(loc, targetType, value).getResult(); - } - if (sourceWidth > targetWidth) - return builder.create(loc, targetType, value) - .getResult(); - return failure(); - }; - - auto sourceInteger = dyn_cast(sourceType); - auto targetInteger = dyn_cast(targetType); - if (sourceInteger && targetInteger) - return extendOrTruncate(sourceInteger.getWidth(), targetInteger.getWidth()); - - auto sourceTensor = dyn_cast(sourceType); - auto targetTensor = dyn_cast(targetType); - if (!sourceTensor || !targetTensor || - sourceTensor.getShape() != targetTensor.getShape() || - sourceTensor.getEncoding() != targetTensor.getEncoding()) - return failure(); - - auto sourceElement = dyn_cast(sourceTensor.getElementType()); - auto targetElement = dyn_cast(targetTensor.getElementType()); - if (!sourceElement || !targetElement) - return failure(); - - return extendOrTruncate(sourceElement.getWidth(), targetElement.getWidth()); -} - // TODO: imply these function below OpFoldResult addOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b, - Type resultType) { - FailureOr resolvedType = resolveIntegerResultType(lhs, rhs, resultType); - if (failed(resolvedType)) - return OpFoldResult(); - + const Location &loc, OpBuilder &b) { auto lhsInt = getConstantOfAttr(lhs); auto rhsInt = getConstantOfAttr(rhs); if (lhsInt && rhsInt) - return b.getIntegerAttr(*resolvedType, *lhsInt + *rhsInt); + return b.getIndexAttr(lhsInt.value() + rhsInt.value()); - if (!lhsInt && rhsInt && *rhsInt == 0) { - FailureOr result = - castIntegerFoldResult(lhs, *resolvedType, loc, b); - return succeeded(result) ? *result : OpFoldResult(); - } - if (!rhsInt && lhsInt && *lhsInt == 0) { - FailureOr result = - castIntegerFoldResult(rhs, *resolvedType, loc, b); - return succeeded(result) ? *result : OpFoldResult(); + if (!lhsInt && rhsInt && rhsInt.value() == 0) + return lhs; + if (!rhsInt && lhsInt && lhsInt.value() == 0) + return rhs; + + auto lhsValue = dyn_cast(lhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); } - FailureOr lhsValue = - materializeIntegerOperand(lhs, lhsInt, *resolvedType, loc, b); - FailureOr rhsValue = - materializeIntegerOperand(rhs, rhsInt, *resolvedType, loc, b); - if (failed(lhsValue) || failed(rhsValue)) - return OpFoldResult(); + auto rhsValue = dyn_cast(rhs); + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } - return b.create(loc, *lhsValue, *rhsValue).getResult(); + return b.create(loc, lhsValue, rhsValue).getResult(); } OpFoldResult subOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, @@ -1181,54 +1040,61 @@ OpFoldResult subOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, if (!lhsInt && rhsInt && rhsInt.value() == 0) return lhs; - FailureOr lhsValue = materializeIndexOperand(lhs, lhsInt, loc, b); - FailureOr rhsValue = materializeIndexOperand(rhs, rhsInt, loc, b); - if (failed(lhsValue) || failed(rhsValue)) - return OpFoldResult(); + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } - return b.create(loc, *lhsValue, *rhsValue).getResult(); + return b.create(loc, lhsValue, rhsValue).getResult(); } OpFoldResult mulOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b, - Type resultType) { - FailureOr resolvedType = resolveIntegerResultType(lhs, rhs, resultType); - if (failed(resolvedType)) - return OpFoldResult(); - + const Location &loc, OpBuilder &b) { auto lhsInt = getConstantOfAttr(lhs); auto rhsInt = getConstantOfAttr(rhs); if (lhsInt && rhsInt) - return b.getIntegerAttr(*resolvedType, *lhsInt * *rhsInt); + return b.getIndexAttr(lhsInt.value() * rhsInt.value()); if (lhsInt) { - if (*lhsInt == 0) - return b.getIntegerAttr(*resolvedType, 0); - if (*lhsInt == 1) { - FailureOr result = - castIntegerFoldResult(rhs, *resolvedType, loc, b); - return succeeded(result) ? *result : OpFoldResult(); - } + if (lhsInt.value() == 0) + return lhs; + if (lhsInt.value() == 1) + return rhs; } if (rhsInt) { - if (*rhsInt == 0) - return b.getIntegerAttr(*resolvedType, 0); - if (*rhsInt == 1) { - FailureOr result = - castIntegerFoldResult(lhs, *resolvedType, loc, b); - return succeeded(result) ? *result : OpFoldResult(); - } + if (rhsInt.value() == 0) + return rhs; + if (rhsInt.value() == 1) + return lhs; } - FailureOr lhsValue = - materializeIntegerOperand(lhs, lhsInt, *resolvedType, loc, b); - FailureOr rhsValue = - materializeIntegerOperand(rhs, rhsInt, *resolvedType, loc, b); - if (failed(lhsValue) || failed(rhsValue)) - return OpFoldResult(); + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } - return b.create(loc, *lhsValue, *rhsValue).getResult(); + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); } OpFoldResult divOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, @@ -1254,12 +1120,22 @@ OpFoldResult divOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, return lhs; } - FailureOr lhsValue = materializeIndexOperand(lhs, lhsInt, loc, b); - FailureOr rhsValue = materializeIndexOperand(rhs, rhsInt, loc, b); - if (failed(lhsValue) || failed(rhsValue)) - return OpFoldResult(); + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } - return b.create(loc, *lhsValue, *rhsValue).getResult(); + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); } OpFoldResult remOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, @@ -1280,12 +1156,22 @@ OpFoldResult remOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, return lhs; } - FailureOr lhsValue = materializeIndexOperand(lhs, lhsInt, loc, b); - FailureOr rhsValue = materializeIndexOperand(rhs, rhsInt, loc, b); - if (failed(lhsValue) || failed(rhsValue)) - return OpFoldResult(); + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } - return b.create(loc, *lhsValue, *rhsValue).getResult(); + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); } OpFoldResult minOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, @@ -1295,12 +1181,22 @@ OpFoldResult minOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, if (lhsInt && rhsInt) return b.getIndexAttr(std::min(lhsInt.value(), rhsInt.value())); - FailureOr lhsValue = materializeIndexOperand(lhs, lhsInt, loc, b); - FailureOr rhsValue = materializeIndexOperand(rhs, rhsInt, loc, b); - if (failed(lhsValue) || failed(rhsValue)) - return OpFoldResult(); + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } - return b.create(loc, *lhsValue, *rhsValue).getResult(); + return b.create(loc, lhsValue, rhsValue).getResult(); } OpFoldResult maxOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, @@ -1310,12 +1206,22 @@ OpFoldResult maxOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, if (lhsInt && rhsInt) return b.getIndexAttr(std::max(lhsInt.value(), rhsInt.value())); - FailureOr lhsValue = materializeIndexOperand(lhs, lhsInt, loc, b); - FailureOr rhsValue = materializeIndexOperand(rhs, rhsInt, loc, b); - if (failed(lhsValue) || failed(rhsValue)) - return OpFoldResult(); + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + rhsValue = convertToIndexIfNeeded(rhsValue, loc, b); + assert(isa(rhsValue.getType())); + } - return b.create(loc, *lhsValue, *rhsValue).getResult(); + return b.create(loc, lhsValue, rhsValue).getResult(); } void addReduceWithIndexAttr(ReduceWithIndexParams params, @@ -1473,18 +1379,19 @@ OpFoldResult getOpFoldResultOfLayoutInfo(Value value, OpBuilder &builder) { return constantFold; } - Type sourceType = value.getType(); - if (!sourceType.isIndex() && !isa(sourceType)) + if (!isa(value.getType())) llvm_unreachable("Illegal data type when parse block data layout info"); - IntegerExtensionKind extension = sourceType.isInteger(/*width=*/1) - ? IntegerExtensionKind::Unsigned - : IntegerExtensionKind::Signed; - FailureOr converted = castIntegerLike( - builder, value.getLoc(), value, builder.getIndexType(), extension); - if (failed(converted)) - llvm_unreachable("Failed to convert block data layout info to index"); - return *converted; + if (!isa(value.getType())) { + if (value.getType().isInteger(/*width*/ 1)) + value = builder.create( + value.getLoc(), builder.getIndexType(), value); + else + value = builder.create(value.getLoc(), + builder.getIndexType(), value); + } + + return value; } // Specialize the Typeless Value (Zero, Min, Max) into a mlir TypedAttr @@ -1687,6 +1594,16 @@ bool isOne(const OpFoldResult ofr) { return staticOfr.has_value() && staticOfr.value() == 1; } +Value convertToIndexIfNeeded(Value input, const Location &loc, OpBuilder &b) { + auto inputType = input.getType(); + if (auto intType = dyn_cast(inputType)) { + if (intType.isInteger(32) || intType.isInteger(64)) { + return b.create(loc, b.getIndexType(), input); + } + } + return input; +} + RankedTensorType getExtractSlicedType(ArrayRef shape, const llvm::SmallBitVector &droppedDims, Type elemType) { diff --git a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_different_base_invalid.mlir b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_different_base_invalid.mlir deleted file mode 100644 index ef8cc4d24b..0000000000 --- a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/block_ptr_different_base_invalid.mlir +++ /dev/null @@ -1,23 +0,0 @@ -// RUN: not triton-opt --triton-control-flow-opt %s 2>&1 | FileCheck %s - -module { - tt.func public @if_block_ptr_different_base(%base0: !tt.ptr, %base1: !tt.ptr, %cond: i1) -> !tt.ptr> { - %c0_i32 = arith.constant 0 : i32 - %c1_i32 = arith.constant 1 : i32 - %c2_i32 = arith.constant 2 : i32 - %c1_i64 = arith.constant 1 : i64 - %c16_i64 = arith.constant 16 : i64 - %ptr0 = tt.make_tensor_ptr %base0, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> - %ptr1 = tt.make_tensor_ptr %base1, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> - %selected = scf.if %cond -> (!tt.ptr>) { - %then_ptr = tt.advance %ptr0, [%c1_i32] : !tt.ptr> - scf.yield %then_ptr : !tt.ptr> - } else { - %else_ptr = tt.advance %ptr1, [%c2_i32] : !tt.ptr> - scf.yield %else_ptr : !tt.ptr> - } - tt.return %selected : !tt.ptr> - } -} - -// CHECK: error: failed to analyze pointer components across control flow diff --git a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/cf_terminal_return.mlir b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/cf_terminal_return.mlir index 0d65f681ec..b4e64b7c46 100644 --- a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/cf_terminal_return.mlir +++ b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/cf_terminal_return.mlir @@ -79,22 +79,6 @@ module { // ----- -module { - tt.func private @entry_return_with_unreachable_block() -> tensor<32xf32> { - %cst = arith.constant dense<0.000000e+00> : tensor<32xf32> - tt.return %cst : tensor<32xf32> - ^bb1: - %poison = ub.poison : tensor<32xf32> - tt.return %poison : tensor<32xf32> - } -} - -// CHECK-LABEL: tt.func private @entry_return_with_unreachable_block -// CHECK: tt.return -// CHECK-NOT: ^bb1 - -// ----- - module { tt.func public @terminal_return_branch_args(%cond: i1, %lhs: i32, %rhs: i32, %bias: i32) -> i32 { cf.cond_br %cond, ^bb1(%lhs : i32), ^bb2(%rhs : i32) diff --git a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/scf_pointer_decouple.mlir b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/scf_pointer_decouple.mlir index 22446047fa..e9c3d413ce 100644 --- a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/scf_pointer_decouple.mlir +++ b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/scf_pointer_decouple.mlir @@ -1,8 +1,4 @@ // RUN: triton-opt --triton-control-flow-opt --split-input-file %s | FileCheck %s -// RUN: triton-opt --triton-control-flow-opt --split-input-file %s | FileCheck %s --check-prefix=NO-ADVANCE - -// NO-ADVANCE: module -// NO-ADVANCE-NOT: tt.advance module { tt.func public @for_block_ptr_dynamic_step(%base: !tt.ptr, %ub: index) -> !tt.ptr> { @@ -33,7 +29,7 @@ module { // ----- module { - tt.func public @for_block_ptr_invariant_delta_carried(%base: !tt.ptr, %ub: index, %stride: i32) -> tensor<32xf16> { + tt.func public @for_block_ptr_invariant_step_closed_form(%base: !tt.ptr, %ub: index, %stride: i32) -> tensor<32xf16> { %c0_i32 = arith.constant 0 : i32 %c1_i64 = arith.constant 1 : i64 %c32_i64 = arith.constant 32 : i64 @@ -51,12 +47,13 @@ module { } } -// CHECK-LABEL: tt.func public @for_block_ptr_invariant_delta_carried -// CHECK: %[[FOR:[^:]+]]:2 = scf.for +// CHECK-LABEL: tt.func public @for_block_ptr_invariant_step_closed_form +// CHECK: %[[FOR:[^:]+]]:2 = scf.for %[[IV:[^ ]+]] = // CHECK-SAME: {{.*}} iter_args(%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) -> (i32, tensor<32xf16>) { -// CHECK-NOT: arith.index_cast -// CHECK-NOT: arith.muli -// CHECK: %[[NEXT:.*]] = arith.addi %{{.*}}, %{{.*}} : i32 +// CHECK: %[[IV_I32:.*]] = arith.index_cast %[[IV]] : index to i32 +// CHECK: %[[SCALED:.*]] = arith.muli %[[IV_I32]], %{{.*}} : i32 +// CHECK: %[[CURRENT:.*]] = arith.addi %{{.*}}, %[[SCALED]] : i32 +// CHECK: %[[NEXT:.*]] = arith.addi %[[CURRENT]], %{{.*}} : i32 // CHECK: %[[LOAD_PTR:.*]] = tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[NEXT]]] {order = array} : > // CHECK: %[[LOADED:.*]] = tt.load %[[LOAD_PTR]] : !tt.ptr> // CHECK: scf.yield %[[NEXT]], %[[LOADED]] : i32, tensor<32xf16> @@ -65,116 +62,6 @@ module { // ----- -module { - tt.func public @for_block_ptr_dynamic_bounds_and_step(%base: !tt.ptr, %lb: index, %ub: index, %step: index, %delta: i32) -> !tt.ptr> { - %c0_i32 = arith.constant 0 : i32 - %c1_i64 = arith.constant 1 : i64 - %c32_i64 = arith.constant 32 : i64 - %ptr0 = tt.make_tensor_ptr %base, [%c32_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> - %final = scf.for %iv = %lb to %ub step %step iter_args(%ptr = %ptr0) -> (!tt.ptr>) { - %next = tt.advance %ptr, [%delta] : !tt.ptr> - scf.yield %next : !tt.ptr> - } - tt.return %final : !tt.ptr> - } -} - -// CHECK-LABEL: tt.func public @for_block_ptr_dynamic_bounds_and_step -// CHECK: %[[FOR:.*]] = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}) -> (i32) { -// CHECK-NOT: arith.muli -// CHECK: %[[NEXT:.*]] = arith.addi %[[OFF]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT]] : i32 -// CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]] {order = array} : > - -// ----- - -module { - tt.func public @for_block_ptr_iter_arg_delta(%base: !tt.ptr, %ub: index) -> !tt.ptr> { - %c0_i32 = arith.constant 0 : i32 - %c1_i32 = arith.constant 1 : i32 - %c1_i64 = arith.constant 1 : i64 - %c32_i64 = arith.constant 32 : i64 - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %ptr0 = tt.make_tensor_ptr %base, [%c32_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> - %res:2 = scf.for %iv = %c0 to %ub step %c1 iter_args(%ptr = %ptr0, %delta = %c1_i32) -> (!tt.ptr>, i32) { - %next_ptr = tt.advance %ptr, [%delta] : !tt.ptr> - %next_delta = arith.addi %delta, %c1_i32 : i32 - scf.yield %next_ptr, %next_delta : !tt.ptr>, i32 - } - tt.return %res#0 : !tt.ptr> - } -} - -// CHECK-LABEL: tt.func public @for_block_ptr_iter_arg_delta -// CHECK: %[[FOR:.*]]:2 = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}, %[[DELTA:.*]] = %{{.*}}) -> (i32, i32) { -// CHECK: %[[NEXT_PTR:.*]] = arith.addi %[[OFF]], %[[DELTA]] : i32 -// CHECK: %[[NEXT_DELTA:.*]] = arith.addi %[[DELTA]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT_PTR]], %[[NEXT_DELTA]] : i32, i32 -// CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]#0] {order = array} : > - -// ----- - -module { - tt.func public @for_block_ptr_multidim_delta(%base: !tt.ptr) -> !tt.ptr> { - %c0_i32 = arith.constant 0 : i32 - %c1_i32 = arith.constant 1 : i32 - %c2_i32 = arith.constant 2 : i32 - %c3_i32 = arith.constant 3 : i32 - %c1_i64 = arith.constant 1 : i64 - %c16_i64 = arith.constant 16 : i64 - %c32_i64 = arith.constant 32 : i64 - %c0 = arith.constant 0 : index - %c4 = arith.constant 4 : index - %c1 = arith.constant 1 : index - %ptr0 = tt.make_tensor_ptr %base, [%c16_i64, %c32_i64], [%c32_i64, %c1_i64], [%c2_i32, %c3_i32] {order = array} : !tt.ptr> - %final = scf.for %iv = %c0 to %c4 step %c1 iter_args(%ptr = %ptr0) -> (!tt.ptr>) { - %next = tt.advance %ptr, [%c1_i32, %c2_i32] : !tt.ptr> - scf.yield %next : !tt.ptr> - } - tt.return %final : !tt.ptr> - } -} - -// CHECK-LABEL: tt.func public @for_block_ptr_multidim_delta -// CHECK: %[[FOR:.*]]:2 = scf.for {{.*}} iter_args(%[[OFF0:.*]] = %{{.*}}, %[[OFF1:.*]] = %{{.*}}) -> (i32, i32) { -// CHECK: %[[NEXT0:.*]] = arith.addi %[[OFF0]], %{{.*}} : i32 -// CHECK: %[[NEXT1:.*]] = arith.addi %[[OFF1]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT0]], %[[NEXT1]] : i32, i32 -// CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}, %{{.*}}], [%{{.*}}, %{{.*}}], [%[[FOR]]#0, %[[FOR]]#1] {order = array} : > - -// ----- - -module { - tt.func public @for_block_ptr_zero_trip(%base: !tt.ptr) -> !tt.ptr> { - %c0_i32 = arith.constant 0 : i32 - %c1_i32 = arith.constant 1 : i32 - %c1_i64 = arith.constant 1 : i64 - %c32_i64 = arith.constant 32 : i64 - %c0 = arith.constant 0 : index - %c4 = arith.constant 4 : index - %c1 = arith.constant 1 : index - %ptr0 = tt.make_tensor_ptr %base, [%c32_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> - %final = scf.for %iv = %c4 to %c0 step %c1 iter_args(%ptr = %ptr0) -> (!tt.ptr>) { - %next = tt.advance %ptr, [%c1_i32] : !tt.ptr> - scf.yield %next : !tt.ptr> - } - tt.return %final : !tt.ptr> - } -} - -// CHECK-LABEL: tt.func public @for_block_ptr_zero_trip -// CHECK: %[[FOR:.*]] = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}) -> (i32) { -// CHECK: %[[NEXT:.*]] = arith.addi %[[OFF]], %{{.*}} : i32 -// CHECK: scf.yield %[[NEXT]] : i32 -// CHECK: } -// CHECK: tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]] {order = array} : > - -// ----- - module { tt.func public @while_block_ptr_basic(%base: !tt.ptr, %n: i32) -> !tt.ptr> { %c0_i32 = arith.constant 0 : i32 @@ -474,6 +361,7 @@ module { // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : i32 // CHECK-DAG: %[[POST_STEP:.*]] = tt.splat %[[C2]] : i32 -> tensor<4xi32> // CHECK: %[[FOR:.*]] = scf.for {{.*}} iter_args(%[[OFF:.*]] = %{{.*}}) -> (tensor<4xi32>) { +// CHECK: %[[CARRIED:.*]] = arith.addi %{{.*}}, %[[OFF]] : tensor<4xi32> // CHECK: %[[SELECTED:.*]] = scf.if %{{.*}} -> (tensor<4xi32>) { // CHECK: arith.addi %{{.*}}, %{{.*}} : tensor<4xi32> // CHECK: scf.yield %{{.*}} : tensor<4xi32> @@ -481,12 +369,42 @@ module { // CHECK: arith.addi %{{.*}}, %{{.*}} : tensor<4xi32> // CHECK: scf.yield %{{.*}} : tensor<4xi32> // CHECK: } -// CHECK: %[[SELECTED_OFFSET:.*]] = arith.addi %{{.*}}, %[[SELECTED]] : tensor<4xi32> -// CHECK: %[[NEXT:.*]] = arith.addi %[[SELECTED_OFFSET]], %[[POST_STEP]] : tensor<4xi32> +// CHECK: %[[MERGED:.*]] = arith.addi %[[CARRIED]], %[[SELECTED]] : tensor<4xi32> +// CHECK: %[[NEXT:.*]] = arith.addi %{{.*}}, %[[POST_STEP]] : tensor<4xi32> // CHECK: scf.yield %[[NEXT]] : tensor<4xi32> // CHECK: } // CHECK: tt.addptr %{{.*}}, %[[FOR]] : tensor<4x!tt.ptr>, tensor<4xi32> +// ----- + +module { + tt.func public @if_tensor_ptr_diff_base_not_decoupled(%base0: !tt.ptr, %base1: !tt.ptr, %cond: i1) -> tensor<4x!tt.ptr> { + %c1_i32 = arith.constant 1 : i32 + %c2_i32 = arith.constant 2 : i32 + %off1 = tt.splat %c1_i32 : i32 -> tensor<4xi32> + %off2 = tt.splat %c2_i32 : i32 -> tensor<4xi32> + %splat0 = tt.splat %base0 : !tt.ptr -> tensor<4x!tt.ptr> + %splat1 = tt.splat %base1 : !tt.ptr -> tensor<4x!tt.ptr> + %selected = scf.if %cond -> (tensor<4x!tt.ptr>) { + %then_ptr = tt.addptr %splat0, %off1 : tensor<4x!tt.ptr>, tensor<4xi32> + scf.yield %then_ptr : tensor<4x!tt.ptr> + } else { + %else_ptr = tt.addptr %splat1, %off2 : tensor<4x!tt.ptr>, tensor<4xi32> + scf.yield %else_ptr : tensor<4x!tt.ptr> + } + tt.return %selected : tensor<4x!tt.ptr> + } +} + +// CHECK-LABEL: tt.func public @if_tensor_ptr_diff_base_not_decoupled +// CHECK: scf.if %{{.*}} -> (tensor<4x!tt.ptr>) { +// CHECK: tt.addptr +// CHECK: } else { +// CHECK: tt.addptr +// CHECK: } + +// ----- + module { tt.func public @for_if_block_ptr_load_after_post_advance(%base: !tt.ptr, %cond: i1) -> tensor<32xf16> { %c0_i32 = arith.constant 0 : i32 @@ -662,16 +580,16 @@ module { // CHECK: } else { // CHECK: scf.yield %{{.*}} : tensor<4xi32> // CHECK: } -// CHECK: tt.addptr %{{.*}}, %[[THEN_INNER]] : tensor<4x!tt.ptr>, tensor<4xi32> -// CHECK: scf.yield %[[THEN_INNER]] : tensor<4xi32> +// CHECK: arith.addi %{{.*}}, %[[THEN_INNER]] : tensor<4xi32> +// CHECK: scf.yield %{{.*}} : tensor<4xi32> // CHECK: } else { // CHECK: %[[ELSE_INNER:.*]] = scf.if %{{.*}} -> (tensor<4xi32>) { // CHECK: scf.yield %{{.*}} : tensor<4xi32> // CHECK: } else { // CHECK: scf.yield %{{.*}} : tensor<4xi32> // CHECK: } -// CHECK: tt.addptr %{{.*}}, %[[ELSE_INNER]] : tensor<4x!tt.ptr>, tensor<4xi32> -// CHECK: scf.yield %[[ELSE_INNER]] : tensor<4xi32> +// CHECK: arith.addi %{{.*}}, %[[ELSE_INNER]] : tensor<4xi32> +// CHECK: scf.yield %{{.*}} : tensor<4xi32> // CHECK: } // CHECK: arith.addi %{{.*}}, %{{.*}} : tensor<4xi32> // CHECK: tt.addptr %{{.*}}, %{{.*}} : tensor<4x!tt.ptr>, tensor<4xi32> @@ -705,6 +623,37 @@ module { // CHECK: } // CHECK: tt.make_tensor_ptr %{{.*}}, [%[[SHAPE]]], [%{{.*}}], [%{{.*}}] {order = array} : > +// ----- + +module { + tt.func public @if_block_ptr_diff_base_not_decoupled(%base0: !tt.ptr, %base1: !tt.ptr, %cond: i1) -> !tt.ptr> { + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + %c2_i32 = arith.constant 2 : i32 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %ptr0 = tt.make_tensor_ptr %base0, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + %ptr1 = tt.make_tensor_ptr %base1, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + %selected = scf.if %cond -> (!tt.ptr>) { + %then_ptr = tt.advance %ptr0, [%c1_i32] : !tt.ptr> + scf.yield %then_ptr : !tt.ptr> + } else { + %else_ptr = tt.advance %ptr1, [%c2_i32] : !tt.ptr> + scf.yield %else_ptr : !tt.ptr> + } + tt.return %selected : !tt.ptr> + } +} + +// CHECK-LABEL: tt.func public @if_block_ptr_diff_base_not_decoupled +// CHECK: scf.if %{{.*}} -> (!tt.ptr>) { +// CHECK: tt.advance +// CHECK: } else { +// CHECK: tt.advance +// CHECK: } + +// ----- + module { tt.func public @if_same_nested_result_not_decoupled(%base: !tt.ptr, %cond0: i1, %cond1: i1) -> !tt.ptr> { %c0_i32 = arith.constant 0 : i32 @@ -742,115 +691,3 @@ module { // CHECK: scf.yield %[[INNER_PTR]] : !tt.ptr> // CHECK: } // CHECK: tt.return %[[OUTER]] : !tt.ptr> - -// ----- - -module { - tt.func public @if_mixed_block_and_tensor_ptr(%block_base: !tt.ptr, %tensor_base: !tt.ptr, %cond: i1) -> (!tt.ptr>, tensor<4x!tt.ptr>) { - %c0_i32 = arith.constant 0 : i32 - %c1_i32 = arith.constant 1 : i32 - %c2_i32 = arith.constant 2 : i32 - %c1_i64 = arith.constant 1 : i64 - %c16_i64 = arith.constant 16 : i64 - %tensor_off1 = tt.splat %c1_i32 : i32 -> tensor<4xi32> - %tensor_off2 = tt.splat %c2_i32 : i32 -> tensor<4xi32> - %block_ptr = tt.make_tensor_ptr %block_base, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> - %tensor_ptr = tt.splat %tensor_base : !tt.ptr -> tensor<4x!tt.ptr> - %result:2 = scf.if %cond -> (!tt.ptr>, tensor<4x!tt.ptr>) { - %then_block = tt.advance %block_ptr, [%c1_i32] : !tt.ptr> - %then_tensor = tt.addptr %tensor_ptr, %tensor_off1 : tensor<4x!tt.ptr>, tensor<4xi32> - scf.yield %then_block, %then_tensor : !tt.ptr>, tensor<4x!tt.ptr> - } else { - %else_block = tt.advance %block_ptr, [%c2_i32] : !tt.ptr> - %else_tensor = tt.addptr %tensor_ptr, %tensor_off2 : tensor<4x!tt.ptr>, tensor<4xi32> - scf.yield %else_block, %else_tensor : !tt.ptr>, tensor<4x!tt.ptr> - } - tt.return %result#0, %result#1 : !tt.ptr>, tensor<4x!tt.ptr> - } -} - -// CHECK-LABEL: tt.func public @if_mixed_block_and_tensor_ptr -// CHECK: %[[RESULT:.*]]:2 = scf.if %{{.*}} -> (i32, tensor<4xi32>) { -// CHECK: scf.yield %{{.*}}, %{{.*}} : i32, tensor<4xi32> -// CHECK: } else { -// CHECK: scf.yield %{{.*}}, %{{.*}} : i32, tensor<4xi32> -// CHECK: } -// CHECK-DAG: %[[BLOCK:.*]] = tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[RESULT]]#0] {order = array} : > -// CHECK-DAG: %[[TENSOR:.*]] = tt.addptr %{{.*}}, %[[RESULT]]#1 : tensor<4x!tt.ptr>, tensor<4xi32> -// CHECK: tt.return %[[BLOCK]], %[[TENSOR]] : !tt.ptr>, tensor<4x!tt.ptr> - -// ----- - -module { - tt.func public @if_without_else_rewrites_nested_block_ptr(%base: !tt.ptr, %cond: i1) { - %c0_i32 = arith.constant 0 : i32 - %c1_i32 = arith.constant 1 : i32 - %c1_i64 = arith.constant 1 : i64 - %c16_i64 = arith.constant 16 : i64 - %c0 = arith.constant 0 : index - %c4 = arith.constant 4 : index - %c1 = arith.constant 1 : index - %ptr0 = tt.make_tensor_ptr %base, [%c16_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> - scf.if %cond { - %final = scf.for %iv = %c0 to %c4 step %c1 iter_args(%ptr = %ptr0) -> (!tt.ptr>) { - %next = tt.advance %ptr, [%c1_i32] : !tt.ptr> - scf.yield %next : !tt.ptr> - } - %loaded = tt.load %final : !tt.ptr> - scf.yield - } - tt.return - } -} - -// CHECK-LABEL: tt.func public @if_without_else_rewrites_nested_block_ptr -// CHECK: scf.if %{{.*}} { -// CHECK: %[[FOR:.*]] = scf.for -// CHECK-SAME: -> (i32) { -// CHECK: scf.yield %{{.*}} : i32 -// CHECK: } -// CHECK: %[[PTR:.*]] = tt.make_tensor_ptr %{{.*}}, [%{{.*}}], [%{{.*}}], [%[[FOR]]] {order = array} : > -// CHECK: tt.load %[[PTR]] : !tt.ptr> -// CHECK: } - -// ----- - -module { - tt.func public @sibling_if_result_initializes_for_tensor_ptr(%base: !tt.ptr, %cond: i1) -> tensor<4x!tt.ptr> { - %c0_i32 = arith.constant 0 : i32 - %c1_i32 = arith.constant 1 : i32 - %c2_i32 = arith.constant 2 : i32 - %c0 = arith.constant 0 : index - %c4 = arith.constant 4 : index - %c1 = arith.constant 1 : index - %off1 = tt.splat %c1_i32 : i32 -> tensor<4xi32> - %off2 = tt.splat %c2_i32 : i32 -> tensor<4xi32> - %ptr0 = tt.splat %base : !tt.ptr -> tensor<4x!tt.ptr> - %selected = scf.if %cond -> (tensor<4x!tt.ptr>) { - %then_ptr = tt.addptr %ptr0, %off1 : tensor<4x!tt.ptr>, tensor<4xi32> - scf.yield %then_ptr : tensor<4x!tt.ptr> - } else { - %else_ptr = tt.addptr %ptr0, %off2 : tensor<4x!tt.ptr>, tensor<4xi32> - scf.yield %else_ptr : tensor<4x!tt.ptr> - } - %final = scf.for %iv = %c0 to %c4 step %c1 iter_args(%ptr = %selected) -> (tensor<4x!tt.ptr>) { - %next = tt.addptr %ptr, %off1 : tensor<4x!tt.ptr>, tensor<4xi32> - scf.yield %next : tensor<4x!tt.ptr> - } - tt.return %final : tensor<4x!tt.ptr> - } -} - -// CHECK-LABEL: tt.func public @sibling_if_result_initializes_for_tensor_ptr -// CHECK: %[[IF_OFFSETS:.*]] = scf.if %{{.*}} -> (tensor<4xi32>) { -// CHECK: scf.yield %{{.*}} : tensor<4xi32> -// CHECK: } else { -// CHECK: scf.yield %{{.*}} : tensor<4xi32> -// CHECK: } -// CHECK: tt.addptr %{{.*}}, %[[IF_OFFSETS]] : tensor<4x!tt.ptr>, tensor<4xi32> -// CHECK: %[[FOR_OFFSETS:.*]] = scf.for -// CHECK-SAME: -> (tensor<4xi32>) { -// CHECK: scf.yield %{{.*}} : tensor<4xi32> -// CHECK: } -// CHECK: %[[FINAL:.*]] = tt.addptr %{{.*}}, %[[FOR_OFFSETS]] : tensor<4x!tt.ptr>, tensor<4xi32> -// CHECK: tt.return %[[FINAL]] : tensor<4x!tt.ptr> diff --git a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/tensor_ptr_different_base_invalid.mlir b/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/tensor_ptr_different_base_invalid.mlir deleted file mode 100644 index 1021b8a535..0000000000 --- a/third_party/ascend/unittest/Conversion/General/TritonControlFlowOpt/tensor_ptr_different_base_invalid.mlir +++ /dev/null @@ -1,22 +0,0 @@ -// RUN: not triton-opt --triton-control-flow-opt %s 2>&1 | FileCheck %s - -module { - tt.func public @if_tensor_ptr_different_base(%base0: !tt.ptr, %base1: !tt.ptr, %cond: i1) -> tensor<4x!tt.ptr> { - %c1_i32 = arith.constant 1 : i32 - %c2_i32 = arith.constant 2 : i32 - %off1 = tt.splat %c1_i32 : i32 -> tensor<4xi32> - %off2 = tt.splat %c2_i32 : i32 -> tensor<4xi32> - %splat0 = tt.splat %base0 : !tt.ptr -> tensor<4x!tt.ptr> - %splat1 = tt.splat %base1 : !tt.ptr -> tensor<4x!tt.ptr> - %selected = scf.if %cond -> (tensor<4x!tt.ptr>) { - %then_ptr = tt.addptr %splat0, %off1 : tensor<4x!tt.ptr>, tensor<4xi32> - scf.yield %then_ptr : tensor<4x!tt.ptr> - } else { - %else_ptr = tt.addptr %splat1, %off2 : tensor<4x!tt.ptr>, tensor<4xi32> - scf.yield %else_ptr : tensor<4x!tt.ptr> - } - tt.return %selected : tensor<4x!tt.ptr> - } -} - -// CHECK: error: failed to analyze pointer components across control flow