feat: support (built-in)indirect_load for unstructure load-triton-ascend-AtomGit

Mascend/backend/compiler.py+1-1

	@@ -108,7 +108,7 @@ def ttir_to_linalg(mod, metadata, opt, *, named_ops=False):
108	"--triton-linearize",	108	"--triton-linearize",
109	"--discrete-mask-access-conversion",	109	"--discrete-mask-access-conversion",
110	"--triton-to-annotation",	110	"--triton-to-annotation",
111	- "--triton-to-unstructure",	111	+ f"--triton-to-unstructure=compile-on-a5={compile_on_a5}",
112	"--triton-to-hivm",	112	"--triton-to-hivm",
113	"--triton-to-hfusion",	113	"--triton-to-hfusion",
114	"--triton-to-llvm",	114	"--triton-to-llvm",

Aascend/test/Conversion/TritonToUnstructure/indirect_load.mlir+73-0

@@ -0,0 +1,73 @@
	1	+// RUN: triton-adapter-opt --triton-to-unstructure=compile-on-a5=true %s \| FileCheck %s
	2	+
	3	+tt.func public @triton_ldst_indirect_05_kernel(%arg0: !tt.ptr<f32>, %arg1: !tt.ptr<i64>, %arg2: !tt.ptr<f32>, %arg3: i32) attributes {noinline = false} {
	4	+ %cst = arith.constant dense<16> : tensor<8x1xi32>
	5	+ %c8_i32 = arith.constant 8 : i32
	6	+ %0 = tt.get_program_id x : i32
	7	+ %1 = arith.muli %0, %c8_i32 : i32
	8	+ %2 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32>
	9	+ %3 = tt.splat %1 : i32 -> tensor<8xi32>
	10	+ %4 = arith.addi %3, %2 : tensor<8xi32>
	11	+ %5 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32>
	12	+ %6 = tt.splat %arg1 : !tt.ptr<i64> -> tensor<16x!tt.ptr<i64>>
	13	+ %7 = tt.addptr %6, %5 : tensor<16x!tt.ptr<i64>>, tensor<16xi32>
	14	+ %8 = tt.load %7 : tensor<16x!tt.ptr<i64>>
	15	+ %9 = tt.expand_dims %2 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32>
	16	+ %10 = tt.splat %arg3 : i32 -> tensor<8x1xi32>
	17	+ %11 = arith.muli %9, %10 : tensor<8x1xi32>
	18	+ %12 = tt.expand_dims %8 {axis = 0 : i32} : tensor<16xi64> -> tensor<1x16xi64>
	19	+ %13 = arith.extsi %11 : tensor<8x1xi32> to tensor<8x1xi64>
	20	+ %14 = tt.broadcast %13 : tensor<8x1xi64> -> tensor<8x16xi64>
	21	+ %15 = tt.broadcast %12 : tensor<1x16xi64> -> tensor<8x16xi64>
	22	+ %16 = arith.addi %14, %15 : tensor<8x16xi64>
	23	+ %17 = tt.splat %arg2 : !tt.ptr<f32> -> tensor<8x16x!tt.ptr<f32>>
	24	+ %18 = tt.addptr %17, %16 : tensor<8x16x!tt.ptr<f32>>, tensor<8x16xi64>
	25	+ %19 = tt.load %18 : tensor<8x16x!tt.ptr<f32>>
	26	+ %20 = math.exp %19 : tensor<8x16xf32>
	27	+ %21 = tt.expand_dims %4 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32>
	28	+ %22 = arith.muli %21, %cst : tensor<8x1xi32>
	29	+ %23 = tt.expand_dims %5 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32>
	30	+ %24 = tt.broadcast %22 : tensor<8x1xi32> -> tensor<8x16xi32>
	31	+ %25 = tt.broadcast %23 : tensor<1x16xi32> -> tensor<8x16xi32>
	32	+ %26 = arith.addi %24, %25 : tensor<8x16xi32>
	33	+ %27 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<8x16x!tt.ptr<f32>>
	34	+ %28 = tt.addptr %27, %26 : tensor<8x16x!tt.ptr<f32>>, tensor<8x16xi32>
	35	+ tt.store %28, %20 : tensor<8x16x!tt.ptr<f32>>
	36	+ tt.return
	37	+ }
	38	+
	39	+
	40	+// CHECK-LABEL: tt.func public @triton_ldst_indirect_05_kernel(
	41	+// CHECK-SAME: %[[VAL_0:.]]: !tt.ptr<f32>, %[[VAL_1:.]]: !tt.ptr<i64>, %[[VAL_2:.]]: !tt.ptr<f32>, %[[VAL_3:.]]: i32) attributes {noinline = false} {
	42	+// CHECK: %[[VAL_4:.*]] = arith.constant dense<16> : tensor<8x1xi32>
	43	+// CHECK: %[[VAL_5:.*]] = arith.constant 8 : i32
	44	+// CHECK: %[[VAL_6:.*]] = tt.get_program_id x : i32
	45	+// CHECK: %[[VAL_7:.*]] = arith.muli %[[VAL_6]], %[[VAL_5]] : i32
	46	+// CHECK: %[[VAL_8:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32>
	47	+// CHECK: %[[VAL_9:.*]] = tt.splat %[[VAL_7]] : i32 -> tensor<8xi32>
	48	+// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_8]] : tensor<8xi32>
	49	+// CHECK: %[[VAL_11:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32>
	50	+// CHECK: %[[VAL_12:.]] = tt.splat %[[VAL_1:.]] : !tt.ptr<i64> -> tensor<16x!tt.ptr<i64>>
	51	+// CHECK: %[[VAL_13:.*]] = tt.addptr %[[VAL_12]], %[[VAL_11]] : tensor<16x!tt.ptr<i64>>, tensor<16xi32>
	52	+// CHECK: %[[VAL_14:.*]] = tt.load %[[VAL_13]] : tensor<16x!tt.ptr<i64>>
	53	+// CHECK: %[[VAL_15:.*]] = tt.expand_dims %[[VAL_8]] {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32>
	54	+// CHECK: %[[VAL_16:.]] = tt.splat %[[VAL_3:.]] : i32 -> tensor<8x1xi32>
	55	+// CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_15]], %[[VAL_16]] : tensor<8x1xi32>
	56	+// CHECK: %[[VAL_18:.*]] = tt.expand_dims %[[VAL_14]] {axis = 0 : i32} : tensor<16xi64> -> tensor<1x16xi64>
	57	+// CHECK: %[[VAL_19:.*]] = arith.extsi %[[VAL_17]] : tensor<8x1xi32> to tensor<8x1xi64>
	58	+// CHECK: %[[VAL_20:.*]] = tt.broadcast %[[VAL_19]] : tensor<8x1xi64> -> tensor<8x16xi64>
	59	+// CHECK: %[[VAL_21:.*]] = tt.broadcast %[[VAL_18]] : tensor<1x16xi64> -> tensor<8x16xi64>
	60	+// CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_20]], %[[VAL_21]] : tensor<8x16xi64>
	61	+// CHECK: %[[VAL_23:.]] = tt.indirect_load %[[VAL_2:.]] : <f32>, %[[VAL_22]] : tensor<8x16xi64> -> tensor<8x16xf32>
	62	+// CHECK: %[[VAL_24:.*]] = math.exp %[[VAL_23]] : tensor<8x16xf32>
	63	+// CHECK: %[[VAL_25:.*]] = tt.expand_dims %[[VAL_10]] {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32>
	64	+// CHECK: %[[VAL_26:.*]] = arith.muli %[[VAL_25]], %[[VAL_4]] : tensor<8x1xi32>
	65	+// CHECK: %[[VAL_27:.*]] = tt.expand_dims %[[VAL_11]] {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32>
	66	+// CHECK: %[[VAL_28:.*]] = tt.broadcast %[[VAL_26]] : tensor<8x1xi32> -> tensor<8x16xi32>
	67	+// CHECK: %[[VAL_29:.*]] = tt.broadcast %[[VAL_27]] : tensor<1x16xi32> -> tensor<8x16xi32>
	68	+// CHECK: %[[VAL_30:.*]] = arith.addi %[[VAL_28]], %[[VAL_29]] : tensor<8x16xi32>
	69	+// CHECK: %[[VAL_31:.]] = tt.splat %[[VAL_0:.]] : !tt.ptr<f32> -> tensor<8x16x!tt.ptr<f32>>
	70	+// CHECK: %[[VAL_32:.*]] = tt.addptr %[[VAL_31]], %[[VAL_30]] : tensor<8x16x!tt.ptr<f32>>, tensor<8x16xi32>
	71	+// CHECK: tt.store %[[VAL_32]], %[[VAL_24]] : tensor<8x16x!tt.ptr<f32>>
	72	+// CHECK: tt.return
	73	+// CHECK: }

Mascend/triton-adapter/include/TritonToLinalg/TritonOpConverter.h+10-0

	@@ -576,6 +576,16 @@ private:
576	static constexpr llvm::StringRef funcNameBase = "triton_embedding_gather";	576	static constexpr llvm::StringRef funcNameBase = "triton_embedding_gather";
577	};	577	};
578		578
		579	+class IndirectLoadConverter : public OpConversionPattern<triton::IndirectLoadOp> {
		580	+public:
		581	+ using OpConversionPattern<triton::IndirectLoadOp>::OpConversionPattern;
		582	+ LogicalResult
		583	+ matchAndRewrite(triton::IndirectLoadOp op, OpAdaptor adaptor,
		584	+ ConversionPatternRewriter &rewriter) const override;
		585	+private:
		586	+ static constexpr llvm::StringRef funcNameBase = "triton_indirect_load";
		587	+};
		588	+
579	} // end of namespace TTOpConverters	589	} // end of namespace TTOpConverters
580		590
581	#endif	591	#endif

Mascend/triton-adapter/include/TritonToUnstructure/Passes.td+3-0

	@@ -9,6 +9,9 @@ def TritonToUnstructure : Pass<"triton-to-unstructure", "mlir::ModuleOp"> {
9	let options = [	9	let options = [
10	Option<"forceScalarizeMode", "force-scalarize-mode", "bool", "false",	10	Option<"forceScalarizeMode", "force-scalarize-mode", "bool", "false",
11	"Scalarize unstructured memory access even if structured dimensions are mixed.">,	11	"Scalarize unstructured memory access even if structured dimensions are mixed.">,
		12	+ Option<"compileOnA5", "compile-on-a5",
		13	+ "bool", /default/"false",
		14	+ "compile on a5">
12	];	15	];
13	}	16	}
14		17

Mascend/triton-adapter/include/TritonToUnstructure/UnstructureConversionPass.h+2-0

	@@ -35,6 +35,8 @@
35	#define GEN_PASS_DEF_TRITONTOUNSTRUCTURE	35	#define GEN_PASS_DEF_TRITONTOUNSTRUCTURE
36	#include "ascend/triton-adapter/include/TritonToUnstructure/Passes.h.inc"	36	#include "ascend/triton-adapter/include/TritonToUnstructure/Passes.h.inc"
37		37
		38	+extern bool compileOnA5Flag;
		39	+
38	namespace mlir {	40	namespace mlir {
39	namespace triton {	41	namespace triton {
40		42

Mascend/triton-adapter/lib/TritonToLinalg/TritonOpConverter.cpp+44-0

	@@ -2036,4 +2036,48 @@ EmbeddingGatherConverter::matchAndRewrite(triton::EmbeddingGatherOp op, OpAdapto
2036	return success();	2036	return success();
2037	}	2037	}
2038		2038
		2039	+
		2040	+LogicalResult
		2041	+IndirectLoadConverter::matchAndRewrite(triton::IndirectLoadOp op, OpAdaptor adaptor,
		2042	+ ConversionPatternRewriter &rewriter) const {
		2043	+ auto loc = op.getLoc();
		2044	+
		2045	+ auto moduleOp = op->getParentOfType<ModuleOp>();
		2046	+ rewriter.setInsertionPoint(moduleOp.getBody(),
		2047	+ std::prev(moduleOp.getBody()->end()));
		2048	+
		2049	+ llvm::SmallString<128> funcName = funcNameBase;
		2050	+ int uniqueId = 0;
		2051	+ while (SymbolTable::lookupSymbolIn(moduleOp, funcName)) {
		2052	+ funcName = funcNameBase;
		2053	+ funcName += ("_" + std::to_string(uniqueId++));
		2054	+ }
		2055	+
		2056	+ auto src = adaptor.getSrc();
		2057	+ auto offsets = op.getOffsets();
		2058	+ auto mask = op.getMask();
		2059	+ auto other = op.getOther();
		2060	+ auto res = op.getResult();
		2061	+ auto resTy = res.getType();
		2062	+
		2063	+ // convert !tt.ptr<f32> to memref<?xf32>
		2064	+ auto srcTy = cast<MemRefType>(src.getType());
		2065	+ SmallVector<Type> inputTypes({srcTy, offsets.getType()});
		2066	+ if (mask) inputTypes.push_back(mask.getType());
		2067	+ if (other) inputTypes.push_back(other.getType());
		2068	+ auto libFnType = rewriter.getFunctionType(inputTypes, {resTy});
		2069	+ auto funcOp = rewriter.create<func::FuncOp>(loc, funcName.str(), libFnType);
		2070	+ SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private);
		2071	+
		2072	+ rewriter.setInsertionPoint(op);
		2073	+ SmallVector<Value> inputVals({src, offsets});
		2074	+ if (mask) inputVals.push_back(mask);
		2075	+ if (other) inputVals.push_back(other);
		2076	+ auto callOp = rewriter.create<func::CallOp>(loc, funcOp.getSymNameAttr(),
		2077	+ TypeRange({resTy}),
		2078	+ inputVals);
		2079	+ rewriter.replaceOp(op, callOp);
		2080	+ return success();
		2081	+}
		2082	+
2039	} // namespace TTOpConverters	2083	} // namespace TTOpConverters

Mascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp+1-0

	@@ -584,6 +584,7 @@ void TritonToLinalgPass::populateTritonToLinalgConversionPatterns(
584	patterns.add<TTOpConverters::SortOpConverter>(patterns.getContext());	584	patterns.add<TTOpConverters::SortOpConverter>(patterns.getContext());
585	patterns.add<TTOpConverters::DotScaledConverter>(patterns.getContext());	585	patterns.add<TTOpConverters::DotScaledConverter>(patterns.getContext());
586	patterns.add<TTOpConverters::PtrToIntConverter>(patterns.getContext());	586	patterns.add<TTOpConverters::PtrToIntConverter>(patterns.getContext());
		587	+ patterns.add<TTOpConverters::IndirectLoadConverter>(patterns.getContext());
587		588
588	if (!this->namedOps) {	589	if (!this->namedOps) {
589	linalg::populateElementwiseToLinalgConversionPatterns(patterns);	590	linalg::populateElementwiseToLinalgConversionPatterns(patterns);

Mascend/triton-adapter/lib/TritonToUnstructure/UnstructureConversionPass.cpp+16-0

	@@ -297,6 +297,20 @@ LogicalResult UnstructuredMemAccessConverter<MemAccOpTy>::matchAndRewrite(
297	// Force scalarize if memory is not aligned	297	// Force scalarize if memory is not aligned
298	if (sizeInByte % 32 != 0)	298	if (sizeInByte % 32 != 0)
299	ptrOffsetInfo.setUnstructured(ptrOffsetInfo.getRank());	299	ptrOffsetInfo.setUnstructured(ptrOffsetInfo.getRank());
		300	+
		301	+ // Fast path on A5: rewrite tl.load to tt.indirect_load directly.
		302	+ if constexpr (std::is_same_v<MemAccOpTy, triton::LoadOp>) {
		303	+ if (compileOnA5Flag && ptrOffsetInfo.isUnstructured()) {
		304	+ assert(!isa<RankedTensorType>(srcPtr.getType()) && "src must be ptr type");
		305	+ Value mask = op.getMask();
		306	+ Value other = op.getOther();
		307	+ auto resultType = op.getType();
		308	+ auto indirect = rewriter.create<triton::IndirectLoadOp>(
		309	+ loc, resultType, srcPtr, ptrOffset, mask, other);
		310	+ rewriter.replaceOp(op, indirect.getResult());
		311	+ return success();
		312	+ }
		313	+ }
300		314
301	Value iterArg = nullptr;	315	Value iterArg = nullptr;
302		316
	@@ -572,6 +586,8 @@ TritonToUnstructurePass::TritonToUnstructurePass(
572	: TritonToUnstructureBase(options) {}	586	: TritonToUnstructureBase(options) {}
573		587
574	void TritonToUnstructurePass::runOnOperation() {	588	void TritonToUnstructurePass::runOnOperation() {
		589	+ compileOnA5Flag = this->compileOnA5;
		590	+
575	ModuleOp moduleOp = getOperation();	591	ModuleOp moduleOp = getOperation();
576	MLIRContext *ctx = &getContext();	592	MLIRContext *ctx = &getContext();
577		593

Mtriton_patch/include/triton/Dialect/Triton/IR/TritonOps.td+50-0

	@@ -1515,4 +1515,54 @@ def TT_CustomOp : TT_Op<"custom", [Pure, MemoryEffects<[MemWrite]>]> {
1515	let assemblyFormat = "$op_name attr-dict ($args^ `:` type($args))?";	1515	let assemblyFormat = "$op_name attr-dict ($args^ `:` type($args))?";
1516	}	1516	}
1517		1517
		1518	+//
		1519	+// Built-in: IndirectLoad Op
		1520	+//
		1521	+def TT_IndirectLoadOp : TT_Op<"indirect_load", [
		1522	+ DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
		1523	+ AttrSizedOperandSegments
		1524	+]> {
		1525	+ let summary = "Built-in: indirect load from global memory using per-element offsets with optional mask/other";
		1526	+
		1527	+ let description = [{
		1528	+ Built-in operation emitted by the compiler for unstructured (discrete) memory
		1529	+ accesses.These are not written directly in the user IR.
		1530	+
		1531	+ Load values from global memory based on per-element offsets. If `mask`
		1532	+ is provided, false lanes return `other`.
		1533	+
		1534	+ The operation takes:
		1535	+ - src: Source pointer
		1536	+ - offsets: Tensor of per-element offsets (relative to `src`) for accessing source memory
		1537	+ - mask (optional): if mask[idx] is false, do not load the data at address pointer[idx]
		1538	+ - other (optional): if mask[idx] is false, return other[idx]
		1539	+ }];
		1540	+
		1541	+ let arguments = (
		1542	+ ins TT_Ptr:$src,
		1543	+ TT_IntTensor:$offsets,
		1544	+ Optional<TT_BoolLike>:$mask,
		1545	+ Optional<TT_Type>:$other
		1546	+ );
		1547	+
		1548	+ let results = (outs TT_Tensor:$result);
		1549	+
		1550	+ let assemblyFormat = [{
		1551	+ $src `:` type($src) `,`
		1552	+ $offsets `:` type($offsets)
		1553	+ (`,` $mask^ `:` type($mask))?
		1554	+ (`,` $other^ `:` type($other))?
		1555	+ attr-dict `->` type($result)
		1556	+ }];
		1557	+
		1558	+ let builders = [
		1559	+ OpBuilder<(ins
		1560	+ "Value":$src,
		1561	+ "Value":$offsets,
		1562	+ "Value":$mask,
		1563	+ "Value":$blocksize
		1564	+ )>
		1565	+ ];
		1566	+}
		1567	+
1518	#endif // Triton_OPS	1568	#endif // Triton_OPS

Mtriton_patch/lib/Dialect/Triton/IR/Ops.cpp+8-0

	@@ -55,6 +55,14 @@ void EmbeddingGatherOp::getEffects(
55	triton::GlobalMemory::get());	55	triton::GlobalMemory::get());
56	}	56	}
57		57
		58	+
		59	+void IndirectLoadOp::getEffects(
		60	+ SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
		61	+ &effects) {
		62	+ effects.emplace_back(MemoryEffects::Read::get(), &getSrcMutable(),
		63	+ triton::GlobalMemory::get());
		64	+}
		65	+
58	} // namespace triton	66	} // namespace triton
59	} // namespace mlir	67	} // namespace mlir
60		68