[WIP][BACKEND] sync Ascend backend with commit c8552741

include/incubated/Conversion/AutoBlockifyIncubated/AutoBlockify.h

@@ -0,0 +1,118 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "mlir/Pass/Pass.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+#include "mlir/IR/PatternMatch.h"
+
+#define GEN_PASS_DECL_AUTOBLOCKIFY
+#include "ascend/include/AutoBlockify/Passes.h.inc"
+
+#define GEN_PASS_DEF_AUTOBLOCKIFY
+#include "ascend/include/AutoBlockify/Passes.h.inc"
+
+namespace mlir {
+namespace triton {
+
+std::unique_ptr<OperationPass<ModuleOp>>
+createAutoBlockifyPass(const AutoBlockifyOptions &options = {});
+
+} // namespace triton
+} // namespace mlir
+
+using namespace mlir;
+using namespace triton;
+
+class PropagateUnrealizedCastDown
+    : public OpRewritePattern<UnrealizedConversionCastOp> {
+public:
+  using OpRewritePattern<UnrealizedConversionCastOp>::OpRewritePattern;
+
+  explicit PropagateUnrealizedCastDown(MLIRContext *context,
+                                       Value logicalBlockId,
+                                       Value logicalBlockNum,
+                                       int autoBlockifySize);
+
+  LogicalResult matchAndRewrite(UnrealizedConversionCastOp op,
+                                PatternRewriter &rewriter) const override;
+
+private:
+  void handleBlockifyLoop(scf::ForOp blockifyLoop, Operation *op, PatternRewriter &rewriter) const;
+  void rewriteSplat(UnrealizedConversionCastOp op, triton::SplatOp splatOp,
+                    PatternRewriter &rewriter) const;
+  void rewriteExpandDims(UnrealizedConversionCastOp op,
+                         triton::ExpandDimsOp expandDimsOp,
+                         PatternRewriter &rewriter) const;
+  void rewriteReduce(UnrealizedConversionCastOp op, triton::ReduceOp reduceOp,
+                     PatternRewriter &rewriter) const;
+  void rewriteScan(UnrealizedConversionCastOp op, triton::ScanOp scanOp,
+                   PatternRewriter &rewriter) const;
+  void rewriteLoad(UnrealizedConversionCastOp op, triton::LoadOp loadOp,
+                   PatternRewriter &rewriter) const;
+  void rewriteStore(UnrealizedConversionCastOp op, triton::StoreOp storeOp,
+                    PatternRewriter &rewriter) const;
+  void rewriteAtomicRMW(UnrealizedConversionCastOp op,
+                        triton::AtomicRMWOp atomicRMWOp,
+                        PatternRewriter &rewriter) const;
+  void rewriteAssert(UnrealizedConversionCastOp op, triton::AssertOp assertOp,
+                     PatternRewriter &rewriter) const;
+  void rewriteExtractSlice(UnrealizedConversionCastOp op,
+                           tensor::ExtractSliceOp extractSliceOp,
+                           PatternRewriter &rewriter) const;
+  void rewriteInsertSlice(UnrealizedConversionCastOp op,
+                          tensor::InsertSliceOp insertSliceOp,
+                          PatternRewriter &rewriter) const;
+  void rewriteWhile(UnrealizedConversionCastOp op, scf::WhileOp whileOp,
+                    PatternRewriter &rewriter) const;
+  void rewriteLoop(UnrealizedConversionCastOp op, LoopLikeOpInterface loopOp,
+                   PatternRewriter &rewriter) const;
+  void rewriteIf(UnrealizedConversionCastOp &op, scf::IfOp ifOp, ArrayRef<int64_t> indices,
+                 PatternRewriter &rewriter) const;
+  void rewriteYield(UnrealizedConversionCastOp &op, scf::YieldOp yieldOp,
+                    PatternRewriter &rewriter) const;
+  void rewriteCondition(UnrealizedConversionCastOp op,
+                        scf::ConditionOp conditionOp,
+                        PatternRewriter &rewriter) const;
+  void rewriteGeneraleOp(UnrealizedConversionCastOp op, Operation *generalOp,
+                         PatternRewriter &rewriter) const;
+
+  Value logicalBlockId;
+  Value logicalBlockNum;
+  int autoBlockifySize;
+};
+
+class AutoBlockifyPass : public ::impl::AutoBlockifyBase<AutoBlockifyPass> {
+public:
+  explicit AutoBlockifyPass(const AutoBlockifyOptions &options);
+  void runOnOperation() override;
+
+private:
+  bool checkBlockifiable(Value v);
+  void preProcess(triton::FuncOp func);
+
+  DenseSet<Value> checkedValues;
+  Value logicalBlockId;
+  Value logicalBlockNum;
+};

include/incubated/Conversion/AutoBlockifyIncubated/CMakeLists.txt

@@ -0,0 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls --name AutoBlockify)
+add_public_tablegen_target(AutoBlockifyPassIncGen)

include/incubated/Conversion/AutoBlockifyIncubated/Passes.h

@@ -0,0 +1,37 @@
+/*
+ * 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_ADAPTER_AUTO_BLOCKIFY_PASSES_H
+#define TRITON_ADAPTER_AUTO_BLOCKIFY_PASSES_H
+
+#include "AutoBlockify.h"
+
+namespace mlir {
+namespace triton {
+
+#define GEN_PASS_REGISTRATION
+#include "ascend/include/AutoBlockify/Passes.h.inc"
+
+} // namespace triton
+} // namespace mlir
+
+#endif // TRITON_ADAPTER_AUTO_BLOCKIFY_PASSES_H

include/incubated/Conversion/AutoBlockifyIncubated/Passes.td

@@ -0,0 +1,21 @@
+#ifndef AUTO_BLOCKIFY_PASSES
+#define AUTO_BLOCKIFY_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def AutoBlockify : Pass<"auto-blockify", "mlir::ModuleOp"> {
+    let summary = "Apply auto blockify v2";
+    let constructor = "triton::createAutoBlockifyPass()";
+    let dependentDialects = [
+        "mlir::arith::ArithDialect",
+        "mlir::tensor::TensorDialect",
+        "mlir::triton::TritonDialect"
+    ];
+    let options = [
+        Option<"autoBlockifySize", "auto-blockify-size", "int", "1", 
+           "Apply auto blockify v2 when TRITON_ALL_BLOCKS_PARALLEL is 1."
+           "Expand highest dimension with blockify size">
+    ];
+}
+
+#endif // AUTO_BLOCKIFY_PASSES

include/incubated/Conversion/AutoBlockifyIncubated/Utils.h

@@ -0,0 +1,66 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "triton/Dialect/Triton/IR/Dialect.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+
+using namespace mlir;
+using namespace triton;
+
+constexpr llvm::StringLiteral autoBlockifySizeAttr = "auto_blockify_size";
+constexpr llvm::StringLiteral logicalBlockIdAttr = "logical_block_id";
+constexpr llvm::StringLiteral autoBlockifyLoopAttr =
+    "auto_blockify_loop";
+constexpr llvm::StringLiteral autoBlockifyRegionOpAttr =
+    "auto_blockify_region_op";
+
+RankedTensorType getExpandedType(Type type, UnrealizedConversionCastOp op);
+
+Value rewriteValue(Value value, UnrealizedConversionCastOp op,
+                   OpBuilder &builder);
+
+void replaceValue(Operation *newOp, Operation *oldOp, Value newMask,
+                  RewriterBase &rewriter,
+                  ArrayRef<int64_t> replaceIndices = {});
+
+Value createMask(Value mask, Value uccMask, ArrayRef<int64_t> targetShape,
+                 RewriterBase &rewriter);
+
+void mapRegionIterArg(IRMapping &mapping, ValueRange oldArgs,
+                      ValueRange newArgs, ArrayRef<int64_t> indices, Value mask,
+                      OpBuilder &builder);
+
+void mapYieldedValue(IRMapping &mapping, scf::YieldOp yieldOp,
+                     ArrayRef<int64_t> indices, UnrealizedConversionCastOp op,
+                     OpBuilder &builder);
+
+Operation *createBlockifyLoop(Operation *targetOp,
+                              UnrealizedConversionCastOp op,
+                              Value logicalBlockId, Value logicalBlockNum,
+                              int autoBlockifySize, RewriterBase &rewriter);
+
+std::optional<scf::ForOp> getBlockifyLoop(Operation *op);

include/incubated/Conversion/DynamicCVPipelineIncubated/AddControlFlowCondition.h

@@ -0,0 +1,69 @@
+/*
+ * 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_ADAPTER_DYNAMIC_CV_PIPELINE_ADD_CONTROLFLOW_CONDITION_PASS_H
+#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_ADD_CONTROLFLOW_CONDITION_PASS_H
+
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/SmallPtrSet.h"
+
+namespace mlir {
+namespace triton {
+
+struct ControlFlowConditionInfo {
+  llvm::DenseMap<scf::ForOp, SmallVector<int>> blockCounters;
+  llvm::DenseMap<scf::ForOp, int> blockCounterNums;
+  llvm::DenseMap<scf::ForOp, SmallVector<int>> innerDepConds;
+
+  llvm::DenseMap<Value, SmallVector<Value>> crossCoreDependentMap;
+  llvm::DenseMap<scf::ForOp, llvm::DenseMap<Value, SmallVector<Value>>> intraCoreDependentMap;
+
+  // unique counter value for each ifblock
+  llvm::DenseMap<scf::IfOp, Value> cntArgs;
+};
+
+class AddControlFlowConditionPass
+    : public PassWrapper<AddControlFlowConditionPass, OperationPass<ModuleOp>> {
+public:
+  AddControlFlowConditionPass() = default;
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override
+  {
+    registry.insert<LLVM::LLVMDialect>();
+  }
+
+  llvm::StringRef getArgument() const override { return "add-control-flow-condition"; }
+};
+
+std::unique_ptr<OperationPass<ModuleOp>> createAddControlFlowConditionPass();
+
+void registerAddControlFlowConditionPasses();
+} // namespace triton
+} // namespace mlir
+#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_ADD_CONTROLFLOW_CONDITION_PASS_H

include/incubated/Conversion/DynamicCVPipelineIncubated/AddControlFlowCondition/CloneOps.h

@@ -0,0 +1,53 @@
+/*
+ * 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_SSBUF_CLONE_OPS_FOR_CONTROL_FLOW_H
+#define TRITON_ASCEND_SSBUF_CLONE_OPS_FOR_CONTROL_FLOW_H
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/DialectRegistry.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace triton {
+
+class CloneOpsPass : public PassWrapper<CloneOpsPass, OperationPass<ModuleOp>> {
+public:
+  CloneOpsPass() = default;
+
+  void runOnOperation() override;
+
+  LogicalResult validateBlockIdsConsecutive(ModuleOp module);
+  LogicalResult cloneOpsInMainLoop(scf::ForOp forOp);
+  LogicalResult cleanupClonedOpsInMainLoop(scf::ForOp forOp);
+
+  llvm::StringRef getArgument() const override
+  {
+    return "clone-ops";
+  }
+};
+
+std::unique_ptr<OperationPass<ModuleOp>> createCloneOpsPass();
+
+} // namespace triton
+} // namespace mlir
+#endif // TRITON_ASCEND_SSBUF_CLONE_OPS_FOR_CONTROL_FLOW_H