Skip to content

Commit 9a1b135

Browse files
zhzhcookiewcleungaj-70
zhzhcookie
and
wcleungaj-70
committed
[Ascend] Support ascend3.5 from TA with commit c8552741 (#56)
* [BACKEND] sync Ascend backend with commit c8552741 * [Incubated] Fix sync from TA and support building with llvm22 (#55) * [TLE] Change __TLE__ to __TLE_STRUCT__ * [Incubated] Fix createTritonToLinalgIncubatedPass * [Incubated] Sync triton3.5 code and fix include path (#57) * [CodeFormat] Fix by pre-commit --------- Co-authored-by: wcleungaj-70 <leungwingchung70@gmail.com>
1 parent 2756c17 commit 9a1b135

84 files changed

Lines changed: 7175 additions & 2510 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.pre-commit-config.yaml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@ repos:
2020
(?x)^(
2121
include/incubated/.*\.(h|hpp|cc|cpp)|
2222
include/mlir-ext/.*\.(h|hpp|cc|cpp)|
23-
include/npu/.*\.(h|hpp|cc|cpp)|
2423
lib/UtilsIncubated/.*\.(h|hpp|cc|cpp)|
25-
lib/Dialect/(MathExt|TritonAscend|TritonStructuredIncubated)/.*\.(h|hpp|cc|cpp)|
26-
lib/Conversion/(DiscreteMaskAccessConversion|NoBufferize_FlagTree|TritonToAnnotation|TritonToLinalgIncubated|TritonToStructuredIncubated|TritonToUnstructureIncubated)/.*\.(h|hpp|cc|cpp)|
24+
lib/Dialect/(MathExt|TritonStructuredIncubated)/.*\.(h|hpp|cc|cpp)|
25+
lib/Conversion/DiscreteMaskAccessConversion/.*\.(h|hpp|cc|cpp)|
26+
lib/Conversion/.*(Incubated|FlagTree)/.*\.(h|hpp|cc|cpp)|
2727
.*_FlagTree\.(h|hpp|cc|cpp)
2828
)$
2929
stages: [pre-commit, pre-push, manual]

‎CMakeLists.txt‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ endif()
1515
add_subdirectory(include)
1616
add_subdirectory(lib)
1717

18-
if(NOT FLAGTREE_BACKEND STREQUAL "tsingmicro")
18+
if(NOT FLAGTREE_BACKEND)
1919
add_subdirectory(test)
2020
add_subdirectory(tools)
2121
endif()
@@ -24,4 +24,3 @@ if (TRITON_SHARED_BUILD_CPU_BACKEND)
2424
add_triton_plugin(TritonShared ${CMAKE_CURRENT_SOURCE_DIR}/triton_shared.cc LINK_LIBS TritonSharedAnalysis TritonToLinalg TritonTilingExtIR)
2525
endif()
2626

27-

‎include/CMakeLists.txt‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
add_subdirectory(mlir-ext)
22
add_subdirectory(triton-shared)
3-
add_subdirectory(npu)
43
if (FLIR_BUILD_INCUBATED)
54
add_subdirectory(incubated)
65
endif()

‎include/incubated/Conversion/CMakeLists.txt‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
add_subdirectory(TritonToAnnotation)
21
add_subdirectory(TritonToLinalgIncubated)
32
add_subdirectory(DiscreteMaskAccessConversion)
43
add_subdirectory(TritonToUnstructureIncubated)

‎include/incubated/Conversion/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.h‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,6 @@
3434
#define GEN_PASS_DEF_DISCRETEMASKACCESSCONVERSION
3535
#include "incubated/Conversion/DiscreteMaskAccessConversion/Passes.h.inc"
3636

37-
extern bool compileOn91095Flag;
38-
extern bool forceSimtTemplateFlag;
39-
4037
namespace mlir {
4138
namespace triton {
4239

@@ -57,7 +54,7 @@ class DiscreteMaskAccessConversionPass
5754
public:
5855
explicit DiscreteMaskAccessConversionPass(
5956
const DiscreteMaskAccessConversionOptions &options);
60-
57+
void getDependentDialects(DialectRegistry &registry) const override;
6158
void runOnOperation() override;
6259
};
6360

‎include/incubated/Conversion/DiscreteMaskAccessConversion/Passes.td‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,15 @@ def DiscreteMaskAccessConversion : Pass<"discrete-mask-access-conversion", "mlir
1212
let summary = "Recognize and convert discrete mask memory access";
1313
let constructor = "triton::createDiscreteMaskAccessConversionPass()";
1414
let options = [
15-
Option<"compileOn91095", "compile-on-910-95",
15+
Option<"compileOn91095", "compile-on-910-95",
1616
"bool", /*default*/"false",
1717
"compile on 910_95">,
18-
Option<"forceSimtTemplate", "force-simt-template",
18+
Option<"forceSimtTemplate", "force-simt-template",
1919
"bool", /*default*/"false",
20-
"force to use simt template">
20+
"force to use simt template">,
21+
Option<"enableSyncBlockLock", "enable-sync-block-lock",
22+
"bool", /*default*/"true",
23+
"enable sync block lock/unlock in store">
2124
];
2225
}
2326

‎include/incubated/Conversion/TritonToAnnotation/CMakeLists.txt‎

Lines changed: 0 additions & 3 deletions
This file was deleted.

‎include/incubated/Conversion/TritonToAnnotation/Passes.td‎

Lines changed: 0 additions & 17 deletions
This file was deleted.

‎include/incubated/Conversion/TritonToLinalgIncubated/ArgMinMaxConverter.h‎

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,11 @@
2323
#ifndef TRITON_ADAPTER_ARGMINMAXCONVERTER_H
2424
#define TRITON_ADAPTER_ARGMINMAXCONVERTER_H
2525

26+
#include "incubated/Conversion/TritonToLinalgIncubated/ConversionPatterns.h"
2627
#include "incubated/Conversion/UtilsIncubated/Utils.h"
2728

2829
#include "triton/Dialect/Triton/IR/Dialect.h"
2930

30-
#include "ConversionPatterns.h"
3131
#include "mlir/Dialect/Linalg/IR/Linalg.h"
3232
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
3333
#include "mlir/Interfaces/FunctionInterfaces.h"
@@ -36,6 +36,7 @@
3636
#define DEBUG_TYPE "triton-to-linalg"
3737

3838
#include "llvm/Support/Debug.h"
39+
#include "llvm/Support/LogicalResult.h"
3940

4041
#include <limits>
4142

@@ -136,7 +137,9 @@ class ArgMinMaxBaseConverter : public OpConversionPattern<triton::ReduceOp> {
136137
public:
137138
ArgMinMaxBaseConverter(MLIRContext *context) : OpConversionPattern(context) {}
138139

139-
LogicalResult match(triton::ReduceOp op) const override final {
140+
LogicalResult
141+
matchAndRewrite(triton::ReduceOp op, OpAdaptor adaptor,
142+
ConversionPatternRewriter &rewriter) const override {
140143
if (op.getBody()->getNumArguments() != 4) {
141144
return failure();
142145
}
@@ -166,15 +169,6 @@ class ArgMinMaxBaseConverter : public OpConversionPattern<triton::ReduceOp> {
166169

167170
LLVM_DEBUG(llvm::dbgs() << "Matching: " << *opsIt << "\n");
168171
auto indexSelectOp = dyn_cast<arith::SelectOp>(*opsIt++);
169-
if (indexSelectOp) {
170-
if (indexSelectOp.getCondition() != shouldUpdate ||
171-
currIndex != indexSelectOp.getTrueValue() ||
172-
reduceIndex != indexSelectOp.getFalseValue()) {
173-
return failure();
174-
}
175-
} else {
176-
return failure();
177-
}
178172
if (!indexSelectOp || indexSelectOp.getCondition() != shouldUpdate ||
179173
currIndex != indexSelectOp.getTrueValue() ||
180174
reduceIndex != indexSelectOp.getFalseValue()) {
@@ -188,17 +182,13 @@ class ArgMinMaxBaseConverter : public OpConversionPattern<triton::ReduceOp> {
188182
ArrayRef<Value>{valueSelectOp, indexSelectOp})) {
189183
return failure();
190184
}
191-
return success();
192-
}
193185

194-
void rewrite(triton::ReduceOp op, OpAdaptor adaptor,
195-
ConversionPatternRewriter &rewriter) const override final {
186+
// Rewrite phase: perform the actual conversion
196187
auto loc = op.getLoc();
197188
auto elemTypes = op.getElementTypes();
198189

199190
auto valueType = elemTypes[0];
200191
// tl.argmin reorder
201-
auto block = op.getBody();
202192
bool isUnsigned = false;
203193
if (isa<mlir::FloatType>(valueType)) {
204194
arith::CmpFOp cmpFOp;
@@ -262,8 +252,8 @@ class ArgMinMaxBaseConverter : public OpConversionPattern<triton::ReduceOp> {
262252
auto valuesAccBaseVal =
263253
rewriter.create<arith::ConstantOp>(loc, valueType, valueAttr);
264254
int indicesInitValue =
265-
(reduceWithIndexParams.has_value() &&
266-
(*reduceWithIndexParams).tieBreakType == TieBreakType::RIGHT)
255+
(llvm::succeeded(reduceWithIndexParams) &&
256+
reduceWithIndexParams->tieBreakType == TieBreakType::RIGHT)
267257
? -1
268258
: std::numeric_limits<int32_t>::max();
269259

@@ -306,7 +296,8 @@ class ArgMinMaxBaseConverter : public OpConversionPattern<triton::ReduceOp> {
306296
// before we rewrite the argmax reduce op, we know it has return value
307297
// so addReduceWithIndexAttrIfNeeded won't fail
308298
// but ignoring it will lead to compiling failure
309-
if (reduceWithIndexParams.has_value()) {
299+
if (llvm::succeeded(reduceWithIndexParams) &&
300+
reduceWithIndexParams->tieBreakType != TieBreakType::None) {
310301
addReduceWithIndexAttr(*reduceWithIndexParams, rewriter, linalgOp);
311302
}
312303

@@ -320,6 +311,8 @@ class ArgMinMaxBaseConverter : public OpConversionPattern<triton::ReduceOp> {
320311
} else {
321312
rewriter.replaceOp(op, linalgOp);
322313
}
314+
315+
return success();
323316
}
324317
};
325318

‎include/incubated/Conversion/TritonToLinalgIncubated/BlockPtrAnalysis.h‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222

2323
#ifndef TRITON_ANALYSIS_BLOCKPTRANALYSIS_H
2424
#define TRITON_ANALYSIS_BLOCKPTRANALYSIS_H
25-
25+
#if __has_include("bishengir/Dialect/HIVM/IR/HIVM.h")
26+
#include "bishengir/Dialect/HIVM/IR/HIVM.h"
27+
#endif
2628
#include "mlir/Dialect/Arith/IR/Arith.h"
2729
#include "mlir/Dialect/Linalg/IR/Linalg.h"
2830
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -184,6 +186,11 @@ class BlockDataParser {
184186
ConversionPatternRewriter &rewriter,
185187
const llvm::SmallDenseMap<Value, BlockData> &known);
186188

189+
static void parseLinalgGenericFromMakeRange(
190+
linalg::GenericOp op, BlockData &data, const Location &loc,
191+
ConversionPatternRewriter &rewriter,
192+
const llvm::SmallDenseMap<Value, BlockData> &known);
193+
187194
static void
188195
parseExpandDims(triton::ExpandDimsOp op, BlockData &data, const Location &loc,
189196
ConversionPatternRewriter &rewriter,
@@ -241,6 +248,11 @@ class BlockDataParser {
241248
ConversionPatternRewriter &rewriter,
242249
const llvm::SmallDenseMap<Value, BlockData> &known);
243250

251+
static void
252+
parseAtomicRmw(triton::AtomicRMWOp op, BlockData &data, const Location &loc,
253+
ConversionPatternRewriter &rewriter,
254+
const llvm::SmallDenseMap<Value, BlockData> &known);
255+
244256
static void parseFill(linalg::FillOp op, BlockData &data, const Location &loc,
245257
ConversionPatternRewriter &rewriter,
246258
const llvm::SmallDenseMap<Value, BlockData> &known);
@@ -250,6 +262,12 @@ class BlockDataParser {
250262
ConversionPatternRewriter &rewriter,
251263
const llvm::SmallDenseMap<Value, BlockData> &known);
252264

265+
static void parseCustomOp(hivm::CustomOp op, BlockData &data,
266+
const Location &loc,
267+
ConversionPatternRewriter &rewriter,
268+
const llvm::SmallDenseMap<Value, BlockData> &known,
269+
unsigned resultIdx);
270+
253271
static void rewriteAddPtr(triton::AddPtrOp op,
254272
triton::AddPtrOp::Adaptor &adaptor,
255273
ConversionPatternRewriter &rewriter,
@@ -264,6 +282,11 @@ class BlockDataParser {
264282
ConversionPatternRewriter &rewriter,
265283
llvm::SmallDenseMap<Value, BlockData> &known);
266284

285+
static void
286+
rewriteCustomOp(hivm::CustomOp op, hivm::CustomOp::Adaptor &adaptor,
287+
ConversionPatternRewriter &rewriter,
288+
const llvm::SmallDenseMap<Value, BlockData> &known);
289+
267290
template <typename T>
268291
static std::enable_if_t<std::is_same_v<T, scf::YieldOp> ||
269292
std::is_same_v<T, scf::ConditionOp>>
@@ -287,7 +310,8 @@ class BlockDataParser {
287310
template <typename OpTy>
288311
void parseIndirectLoad(OpTy op, BlockData &data, const Location &loc,
289312
ConversionPatternRewriter &rewriter,
290-
const llvm::SmallDenseMap<Value, BlockData> &known);
313+
const llvm::SmallDenseMap<Value, BlockData> &known,
314+
unsigned resultIdx = 0);
291315

292316
} // namespace triton
293317

0 commit comments

Comments
 (0)