feat(linearize): triton-linearize integrate with master-triton-ascend-AtomGit

Mascend/backend/compiler.py+1-0

	@@ -100,6 +100,7 @@ def ttir_to_linalg(mod, metadata, opt, *, named_ops=False):
100	cmd_list = [	100	cmd_list = [
101	triton_adapter_opt_path,	101	triton_adapter_opt_path,
102	src_path,	102	src_path,
		103	+ "--triton-linearize",
103	"--discrete-mask-access-conversion",	104	"--discrete-mask-access-conversion",
104	"--triton-to-annotation",	105	"--triton-to-annotation",
105	"--triton-to-unstructure",	106	"--triton-to-unstructure",

Aascend/examples/pytest_ut/test_linearize.py+129-0

@@ -0,0 +1,129 @@
	1	+import pytest
	2	+
	3	+import triton
	4	+import triton.language as tl
	5	+import test_common
	6	+
	7	+import torch
	8	+import torch_npu
	9	+from triton.runtime.driver import driver
	10	+
	11	+
	12	+
	13	+# npu hardware params from trion
	14	+target = driver.active.get_current_target()
	15	+device = driver.active.get_current_device()
	16	+prop = driver.active.utils.get_device_properties(device)
	17	+
	18	+num_cube_core = prop["num_aicore"]
	19	+num_vector_core = prop["num_aicore"]
	20	+if (target.arch in ("Ascend910B","Ascend910_9382" )):
	21	+ num_vector_core = num_cube_core * 2
	22	+ print(target.arch, "vector_core",num_vector_core)
	23	+
	24	+def foo(a, d ,shape ):
	25	+ y = a.reshape(shape)
	26	+ y = y.permute(0,2,1) + d
	27	+ return y
	28	+
	29	+
	30	+@triton.jit
	31	+def triton_gpu_revised(in_ptr0, in_ptr1, out_ptr0, ynumel, xnumel, YBLOCK : tl.constexpr,
	32	+ SHAPE0:tl.constexpr, SHAPE1:tl.constexpr,SHAPE2:tl.constexpr,
	33	+ XBLOCK : tl.constexpr):
	34	+ yoffset = tl.program_id(1) * (tl.program_id(2) + 1) * YBLOCK
	35	+ yindex = yoffset + tl.arange(0, YBLOCK)
	36	+ ymask = yindex < ynumel
	37	+ xoffset = tl.program_id(0) * XBLOCK
	38	+ xindex = xoffset + tl.arange(0, XBLOCK)
	39	+ xmask = xindex < xnumel
	40	+ x2 = xindex[:, None]
	41	+ x2_1 = xindex[None, :]
	42	+ y3 = yindex[ :, None]
	43	+ y0 = (yindex % SHAPE1)[None,:]
	44	+ y1 = (yindex // SHAPE1)[None,:]
	45	+ tmp0 = tl.load(in_ptr0 + (x2_1 + (SHAPE2*y3)), (x2_1 < xnumel) & (y3 < ynumel) )
	46	+ tmp1 = tl.load(in_ptr1 + (y0 + (SHAPE1x2) + (SHAPE1SHAPE2*y1)), \
	47	+ (xindex[:,None] < xnumel) & (yindex[None,:] < ynumel))
	48	+ # (x2 < xnumel) & (y0 < SHAPE1))
	49	+ # (xindex[:,None] < xnumel) & (yindex[None,:] < ynumel))
	50	+ tmp10 = tmp1.permute(1,0)
	51	+ tmp2 = tmp0 + tmp10
	52	+ tl.store(out_ptr0 + (x2_1 + (SHAPE2*y3)), tmp2, (x2_1 < xnumel) & (y3 < ynumel))
	53	+
	54	+def biggest_divisor(num):
	55	+ for i in range(2,num):
	56	+ if num % i == 0:
	57	+ return num // i
	58	+ return num
	59	+
	60	+
	61	+
	62	+def find_good_yblock(ynumel,xnumel, y_upper, dtype) :
	63	+ y = ynumel
	64	+ x = xnumel
	65	+
	66	+ align_numel = 4 if dtype==torch.int64 else 8
	67	+ ub_upper = 3900 if dtype==torch.int64 else 8000
	68	+
	69	+
	70	+ # optimize block_dim
	71	+ def get_block_dim(y,x) :
	72	+ return ((xnumel + x -1 )//x) * ((y_upper+ y -1 ) // y)
	73	+
	74	+ count = 0
	75	+ while( get_block_dim(y,x) < num_vector_core and y > 8 and count < 20 ) :
	76	+ y_1 = biggest_divisor(y)
	77	+ if get_block_dim(y_1, x) > num_vector_core :
	78	+ break
	79	+ y = y_1
	80	+ if get_block_dim(y,x) < num_vector_core and x > align_numel :
	81	+ x = x // 2
	82	+ count = count + 1
	83	+
	84	+ # optimize block_size to avoid ub-overflow
	85	+ while( y * x > ub_upper) :
	86	+ y_1 = biggest_divisor(y)
	87	+ if y_1 == y or y_1 <=align_numel:
	88	+ break
	89	+ y = y_1
	90	+
	91	+ while( y * x > ub_upper and x > align_numel) :
	92	+ x_1 = x // 2
	93	+ if x_1 <= align_numel :
	94	+ break
	95	+ x = x_1
	96	+
	97	+ return (x,y)
	98	+
	99	+def triton_foo(a, d ,shape, dtype) :
	100	+ z, y, x = shape
	101	+ out = torch.empty_strided((z, x, y), (x*y, 1, x), device='npu', dtype=dtype)
	102	+ XBLOCK, YBLOCK= find_good_yblock(y,x, y*z, dtype=dtype)
	103	+ print(f"XBLOCK={XBLOCK},YBLOCK={YBLOCK}, block_dim={((x + XBLOCK -1 )//XBLOCK) * (((y*z) + YBLOCK -1 ) // YBLOCK)}")
	104	+ grid = ((x + XBLOCK -1 )//XBLOCK, ((y*z) + YBLOCK -1 ) // YBLOCK, 1)
	105	+
	106	+ triton_gpu_revised[grid](a, d, out, y*z, x,
	107	+ SHAPE0=z,SHAPE1=y,SHAPE2=x,
	108	+ YBLOCK=YBLOCK, XBLOCK = XBLOCK)
	109	+ return out
	110	+
	111	+
	112	+types = [
	113	+ (torch.float32, 'float32'),
	114	+]
	115	+
	116	+
	117	+shapes = [(8,2048,4)]
	118	+
	119	+@pytest.mark.parametrize('dtype,sigtype', types)
	120	+@pytest.mark.parametrize('Z,Y,X', shapes)
	121	+def test_linearize(Z,Y,X, dtype, sigtype) :
	122	+ shape = (Z,Y,X)
	123	+ print(f"start data validation on shape:{shape}")
	124	+ a = test_common.generate_tensor(shape=(Z, Y * X), dtype=sigtype).npu()
	125	+ d = test_common.generate_tensor(shape=(Z, X, Y), dtype=sigtype).npu()
	126	+ r = triton_foo( a, d, shape, dtype)
	127	+ r1 = foo(a, d ,shape )
	128	+ test_common.validate_cmp(sigtype, r1, r)
	129	+ print(f"data validation passed")

Mascend/triton-adapter/include/CMakeLists.txt+2-0

	@@ -5,3 +5,5 @@ add_subdirectory(TritonToLinalg)
5	add_subdirectory(TritonToLLVM)	5	add_subdirectory(TritonToLLVM)
6	add_subdirectory(DiscreteMaskAccessConversion)	6	add_subdirectory(DiscreteMaskAccessConversion)
7	add_subdirectory(TritonToUnstructure)	7	add_subdirectory(TritonToUnstructure)
		8	+add_subdirectory(Dialect)
		9	+add_subdirectory(TritonLinearize)

Aascend/triton-adapter/include/Dialect/CMakeLists.txt+1-0

	@@ -0,0 +1 @@
		1	+add_subdirectory(TritonStructured)

Aascend/triton-adapter/include/Dialect/TritonStructured/CMakeLists.txt+1-0

	@@ -0,0 +1 @@
		1	+add_subdirectory(IR)

Aascend/triton-adapter/include/Dialect/TritonStructured/IR/CMakeLists.txt+8-0

@@ -0,0 +1,8 @@
	1	+set(LLVM_TARGET_DEFINITIONS TritonStructuredDialect.td)
	2	+mlir_tablegen(TritonStructuredDialect.h.inc -gen-dialect-decls -dialect=tts)
	3	+mlir_tablegen(TritonStructuredDialect.cpp.inc -gen-dialect-defs -dialect=tts)
	4	+mlir_tablegen(TritonStructuredOps.h.inc -gen-op-decls)
	5	+mlir_tablegen(TritonStructuredOps.cpp.inc -gen-op-defs)
	6	+
	7	+
	8	+add_public_tablegen_target(TritonStructuredTableGen)

Aascend/triton-adapter/include/Dialect/TritonStructured/IR/TritonStructuredDialect.h+30-0

@@ -0,0 +1,30 @@
	1	+#ifndef MLIR_DIALECT_TRITON_STRUCTURED_IR_TRITON_STRUCTURED_DIALECT_H_
	2	+#define MLIR_DIALECT_TRITON_STRUCTURED_IR_TRITON_STRUCTURED_DIALECT_H_
	3	+
	4	+#include "mlir/IR/BuiltinOps.h"
	5	+#include "mlir/IR/BuiltinTypes.h"
	6	+#include "mlir/IR/Dialect.h"
	7	+#include "mlir/IR/MLIRContext.h"
	8	+#include "mlir/IR/OpDefinition.h"
	9	+#include "mlir/IR/SymbolTable.h"
	10	+#include "mlir/IR/TypeSupport.h"
	11	+#include "mlir/IR/Types.h"
	12	+#include "mlir/Interfaces/SideEffectInterfaces.h"
	13	+#include "triton/Dialect/Triton/IR/Dialect.h"
	14	+
	15	+#include "mlir/IR/Dialect.h"
	16	+
	17	+using namespace mlir;
	18	+using namespace mlir::triton;
	19	+//===----------------------------------------------------------------------===//
	20	+// TritonStructured Operations
	21	+//===----------------------------------------------------------------------===//
	22	+#include "Dialect/TritonStructured/IR/TritonStructuredDialect.h.inc"
	23	+
	24	+// Include the auto-generated header file containing the declarations of the
	25	+// TritonStructured operations.
	26	+#define GET_OP_CLASSES
	27	+
	28	+#include "Dialect/TritonStructured/IR/TritonStructuredOps.h.inc"
	29	+
	30	+#endif

Aascend/triton-adapter/include/Dialect/TritonStructured/IR/TritonStructuredDialect.td+66-0

@@ -0,0 +1,66 @@
	1	+#ifndef TRITON_STRUCTURED_DIALECT
	2	+#define TRITON_STRUCTURED_DIALECT
	3	+
	4	+include "mlir/IR/OpBase.td"
	5	+include "triton/Dialect/Triton/IR/TritonDialect.td"
	6	+include "triton/Dialect/Triton/IR/TritonTypes.td"
	7	+include "triton/Dialect/Triton/IR/TritonAttrDefs.td"
	8	+//include "triton/Dialect/Triton/IR/TritonTypes.td"
	9	+include "mlir/Interfaces/SideEffectInterfaces.td"
	10	+
	11	+
	12	+
	13	+def Triton_Structured_Dialect : Dialect {
	14	+ let name = "tts";
	15	+
	16	+ let cppNamespace = "::mlir::tts";
	17	+
	18	+ let summary = "Structured Triton operations";
	19	+
	20	+ let description = [{
	21	+ Triton Structured Dialect.
	22	+ }];
	23	+
	24	+ let dependentDialects = [
	25	+ "triton::TritonDialect"
	26	+ ];
	27	+
	28	+ let usePropertiesForAttributes = 1;
	29	+}
	30	+
	31	+//
	32	+// Op Base
	33	+//
	34	+class TTS_Op<string mnemonic, list<Trait> traits = []> :
	35	+ Op<Triton_Structured_Dialect, mnemonic, traits> {
	36	+}
	37	+
	38	+
	39	+// SameVariadicResultSize
	40	+// AttrSizedResultSegments
	41	+def TTS_GetStructuredStateOp : TTS_Op<"get_structured_state", [AttrSizedResultSegments, Pure]> {
	42	+ let summary = "Placeholder for the structured pointer states computed during PtrAnalysis.";
	43	+ let description = "Used to pass the offsets and strides to scf.for op to simplify IR rewrites.";
	44	+
	45	+ let arguments = (ins AnyTypeOf<[TT_PtrLike, I32Tensor, I64Tensor,I16Tensor,I8Tensor,I1Tensor]>:$input);
	46	+ let results = (outs AnyTypeOf<[TT_PtrLike, I32Tensor, I64Tensor,I16Tensor,I8Tensor,I1Tensor]>:$structured, Variadic<Index>:$offsets, Variadic<Index>:$strides);
	47	+
	48	+ let builders = [
	49	+ OpBuilder<(ins "Value":$input)>,
	50	+ ];
	51	+
	52	+ let extraClassDeclaration = [{
	53	+ static std::optional<std::pair<SmallVector<Type>, SmallVector<Type>>>
	54	+ getOffsetAndStrideTypes(MLIRContext *context, Type ptrLikeType);
	55	+
	56	+ static std::optional<std::pair<int32_t, int32_t>>
	57	+ getOffsetAndStrideSegmentSizes(Type ptrLikeType);
	58	+ }];
	59	+
	60	+ let hasFolder = 0;
	61	+ let hasVerifier = 1;
	62	+}
	63	+
	64	+
	65	+
	66	+#endif // TRITON_STRUCTURED_DIALECT

Aascend/triton-adapter/include/TritonLinearize/CMakeLists.txt+3-0

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

Aascend/triton-adapter/include/TritonLinearize/MaskAnalysis.h+231-0

@@ -0,0 +1,231 @@
	1	+/*
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	3	+ *
	4	+ * Permission is hereby granted, free of charge, to any person obtaining a copy
	5	+ * of this software and associated documentation files (the "Software"), to deal
	6	+ * in the Software without restriction, including without limitation the rights
	7	+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	8	+ * copies of the Software, and to permit persons to whom the Software is
	9	+ * furnished to do so, subject to the following conditions:
	10	+ *
	11	+ * The above copyright notice and this permission notice shall be included in
	12	+ * all copies or substantial portions of the Software.
	13	+ *
	14	+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	15	+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	16	+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	17	+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	18	+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	19	+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
	20	+ * THE SOFTWARE.
	21	+ */
	22	+
	23	+#ifndef TRITON_LINERIZE_ANALYSIS_MASKANALYSIS_H
	24	+#define TRITON_LINERIZE_ANALYSIS_MASKANALYSIS_H
	25	+
	26	+#include "mlir/Dialect/Arith/IR/Arith.h"
	27	+#include "mlir/Dialect/MemRef/IR/MemRef.h"
	28	+#include "mlir/Dialect/Tensor/IR/Tensor.h"
	29	+#include "TritonLinearize/OpFoldResultUtils.h"
	30	+
	31	+#include "mlir/IR/OpDefinition.h"
	32	+#include "triton/Dialect/Triton/IR/Dialect.h"
	33	+#include "llvm/ADT/SmallVector.h"
	34	+
	35	+#include "triton/Dialect/Triton/IR/Dialect.h"
	36	+
	37	+#include <utility>
	38	+
	39	+
	40	+namespace mlir {
	41	+
	42	+class OpBuilder;
	43	+
	44	+namespace triton_linearize {
	45	+
	46	+struct dimInfo {
	47	+ OpFoldResult div;
	48	+ OpFoldResult shape;
	49	+ bool isSlt;
	50	+ Value rhs; // rhs value of CmpIOp
	51	+ bool isRealDim;
	52	+ int64_t dim;
	53	+ dimInfo() : isSlt(false), isRealDim(false), dim(0) {}
	54	+ dimInfo(OpFoldResult div, OpFoldResult shape, int64_t dim = 0) : div(div), shape(shape), isSlt(false), isRealDim(true), dim(dim) {}
	55	+ bool operator==(const dimInfo &other) const {
	56	+ auto staticDiv = getIntAttr(div);
	57	+ auto staticShape = getIntAttr(shape);
	58	+ auto otherShape = getIntAttr(other.shape);
	59	+ auto otherDiv = getIntAttr(other.div);
	60	+ assert(staticDiv.has_value() && staticDiv.has_value() && otherDiv.has_value() && otherShape.has_value() && "MaskAnalysis: do not support dynamic shape/div");
	61	+ return staticDiv.value() == otherDiv.value() && staticShape.value() == otherShape.value() && dim == other.dim;
	62	+ }
	63	+ void dump() const ;
	64	+ bool hasModulo() const {
	65	+ auto intAttr = getIntAttr(shape);
	66	+ if (!intAttr.has_value()) {
	67	+ return false;
	68	+ }
	69	+ return intAttr.value() != 0;
	70	+ };
	71	+ bool hasDivision() const {
	72	+ auto intAttr = getIntAttr(div);
	73	+ if (!intAttr.has_value()) {
	74	+ return false;
	75	+ }
	76	+ return intAttr.value() != 0;
	77	+ };
	78	+};
	79	+// Data structure used to decode the pattern in a mask used for load and store.
	80	+// start and end field represent the start and end index of a range (produced
	81	+// by make_range, addi, etc.). While multi-dimensional data is possible, we
	82	+// assume range comparison can only be done on 1 dimension at a time (and
	83	+// results of range comparions across dimensions can be combined), hence start
	84	+// and end are not vectors. dims represents the real access size for ld/st
	85	+// (instead of the tensor/memref size specified by the IR). scalar is a shortcut
	86	+// used when the entire state contains a single scalar value.
	87	+//
	88	+// The general lifetime of this data structure is roughly:
	89	+// 1. A range is created by make_range and optionally operated on by addi w/
	90	+// result of splat, expand_dims, etc. During this phase, either (1) both start
	91	+// and end are populated, or (2) scalar is populated. Only one of the dimensions
	92	+// (that contains the range) can have dim > 1.
	93	+// 2. Result from step 1 is compared with a another MaskState that represents a
	94	+// scalar value. The resulting state only has dims populated.
	95	+// 3. Optionally, result from step 2 can be broadcasted and anded with other
	96	+// results from step 2. The resulting state only has dims populated.
	97	+//
	98	+// Example of creating 2D mask:
	99	+// mask = (rows[:, None] < M) & (cols[None, :] < N)
	100	+struct MaskState {
	101	+ OpFoldResult start;
	102	+ OpFoldResult end;
	103	+ SmallVector<OpFoldResult> dims;
	104	+ OpFoldResult scalar;
	105	+ // OpFoldResult div;
	106	+ // OpFoldResult shape;
	107	+ SmallVector<dimInfo> stateInfo;
	108	+
	109	+ void dump() const;
	110	+
	111	+ bool hasModulo() const;
	112	+
	113	+ bool hasDivision() const;
	114	+
	115	+ bool dimHasDivision(uint32_t dim) const;
	116	+
	117	+ bool dimHasModulo(uint32_t dim) const;
	118	+
	119	+ int64_t getRank() const { return dims.size(); }
	120	+
	121	+ bool isEmpty() const { return getRank() == 0 && !scalar && !start && !end; }
	122	+
	123	+ bool isMask() const { return !start && !end && !scalar && dims.size() != 0; }
	124	+
	125	+ // Recursively parse a Value; call the coresponding function based on the
	126	+ // defining operation and Value type
	127	+ LogicalResult parse(Value operand, const Location loc, OpBuilder &builder);
	128	+
	129	+ tensor::ExtractSliceOp
	130	+ getExtractSlice(Value source, const Location loc,
	131	+ OpBuilder &builder) const;
	132	+
	133	+ memref::SubViewOp getSubview(Value source, const Location loc,
	134	+ OpBuilder &builder) const;
	135	+
	136	+ std::pair<memref::SubViewOp, memref::SubViewOp>
	137	+ getSideBySideSubviews(Value block1, Value block2, const Location loc,
	138	+ OpBuilder &builder) const;
	139	+
	140	+ std::pair<memref::SubViewOp, memref::SubViewOp>
	141	+ getStackedSubviews(Value block1, Value block2, const Location loc,
	142	+ OpBuilder &builder) const;
	143	+
	144	+ tensor::InsertSliceOp getInsertSlice(Value source, Value dest,
	145	+ const Location &loc,
	146	+ OpBuilder &builder) const;
	147	+
	148	+
	149	+ void eraseInsertedOps(Operation *rawOp, PatternRewriter &rewriter);
	150	+
	151	+
	152	+private:
	153	+ // -------
	154	+ // Utility functions to operate on MaskState
	155	+ // -------
	156	+ LogicalResult addStateScalar(const MaskState &state,
	157	+ const OpFoldResult scalar, Location loc,
	158	+ OpBuilder &builder);
	159	+
	160	+ LogicalResult addStates(const MaskState &lhsState, const MaskState &rhsState,
	161	+ Location loc, OpBuilder &builder);
	162	+
	163	+ LogicalResult minStates(MaskState &lhsState, MaskState &rhsState,
	164	+ Location loc, OpBuilder &builder);
	165	+ // -------
	166	+ // Helper functions to parse values to populate MaskState
	167	+ // -------
	168	+
	169	+ // Operand is the result of a constant
	170	+ // Get the value of the constant and assign it to scalar.
	171	+ LogicalResult parseConstant(arith::ConstantOp constOp, const Location loc,
	172	+ OpBuilder &builder);
	173	+
	174	+ // Operand is an integer scalar
	175	+ LogicalResult parseIntScalar(Value scalar, const Location loc,
	176	+ OpBuilder &builder);
	177	+
	178	+ // Operand is the result of addi
	179	+ // One and only one of the operands should be a scalar. Increment both start
	180	+ // and end, dims remains unchanged, and scalar is empty.
	181	+ LogicalResult parseAdd(arith::AddIOp addOp, const Location loc,
	182	+ OpBuilder &builder);
	183	+ // Operand is the result of andi
	184	+ // Each of the result state dims is smaller of the two operands' dims.
	185	+ // Insert instruction if needed to get new dims.
	186	+ LogicalResult parseAnd(arith::AndIOp andOp, const Location loc,
	187	+ OpBuilder &builder);
	188	+
	189	+ // Operand is the result of cmpi
	190	+ // Assume only of the dimensions have size > 1. Only support slt for now.
	191	+ // For that dimension, calculate this new dim as: dim = min(end, value) -
	192	+ // start
	193	+ LogicalResult parseCmp(arith::CmpIOp cmpOp, const Location loc,
	194	+ OpBuilder &builder);
	195	+ // Operand is the result of make_range
	196	+ // Set start and end accordingly; step size must be 1.
	197	+ LogicalResult parseMakeRange(triton::MakeRangeOp rangeOp, const Location loc,
	198	+ OpBuilder &builder);
	199	+ // Operand is the result of broadcast
	200	+ // Change dims only; assume only applies to tensors.
	201	+ LogicalResult parseBroadcast(triton::BroadcastOp broadcastOp,
	202	+ const Location loc, OpBuilder &builder);
	203	+ // Operand is the result of splat
	204	+ // Assume only applies to scalar. start and end are left empty; scalar will
	205	+ // be assigned, and dims will be updated.
	206	+ LogicalResult parseSplat(triton::SplatOp splatOp, const Location loc,
	207	+ OpBuilder &builder);
	208	+ // Operand is the result of expand_dims
	209	+ // Insert additional dims; start and end do not change and correspond to the
	210	+ // dimension that contains the range.
	211	+ LogicalResult parseExpandDims(triton::ExpandDimsOp expandDimsOp,
	212	+ const Location loc, OpBuilder &builder);
	213	+
	214	+ LogicalResult parseExtSI(arith::ExtSIOp op, const Location loc,
	215	+ OpBuilder &builder);
	216	+ LogicalResult parseRemsi(arith::RemSIOp remsiOp,
	217	+ const Location loc,
	218	+ OpBuilder &builder);
	219	+
	220	+ LogicalResult parseDivsi(arith::DivSIOp divsiOp,
	221	+ const Location loc,
	222	+ OpBuilder &builder) ;
	223	+ LogicalResult parseLoopIterArg(Value v, const Location loc,
	224	+ OpBuilder &builder);
	225	+};
	226	+
	227	+} // namespace triton
	228	+
	229	+} // namespace mlir
	230	+
	231	+#endif

Aascend/triton-adapter/include/TritonLinearize/OpFoldResultUtils.h+92-0

@@ -0,0 +1,92 @@
	1	+/*
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	3	+ *
	4	+ * Permission is hereby granted, free of charge, to any person obtaining a copy
	5	+ * of this software and associated documentation files (the "Software"), to deal
	6	+ * in the Software without restriction, including without limitation the rights
	7	+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	8	+ * copies of the Software, and to permit persons to whom the Software is
	9	+ * furnished to do so, subject to the following conditions:
	10	+ *
	11	+ * The above copyright notice and this permission notice shall be included in
	12	+ * all copies or substantial portions of the Software.
	13	+ *
	14	+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	15	+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	16	+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	17	+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	18	+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	19	+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
	20	+ * THE SOFTWARE.
	21	+ */
	22	+
	23	+#ifndef TRITON_LINERIZE_ANALYSIS_OPFOLDRESULT_UTILS_H
	24	+#define TRITON_LINERIZE_ANALYSIS_OPFOLDRESULT_UTILS_H
	25	+
	26	+#include "mlir/IR/Location.h"
	27	+#include "mlir/IR/OpDefinition.h"
	28	+
	29	+#include <optional>
	30	+
	31	+namespace mlir {
	32	+
	33	+class OpBuilder;
	34	+
	35	+Value materializeValue(OpBuilder &builder, Location loc, OpFoldResult ofr);
	36	+
	37	+// Return integer if ofr is an IntegerAttr. Note that this function differs
	38	+// from getConstantIntValue, which returns an integer if ofr is the constant
	39	+// result of an operation too.
	40	+std::optional<int64_t> getIntAttr(const OpFoldResult ofr);
	41	+
	42	+// Return if ofr contains a constant zero, either represented by an integer
	43	+// attribute or a constant value.
	44	+bool hasConstZero(const OpFoldResult ofr);
	45	+
	46	+// Create a value of index type if necessary from an OpFoldResult.
	47	+Value ofrToIndexValue(const OpFoldResult ofr, const Location loc, OpBuilder &b);
	48	+
	49	+// Create a vector of values of index type if necessary from an array of
	50	+// OpFoldResults.
	51	+SmallVector<Value> ofrsToIndexValues(ArrayRef<OpFoldResult> ofrs,
	52	+ const Location loc, OpBuilder &b);
	53	+
	54	+// Process addition of two OFRs. If both OFRs are Integer Attributes, result
	55	+// is an Integer Attribute. Otherwise, insert the arith.addi instruction if
	56	+// needed and use its result Value.
	57	+OpFoldResult addOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	58	+ const Location loc, OpBuilder &b);
	59	+
	60	+// Produce result = lhs - rhs. If both OFRs are Integer Attributes, result
	61	+// is an Integer Attribute. Otherwise, insert the arith.addi instruction if
	62	+// needed and use its result Value.
	63	+OpFoldResult subOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	64	+ const Location loc, OpBuilder &b);
	65	+
	66	+// Process multiplication of two OFRs. If both OFRs are Integer Attributes,
	67	+// result is an Integer Attribtue. Otherwise, insert the arith.muli
	68	+// instruction if needed and use its result Value.
	69	+OpFoldResult mulOFRValue(const OpFoldResult lhs, const Value rhs,
	70	+ const Location loc, OpBuilder &b);
	71	+
	72	+// Process division of two OFRs. If both OFRs are Integer Attributes,
	73	+// the result is an Integer Attribute. Otherwise, insert the arith.divi
	74	+// instruction if needed and use its result Value.
	75	+OpFoldResult divOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	76	+ const Location loc, OpBuilder &b);
	77	+
	78	+OpFoldResult remOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	79	+ const Location loc, OpBuilder &b);
	80	+
	81	+
	82	+OpFoldResult mulOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	83	+ const Location loc, OpBuilder &b);
	84	+
	85	+OpFoldResult minOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	86	+ const Location loc, OpBuilder &b);
	87	+
	88	+OpFoldResult maxOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	89	+ const Location loc, OpBuilder &b);
	90	+} // namespace mlir
	91	+
	92	+#endif

Aascend/triton-adapter/include/TritonLinearize/Passes.h+15-0

@@ -0,0 +1,15 @@
	1	+#ifndef TRITON_ADAPTER_TRITON_LINEARIZE_PASSES_H
	2	+#define TRITON_ADAPTER_TRITON_LINEARIZE_PASSES_H
	3	+
	4	+#include "TritonLinearize.h"
	5	+
	6	+namespace mlir {
	7	+namespace triton {
	8	+
	9	+#define GEN_PASS_REGISTRATION
	10	+#include "ascend/triton-adapter/include/TritonLinearize/Passes.h.inc"
	11	+
	12	+} // namespace triton
	13	+} // namespace mlir
	14	+
	15	+#endif // TRITON_ADAPTER_TRITON_LINEARIZE_PASSES_H

Aascend/triton-adapter/include/TritonLinearize/Passes.td+17-0

@@ -0,0 +1,17 @@
	1	+#ifndef TRITON_LINEARIZE_PASSES
	2	+#define TRITON_LINEARIZE_PASSES
	3	+
	4	+include "mlir/Pass/PassBase.td"
	5	+
	6	+def TritonLinearize : Pass<"triton-linearize", "mlir::ModuleOp"> {
	7	+ let summary = "remove reminder/divider and reproduce addptr/mask expression ";
	8	+ let constructor = "triton::createTritonLinearizePass()";
	9	+ let options = [
	10	+ Option<"reserved", "reserved",
	11	+ "bool", /default/"true",
	12	+ "reserved option">,
	13	+
	14	+ ];
	15	+}
	16	+
	17	+#endif // TRITON_LINEARIZE_PASSES

Aascend/triton-adapter/include/TritonLinearize/PtrAnalysis.h+477-0

@@ -0,0 +1,477 @@
	1	+/*
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	3	+ *
	4	+ * Permission is hereby granted, free of charge, to any person obtaining a copy
	5	+ * of this software and associated documentation files (the "Software"), to deal
	6	+ * in the Software without restriction, including without limitation the rights
	7	+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	8	+ * copies of the Software, and to permit persons to whom the Software is
	9	+ * furnished to do so, subject to the following conditions:
	10	+ *
	11	+ * The above copyright notice and this permission notice shall be included in
	12	+ * all copies or substantial portions of the Software.
	13	+ *
	14	+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	15	+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	16	+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	17	+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	18	+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	19	+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
	20	+ * THE SOFTWARE.
	21	+ */
	22	+
	23	+#ifndef TRITON_LINERIZE_ANALYSISSTRUCTURED_PTRANALYSIS_H
	24	+#define TRITON_LINERIZE_ANALYSISSTRUCTURED_PTRANALYSIS_H
	25	+
	26	+#include "mlir/Dialect/Arith/IR/Arith.h"
	27	+#include "mlir/Dialect/MemRef/IR/MemRef.h"
	28	+#include "mlir/Dialect/SCF/IR/SCF.h"
	29	+
	30	+#include "mlir/IR/Value.h"
	31	+#include "mlir/Support/LLVM.h"
	32	+#include "mlir/Support/LogicalResult.h"
	33	+
	34	+#include "triton/Dialect/Triton/IR/Dialect.h"
	35	+#include "Dialect/TritonStructured/IR/TritonStructuredDialect.h"
	36	+
	37	+#include <cstddef>
	38	+#include <set>
	39	+
	40	+namespace mlir {
	41	+
	42	+class OpBuilder;
	43	+
	44	+namespace triton {
	45	+
	46	+const extern std::string ptrAnalysisAttr;
	47	+
	48	+// Data structure used to decode pointer arithmetics. offsets, sizes, and
	49	+// strides are in unit of elements in a linearly laid-out memory, which is the
	50	+// same as pointer arithmetic operations in Triton language. scalar is a
	51	+// shortcut used when the entire state describes a single scalar value. source
	52	+// is the base pointer. If order is present, PtrState describes block pointer;
	53	+// otherwise it describes non-block pointers. When it describes block pointer,
	54	+// shape field means the same field as tt.make_tensor_ptr; when it describes a
	55	+// non-block pointer, shape field indicates how address wraps around (i.e.,
	56	+// modulo); a constant 0 indicates no modulo for the dimension.
	57	+
	58	+struct StateInfo{
	59	+ OpFoldResult offset; // final offset with base address
	60	+ OpFoldResult stride;
	61	+ OpFoldResult shape; // rem value
	62	+ OpFoldResult mask; // div value
	63	+ Value dimVar; // new dim var (for div/rem)
	64	+ OpFoldResult dimOffset; // offset for new dimVar
	65	+ size_t dim;
	66	+ bool remIsBeforeDiv = false;
	67	+
	68	+ StateInfo() : dim(0) {}
	69	+ StateInfo(OpFoldResult offset, OpFoldResult stride,
	70	+ OpFoldResult shape, OpFoldResult mask, size_t dim = 0)
	71	+ : offset(offset), stride(stride), shape(shape), mask(mask), dim(dim) {}
	72	+ void dump() const;
	73	+
	74	+ bool hasModulo() const {
	75	+ auto intAttr = getIntAttr(shape);
	76	+ if (!intAttr.has_value()) {
	77	+ return false;
	78	+ }
	79	+ return intAttr.value() != 0;
	80	+ };
	81	+ bool hasDivision() const {
	82	+ auto intAttr = getIntAttr(mask);
	83	+ if (!intAttr.has_value()) {
	84	+ return false;
	85	+ }
	86	+ return intAttr.value() != 0;
	87	+ };
	88	+
	89	+};
	90	+
	91	+// Group StateInfo by original dimension
	92	+struct StateInfoGroup {
	93	+ size_t dim; // original dim
	94	+ SmallVector<size_t> idxes; // indexes of stateInfo belonging to this dim (in original order)
	95	+ int64_t minStride = std::numeric_limits<int64_t>::max(); // The smallest stride within this group
	96	+
	97	+ void dump() const;
	98	+};
	99	+
	100	+enum class MemAccVal { Undefined = 0, StrucMemAcc = 1, UnstrucMemAcc = 2, Fallback = 3 };
	101	+
	102	+struct MemAccType {
	103	+
	104	+ MemAccVal value;
	105	+
	106	+ explicit constexpr MemAccType(MemAccVal v = MemAccVal::Undefined)
	107	+ : value(v) {}
	108	+
	109	+ constexpr operator MemAccVal() const { return value; }
	110	+ explicit operator bool() = delete;
	111	+
	112	+ constexpr bool isUndefined() const { return value == MemAccVal::Undefined; }
	113	+ constexpr bool isStructured() const {
	114	+ return value == MemAccVal::StrucMemAcc;
	115	+ }
	116	+ constexpr bool isUnstructured() const {
	117	+ return value == MemAccVal::UnstrucMemAcc;
	118	+ }
	119	+ constexpr bool isFallback() const {
	120	+ return value == MemAccVal::Fallback;
	121	+ }
	122	+
	123	+ void merge(MemAccType &other) {
	124	+ this->value = (this->value > other.value) ? this->value : other.value;
	125	+ }
	126	+
	127	+ std::string_view toString() const {
	128	+ static constexpr std::string_view names[] = {"Undefined", "StrucMemAcc",
	129	+ "UnstrucMemAcc"};
	130	+ return names[static_cast<int>(value)];
	131	+ }
	132	+};
	133	+
	134	+
	135	+/*
	136	+limits:
	137	+
	138	+sizes.size() can greater than stateInfo.size()
	139	+dimLenth.size() = sizes.size() + 1, which initialized in countDims() called by addPtrState(), only triggered when meeting addptrOp.
	140	+when sizes.size()==1 or stateInfo.size()==0 : (scalar \|\| source) must be True.
	141	+*/
	142	+struct PtrState {
	143	+
	144	+ SmallVector<StateInfo> stateInfo; // shape info when load , maintained with visitOps
	145	+ SmallVector<OpFoldResult> sizes; // original shape, maintained with visitOps
	146	+ Attribute blockedAttr; // mark if isa tt.make_block_ptr, maintained TODO: maybe unused
	147	+
	148	+
	149	+ SmallVector<int32_t> order; // maintained in makeTensorPtr and advance
	150	+ SmallVector<size_t> dimLenth; // maintained only in countDim()
	151	+ SmallVector<int32_t> permute; // used if hasPermute()
	152	+ SmallVector<StateInfoGroup> stateInfoGroups; // stateInfo grouped by stride descendingly, maintained in analyzePermute, sortStateByStride
	153	+ bool ptrIsTensor = true; // maintained in addPtrState and Splat // TODO refactor ptrIsTensor
	154	+ MemAccType memAccTy;
	155	+
	156	+ Value source; // base address (ptr), maintained with visitOps
	157	+ Value scalar; // scalar offset (int), maintained with visitOps
	158	+
	159	+ void dump() const ;
	160	+ int32_t getRank() const;
	161	+ bool isLegal() const;
	162	+
	163	+ bool isSameSizeAs(const PtrState& x) const;
	164	+
	165	+ bool shouldRemove(const StateInfo& x) const;
	166	+
	167	+ bool opFoldResultIsZero(OpFoldResult op) const;
	168	+
	169	+ bool isEmpty() const;
	170	+
	171	+ bool hasModulo() const;
	172	+
	173	+ bool hasDivision() const;
	174	+
	175	+ bool hasPermute() const;
	176	+
	177	+ bool hasBroadcast() const;
	178	+
	179	+ bool dimHasDivision(uint32_t dim) const;
	180	+
	181	+ bool dimHasModulo(uint32_t dim) const;
	182	+
	183	+ bool isBlockPtr() const;
	184	+
	185	+ bool countDims();
	186	+
	187	+ void setMemAccTy(const MemAccType &);
	188	+ void setMemAccVal(const MemAccVal);
	189	+ MemAccType getMemAccType() const ;
	190	+ MemAccType &getMemAccTypeRef() ;
	191	+ void removeSource() { this->source = nullptr; };
	192	+ bool hasSource() const { return this->source != nullptr; }
	193	+
	194	+ // This function transforms constant dimensions in a matrix into offsets,
	195	+ // effectively removing those dimensions. It simplifies the matrix by
	196	+ // eliminating unnecessary constant dimensions.
	197	+ LogicalResult removeConstDim(SmallVector<StateInfo> &InfoPerDim,
	198	+ Operation *op, OpBuilder &builder);
	199	+
	200	+ // Check whether the current read count is less than the batch size and
	201	+ // broadcast the data along the highest dimension if needed.
	202	+ LogicalResult broadcastIfNeeded(SmallVector<StateInfo> &InfoPerDim,
	203	+ Operation *op, OpBuilder &builder);
	204	+
	205	+ // Process addition of two PtrStates.
	206	+ LogicalResult addState(const PtrState &lhsState, const PtrState &rhsState,
	207	+ Operation *op, OpBuilder &builder);
	208	+
	209	+ // Process multiplication of two PtrStates
	210	+ LogicalResult mulState(const PtrState &lhsState, const PtrState &rhsState,
	211	+ Operation *op, OpBuilder &builder);
	212	+
	213	+ LogicalResult subState(const PtrState &lhsState,
	214	+ const PtrState &rhsState, Operation *op,
	215	+ OpBuilder &builder);
	216	+
	217	+ // Process addition of ptr and offset.
	218	+ LogicalResult addPtrState(const PtrState &lhsState, const PtrState &rhsState,
	219	+ Operation *op, OpBuilder &builder);
	220	+
	221	+ LogicalResult ExpandInfo(SmallVector<StateInfo> &InfoPerDim,
	222	+ Operation *op, OpBuilder &builder);
	223	+
	224	+ // analyze Permute info.
	225	+ LogicalResult analyzePermute(SmallVector<StateInfo> &InfoPerDim,
	226	+ Operation *op, OpBuilder &builder);
	227	+
	228	+ // sort stateInfo by stride descendingly.
	229	+ LogicalResult sortStateByStride(SmallVector<StateInfo> &InfoPerDim,
	230	+ Operation *op, OpBuilder &builder);
	231	+
	232	+
	233	+ triton::AddPtrOp createAddPtrOp(OpBuilder &builder, Location loc);
	234	+};
	235	+
	236	+class PtrAnalysis {
	237	+ // This function is internally used by getLoopIterArgPtrState and
	238	+ // getLoopResultPtrState to get the correct PtrState for either an iter-arg or
	239	+ // a loop's result.
	240	+ //
	241	+ // A PtrState of an scf.for's iter-arg is the same as its corresponding
	242	+ // init-arg, except that the strides and offsets have to point to the loop's
	243	+ // iter-args that were created to carry the offsets and strides.
	244	+ //
	245	+ // For instance, for a pointer with index i and rank 2, 4 additional args
	246	+ // starting at index i + 1 are created. The PtrState's strides and offsets
	247	+ // value of the pointer's iter-arg must point to these 4 additionally created
	248	+ // iter-args.
	249	+ //
	250	+ // A similar process is used for getting the PtrState of the loop's i'th
	251	+ // result: its strides and offsets have to point to the corresponding stride
	252	+ // and offset values returned by the loop.
	253	+ PtrState reconcileLoopPtrState(
	254	+ scf::ForOp forOp, size_t ptrArgIndex, const PtrState &state,
	255	+ llvm::function_ref<Value(scf::ForOp op, size_t)> getReplacementVal);
	256	+
	257	+ PtrState reconcileWhilePtrState(
	258	+ scf::WhileOp whileOp, size_t argIndex, const PtrState &state,
	259	+ llvm::function_ref<Value(scf::WhileOp op, size_t)> getReplacementVal);
	260	+
	261	+ DenseSet<Value> maybeStructuredArgs;
	262	+
	263	+public:
	264	+ void initializeMaybeStructuredArgs(Operation *op);
	265	+ using IndexMapSet = std::map<int, std::set<int>>;
	266	+ IndexMapSet levelToBlockArgIndex;
	267	+ int level = 0;
	268	+ // AddptrOp result -> PtrState
	269	+ llvm::SmallDenseMap<Value, PtrState> knownPtrs;
	270	+ // AddptrOp result -> new AddPtrOp result
	271	+ IRMapping ptrMap;
	272	+
	273	+ // Recursively parse a Value; call the corresponding
	274	+ // function based on the defining operation and argument type.
	275	+ LogicalResult visitOperand(Value operand, PtrState &state, const Location loc, OpBuilder &builder);
	276	+
	277	+ // Operand is a result of an scf.for. Such cases occur when there are multiple
	278	+ // levels of nested loops where the results of the inner scf.for (pointer) are
	279	+ // yielded by the outer loop.
	280	+ LogicalResult visitOperandForOp(scf::ForOp forOp, Value operand,
	281	+ PtrState &state, const Location loc,
	282	+ OpBuilder &builder);
	283	+
	284	+ LogicalResult visitOperandWhileOp(scf::WhileOp whileOp, Value operand,
	285	+ PtrState &state, const Location loc,
	286	+ OpBuilder &builder);
	287	+
	288	+ // Operand is the result of arith.addi. Process both arguments and insert any
	289	+ // arith.addi instruction as needed.
	290	+ // Main assumptions:
	291	+ // Only one of lhsState and rhsState has source field set
	292	+ // Current PtrState should be empty
	293	+ // Expected result:
	294	+ // source = lhsState.source ? lhsState.source : rhsState.source
	295	+ // sizes[i] = lhsState.sizes[i] (which should match rhsState.sizes[i])
	296	+ // offsets[i] = lhsState.offsets[i] + rhsState.offsets[i]
	297	+ // strides[i] = lhsState.strides[i] + rhsState.strides[i]
	298	+ LogicalResult visitOperandAdd(arith::AddIOp addOp, PtrState &state,
	299	+ const Location loc, OpBuilder &builder);
	300	+
	301	+ // Operand is the result of arith.muli. Process both arguments and insert any
	302	+ // arith.muli instruction as needed.
	303	+ // Main assumptions:
	304	+ // Neither lhsState nor rhsState has source field set
	305	+ // Current PtrState should be empty
	306	+ // Currently only support one of the operand is a scalar index
	307	+ // Expected result (scalar and tensorState represent the two operands):
	308	+ // source = null
	309	+ // sizes[i] = tensorState.sizes[i]
	310	+ // offsets[i] = tensorState.offsets[i] * scalar
	311	+ // strides[i] = tensorState.strides[i] * scalar
	312	+ LogicalResult visitOperandMul(arith::MulIOp mulOp, PtrState &state,
	313	+ const Location loc, OpBuilder &builder);
	314	+
	315	+ LogicalResult visitOperandSub(arith::SubIOp subOp, PtrState &state,
	316	+ const Location loc,
	317	+ OpBuilder &builder);
	318	+
	319	+ LogicalResult visitOperandRem(arith::RemSIOp mulOp, PtrState &state,
	320	+ const Location loc, OpBuilder &builder);
	321	+
	322	+ LogicalResult visitOperandDiv(arith::DivSIOp divOp, PtrState &state,
	323	+ const Location loc, OpBuilder &builder);
	324	+
	325	+ // Operand is the result of make_range.
	326	+ // Main assumptions:
	327	+ // start, end, and shape are all statically known
	328	+ // The output of make_range is 1-dimensional
	329	+ // Does not check validity of inputs (e.g., stride > 0)
	330	+ // Expected result:
	331	+ // source = null
	332	+ // sizes[0] = shape[0]
	333	+ // offset[0] = start
	334	+ // strides[0] = ceiling( (end - start) / shape[0] )
	335	+ LogicalResult visitOperandMakeRange(triton::MakeRangeOp rangeOp,
	336	+ PtrState &state, Location loc,
	337	+ OpBuilder &builder);
	338	+
	339	+ // Operand is the result of expand_dims
	340	+ // Main assumptions:
	341	+ // Only 1 dimension changes for each invocation of reshape
	342	+ // The changed dimension must have size of 1
	343	+ // Expected result:
	344	+ // Insert a dimension of size 1, stride 0, and offset 0
	345	+ LogicalResult visitOperandExpandDims(triton::ExpandDimsOp expandDimsOp,
	346	+ PtrState &state, const Location loc,
	347	+ OpBuilder &builder);
	348	+
	349	+ // Operand is the result of broadcast
	350	+ // Main assumptions:
	351	+ // Rank of soure and result is the same
	352	+ // Expected result:
	353	+ // Update sizes[i] only, no changes to other fields
	354	+ LogicalResult visitOperandBroadcast(triton::BroadcastOp broadcastOp,
	355	+ PtrState &state, const Location loc,
	356	+ OpBuilder &builder);
	357	+
	358	+ // Operand is the result of splat
	359	+ // Main assumptions:
	360	+ // Source is a scalar value (i.e., an integer or a pointer, not a tensor)
	361	+ // Expected result:
	362	+ // sizes[i] reflect the shape of the result, strides[i] = 0, offsets[i] = 0
	363	+ // if source is an integer, offset[0] = scalar = source
	364	+ LogicalResult visitOperandSplat(triton::SplatOp splatOp, PtrState &state,
	365	+ const Location loc, OpBuilder &builder);
	366	+
	367	+ // Operand is the result of arith.constant that is a splat
	368	+ // Main assumptions:
	369	+ // Source is a constant op that produces a constant dense tensor where all
	370	+ // elements are the same (i.e.: a constant that is splatted)
	371	+ // Expected result:
	372	+ // sizes[i] reflect the shape of the result, strides[i] = 0, offsets[i] =
	373	+ // splat value if i == 0, otherwise 0
	374	+ LogicalResult visitOperandConstSplat(arith::ConstantOp op, PtrState &state,
	375	+const Location loc, OpBuilder &builder);
	376	+
	377	+ LogicalResult visitOperandExtSI(arith::ExtSIOp, PtrState &state,
	378	+ const Location loc, OpBuilder &builder);
	379	+
	380	+ // Operand is the result of addptr.
	381	+ // Main assumptions:
	382	+ // The ptr field should populate the source field
	383	+ // ptr and offset fields should result in same rank
	384	+ // Expected result:
	385	+ // The resulting state for ptr and offset wil be added
	386	+ LogicalResult visitOperandAddptr(triton::AddPtrOp addptrOp, PtrState &state,
	387	+ const Location loc, OpBuilder &builder);
	388	+
	389	+ // Operand is the result of tt.make_tensor_ptr.
	390	+ // Expected result:
	391	+ // Parse source pointer and grab results
	392	+ LogicalResult visitOperandMakeTensorPtr(triton::MakeTensorPtrOp makeTPtrOp,
	393	+ PtrState &state, const Location loc,
	394	+ OpBuilder &builder);
	395	+
	396	+ template <typename OpType>
	397	+ LogicalResult visitOperandIndirectLoad(OpType op, PtrState &state,
	398	+ const Location &loc,
	399	+ OpBuilder &builder) ;
	400	+
	401	+ LogicalResult visitOperandMakeTensorDescOp(triton::MakeTensorDescOp tensorDescType,
	402	+ PtrState &state, const Location loc,
	403	+ OpBuilder &builder) ;
	404	+
	405	+ LogicalResult visitOperandDescriptorLoad(triton::DescriptorLoadOp descLoadOp,
	406	+ PtrState &state, const Location loc,
	407	+ OpBuilder &builder);
	408	+
	409	+// Get the computed PtrState for the forOp's init-arg at the provided index.
	410	+ FailureOr<PtrState> getLoopInitArgPtrState(scf::ForOp forOp, size_t index);
	411	+
	412	+ // Get the computed PtrState for the forOp's iter-arg at the provided index.
	413	+ FailureOr<PtrState> getLoopIterArgPtrState(scf::ForOp forOp, size_t index);
	414	+
	415	+ // Get the computed PtrState for the forOp's result at the provided index.
	416	+ FailureOr<PtrState> getLoopResultPtrState(scf::ForOp forOp, size_t index);
	417	+
	418	+ // After PtrAnalysis finishes, rewrite the GetStructuredStateOp by creating
	419	+ // the correct initialization ops for offsets and strides and passing them to
	420	+ // any loop's init-args.
	421	+ LogicalResult rewriteGetStructuredStateOp(tts::GetStructuredStateOp op);
	422	+
	423	+ // Parse the state of AddPtrOp, insert any instruction needed to
	424	+ // calculate strides and offsets, build PtrState for this operand, and record
	425	+ // PtrState for knownPtrs.
	426	+ LogicalResult rewriteAddptrOp(triton::AddPtrOp op);
	427	+
	428	+
	429	+ // Parse the state of YieldOp, insert any instruction needed to calculate
	430	+ // strides and offsets, build PtrState for this operand, and record PtrState
	431	+ // in knownPtrs.
	432	+ LogicalResult
	433	+ rewriteYieldOp(scf::YieldOp op,
	434	+ llvm::SmallDenseMap<int, PtrState> &knownPtrsFor);
	435	+
	436	+ // Rewrite eligible tt.addptr in loop init args so loop can update the ; such
	437	+ // pointers over iterations. Insert any instruction needed to calculate
	438	+ // strides, offsets, and modulos.
	439	+ LogicalResult rewriteForOp(scf::ForOp op);
	440	+ LogicalResult rewriteWhileOp(scf::WhileOp op);
	441	+ FailureOr<PtrState> getWhileInitArgPtrState(scf::WhileOp whileOp, size_t index);
	442	+ FailureOr<PtrState> getWhileResultPtrState(scf::WhileOp whileOp, size_t index);
	443	+ FailureOr<PtrState> getWhileAfterArgPtrState(scf::WhileOp whileOp, size_t index);
	444	+ FailureOr<PtrState> getWhileBeforeArgPtrState(scf::WhileOp whileOp, size_t index);
	445	+
	446	+ LogicalResult analysisSplat(Operation *op, OpBuilder &builder, Value &ptr, PtrState &ptrState);
	447	+ LogicalResult extractScalarFromLoadedTensor(Operation* op, OpBuilder &builder,
	448	+ Value &loadResult, const Location loc);
	449	+
	450	+ LogicalResult rewriteScalarLoadOp(triton::LoadOp op, OpBuilder &builder,
	451	+ Value &loadResult, const Location loc);
	452	+
	453	+ LogicalResult createBroadcast(Operation *op, SmallVector<int64_t> &loadShape,
	454	+ Value &loadResult);
	455	+
	456	+ LogicalResult createReshape(Operation *op, Value &loadResult, SmallVector<int64_t> &srcShape);
	457	+
	458	+ LogicalResult analyzeMask(Operation * op, PtrState &ptrState, triton_linearize::MaskState& maskState,
	459	+ SmallVector<OpFoldResult> &dims, SmallVector<int64_t> &dimMode);
	460	+
	461	+ Value buildMaskValue(Value& ptr, PtrState& ptrState, triton_linearize::MaskState& maskState, OpBuilder& builder, Location loc);
	462	+
	463	+ Value alignMaskForFallback(Value mask, ArrayRef<int64_t> targetShape, const PtrState &ptrState, OpBuilder &builder, Location loc);
	464	+
	465	+ LogicalResult checkModifiedByAddPtrConverter(triton::LoadOp &op) const;
	466	+ LogicalResult rewriteLoadOp(triton::LoadOp op);
	467	+
	468	+ LogicalResult rewriteStoreOp(triton::StoreOp op);
	469	+
	470	+ LogicalResult rewriteOp(Operation *op);
	471	+};
	472	+
	473	+} // namespace triton
	474	+
	475	+} // namespace mlir
	476	+
	477	+#endif

Aascend/triton-adapter/include/TritonLinearize/TritonLinearize.h+19-0

@@ -0,0 +1,19 @@
	1	+#ifndef TRITON_LINERIZE_H
	2	+#define TRITON_LINERIZE_H
	3	+
	4	+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
	5	+#include "mlir/Dialect/Linalg/IR/Linalg.h"
	6	+#include "mlir/Pass/Pass.h"
	7	+#include "mlir/Transforms/DialectConversion.h"
	8	+#include "triton/Dialect/Triton/IR/Dialect.h"
	9	+
	10	+
	11	+namespace mlir {
	12	+namespace triton {
	13	+
	14	+std::unique_ptr<OperationPass<ModuleOp>> createTritonLinearizePass();
	15	+
	16	+} // namespace triton
	17	+} // namespace mlir
	18	+
	19	+#endif // TRITON_LINERIZE_H

Mascend/triton-adapter/lib/CMakeLists.txt+2-0

	@@ -5,4 +5,6 @@ add_subdirectory(TritonToLinalg)
5	add_subdirectory(TritonToLLVM)	5	add_subdirectory(TritonToLLVM)
6	add_subdirectory(DiscreteMaskAccessConversion)	6	add_subdirectory(DiscreteMaskAccessConversion)
7	add_subdirectory(TritonToUnstructure)	7	add_subdirectory(TritonToUnstructure)
		8	+add_subdirectory(Dialect)
		9	+add_subdirectory(TritonLinearize)
8	add_subdirectory(Utils)	10	add_subdirectory(Utils)

Aascend/triton-adapter/lib/Dialect/CMakeLists.txt+1-0

	@@ -0,0 +1 @@
		1	+add_subdirectory(TritonStructured)

Aascend/triton-adapter/lib/Dialect/TritonStructured/CMakeLists.txt+1-0

	@@ -0,0 +1 @@
		1	+add_subdirectory(IR)

Aascend/triton-adapter/lib/Dialect/TritonStructured/IR/CMakeLists.txt+11-0

@@ -0,0 +1,11 @@
	1	+add_triton_library(TritonStructuredIR
	2	+ TritonStructuredOps.cpp
	3	+ TritonStructuredDialect.cpp
	4	+
	5	+ DEPENDS
	6	+ TritonStructuredTableGen
	7	+
	8	+ LINK_LIBS PUBLIC
	9	+ TritonIR
	10	+ MLIRIR
	11	+ )

Aascend/triton-adapter/lib/Dialect/TritonStructured/IR/TritonStructuredDialect.cpp+22-0

@@ -0,0 +1,22 @@
	1	+#include "Dialect/TritonStructured/IR/TritonStructuredDialect.h"
	2	+
	3	+using namespace mlir;
	4	+using namespace mlir::tts;
	5	+
	6	+/// Dialect creation, the instance will be owned by the context. This is the
	7	+/// point of registration of custom types and operations for the dialect.
	8	+void TritonStructuredDialect::initialize() {
	9	+ addOperations<
	10	+#define GET_OP_LIST
	11	+#include "Dialect/TritonStructured/IR/TritonStructuredOps.cpp.inc"
	12	+ >();
	13	+}
	14	+
	15	+//===----------------------------------------------------------------------===//
	16	+// TableGen'd op method definitions
	17	+//===----------------------------------------------------------------------===//
	18	+
	19	+#define GET_OP_CLASSES
	20	+#include "Dialect/TritonStructured/IR/TritonStructuredOps.cpp.inc"
	21	+
	22	+#include "Dialect/TritonStructured/IR/TritonStructuredDialect.cpp.inc"

Aascend/triton-adapter/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp+109-0

@@ -0,0 +1,109 @@
	1	+#include "mlir/Bytecode/BytecodeOpInterface.h"
	2	+#include "mlir/Dialect/Utils/StaticValueUtils.h"
	3	+#include "mlir/IR/Builders.h"
	4	+#include "mlir/IR/BuiltinAttributes.h"
	5	+#include "mlir/IR/BuiltinTypes.h"
	6	+#include "mlir/IR/MLIRContext.h"
	7	+#include "mlir/IR/OperationSupport.h"
	8	+#include "mlir/Interfaces/SideEffectInterfaces.h"
	9	+#include "mlir/Support/LogicalResult.h"
	10	+#include "triton/Dialect/Triton/IR/Types.h"
	11	+#include "llvm/ADT/STLExtras.h"
	12	+#include "llvm/ADT/SmallVector.h"
	13	+#include "llvm/Support/Casting.h"
	14	+#include "llvm/Support/LogicalResult.h"
	15	+#include <cstdint>
	16	+#include <optional>
	17	+#include <utility>
	18	+
	19	+#define GET_OP_CLASSES
	20	+#include "Dialect/TritonStructured/IR/TritonStructuredDialect.h"
	21	+using namespace mlir;
	22	+using namespace mlir::tts;
	23	+
	24	+namespace mlir {
	25	+namespace tts {
	26	+
	27	+LogicalResult GetStructuredStateOp::verify() {
	28	+ auto expectedOffsetAndStrideTypes =
	29	+ getOffsetAndStrideTypes(getContext(), getInput().getType());
	30	+
	31	+ if (!expectedOffsetAndStrideTypes.has_value()) {
	32	+ return failure();
	33	+ }
	34	+
	35	+ auto [expectedOffsetTypes, expectedStrideTypes] =
	36	+ *expectedOffsetAndStrideTypes;
	37	+
	38	+ return success(expectedOffsetTypes.size() == getOffsets().size() &&
	39	+ llvm::equal(expectedOffsetTypes, getOffsets().getTypes()) &&
	40	+ expectedStrideTypes.size() == getStrides().size() &&
	41	+ llvm::equal(expectedStrideTypes, getStrides().getTypes()));
	42	+}
	43	+
	44	+void GetStructuredStateOp::build(OpBuilder &b, OperationState &state,
	45	+ Value val) {
	46	+ auto type = val.getType();
	47	+
	48	+ // Builder cannot fail, so we default to empty offset and stride types.
	49	+ // The invalid op will be rejected by the verifier later.
	50	+ auto [offsetTypes, strideTypes] =
	51	+ getOffsetAndStrideTypes(b.getContext(), type)
	52	+ .value_or(std::make_pair(SmallVector<Type>{}, SmallVector<Type>{}));
	53	+
	54	+ build(b, state, val.getType(), offsetTypes, strideTypes, val);
	55	+}
	56	+
	57	+std::optional<std::pair<SmallVector<Type>, SmallVector<Type>>>
	58	+GetStructuredStateOp::getOffsetAndStrideTypes(MLIRContext *context, Type type) {
	59	+ auto sizes = getOffsetAndStrideSegmentSizes(type);
	60	+ if (!sizes.has_value()) {
	61	+ return std::nullopt;
	62	+ }
	63	+ return std::make_pair(
	64	+ SmallVector<Type>(sizes->first, IndexType::get(context)),
	65	+ SmallVector<Type>(sizes->second, IndexType::get(context)));
	66	+}
	67	+
	68	+std::optional<std::pair<int32_t, int32_t>>
	69	+GetStructuredStateOp::getOffsetAndStrideSegmentSizes(Type type) {
	70	+ int32_t offsetSegmentSize = 0;
	71	+ int32_t strideSegmentSize = 0;
	72	+
	73	+ if (auto tensorType = llvm::dyn_cast<RankedTensorType>(type)) {
	74	+ if (tensorType.getElementType().isIntOrIndex()) {
	75	+ // Tensors of offsets
	76	+ // Important note:
	77	+ // We only care about tensor of index / int (in addition to pointer type)
	78	+ // because only values of int and index type can potentially be part of a
	79	+ // pointer arithmetic sequence.
	80	+ offsetSegmentSize = strideSegmentSize = tensorType.getRank();
	81	+ } else if (auto ptrType =
	82	+ dyn_cast<triton::PointerType>(tensorType.getElementType())) {
	83	+ // Unstructured pointers (tensor<!tt.ptr<type>>)
	84	+ // Each tensor of rank k gets k values for its offsets and k values for
	85	+ // its strides, all of which has Index type.
	86	+ offsetSegmentSize = strideSegmentSize = tensorType.getRank();
	87	+ }
	88	+ }
	89	+ // Block pointers (!tt.ptr<tensor<type>> or !tt.ptr<type>)
	90	+ else if (auto ptrType = llvm::dyn_cast<triton::PointerType>(type)) {
	91	+ if (auto tensorType =
	92	+ llvm::dyn_cast<RankedTensorType>(ptrType.getPointeeType())) {
	93	+ // Each tensor of rank k gets k values for its offsets and k values for
	94	+ // its strides, all of which has Index type.
	95	+ offsetSegmentSize = strideSegmentSize = tensorType.getRank();
	96	+ } else {
	97	+ // The only relevant state that can be updated in loops for scalar
	98	+ // pointers are offset. No need to include stride here.
	99	+ offsetSegmentSize = 1;
	100	+ }
	101	+ } else {
	102	+ return std::nullopt;
	103	+ }
	104	+
	105	+ return std::make_pair(offsetSegmentSize, strideSegmentSize);
	106	+}
	107	+
	108	+} // namespace tts
	109	+} // namespace mlir

Aascend/triton-adapter/lib/TritonLinearize/CMakeLists.txt+26-0

@@ -0,0 +1,26 @@
	1	+add_triton_library(TritonLinearize
	2	+ TritonLinearize.cpp
	3	+ TritonLinearizePass.cpp
	4	+ PtrAnalysis.cpp
	5	+ MaskAnalysis.cpp
	6	+ OpFoldResultUtils.cpp
	7	+
	8	+
	9	+ DEPENDS
	10	+ TritonStructuredTableGen
	11	+ TritonLinearizePassIncGen
	12	+
	13	+ LINK_LIBS PUBLIC
	14	+ MLIRArithDialect
	15	+ MLIRDialectUtils
	16	+ MLIRIR
	17	+ MLIRMathDialect
	18	+ MLIRPass
	19	+ MLIRTensorDialect
	20	+ MLIRTransforms
	21	+ MLIRSupport
	22	+ MLIRReconcileUnrealizedCasts
	23	+ TritonIR
	24	+ TritonTransforms
	25	+ TritonStructuredIR
	26	+)

Aascend/triton-adapter/lib/TritonLinearize/MaskAnalysis.cpp+809-0

Aascend/triton-adapter/lib/TritonLinearize/OpFoldResultUtils.cpp+395-0

@@ -0,0 +1,395 @@
	1	+/*
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	3	+ *
	4	+ * Permission is hereby granted, free of charge, to any person obtaining a copy
	5	+ * of this software and associated documentation files (the "Software"), to deal
	6	+ * in the Software without restriction, including without limitation the rights
	7	+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	8	+ * copies of the Software, and to permit persons to whom the Software is
	9	+ * furnished to do so, subject to the following conditions:
	10	+ *
	11	+ * The above copyright notice and this permission notice shall be included in
	12	+ * all copies or substantial portions of the Software.
	13	+ *
	14	+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	15	+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	16	+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	17	+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	18	+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	19	+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
	20	+ * THE SOFTWARE.
	21	+ */
	22	+
	23	+#include "TritonLinearize/OpFoldResultUtils.h"
	24	+
	25	+#include "mlir/Dialect/Arith/IR/Arith.h"
	26	+#include "mlir/IR/BuiltinAttributes.h"
	27	+#include "mlir/IR/BuiltinTypes.h"
	28	+#include "mlir/IR/OpDefinition.h"
	29	+#include "mlir/Transforms/DialectConversion.h"
	30	+
	31	+namespace mlir {
	32	+
	33	+Value materializeValue(OpBuilder &builder, Location loc, OpFoldResult ofr) {
	34	+ if (auto val = ofr.dyn_cast<Value>()) {
	35	+ return val;
	36	+ }
	37	+
	38	+ auto intVal = getIntAttr(ofr);
	39	+ if (intVal.has_value()) {
	40	+ return builder.create<arith::ConstantOp>(loc, builder.getI32IntegerAttr(intVal.value()));
	41	+ }
	42	+ assert(intVal.has_value());
	43	+ return Value();
	44	+
	45	+ // return builder.create<arith::ConstantIndexOp>(
	46	+ // loc, dyn_cast<IntegerAttr>(attr).getInt());
	47	+}
	48	+
	49	+std::optional<int64_t> getIntAttr(const OpFoldResult ofr) {
	50	+ if (ofr.is<Attribute>() && isa<IntegerAttr>(ofr.get<Attribute>()))
	51	+ return dyn_cast<IntegerAttr>(ofr.get<Attribute>()).getInt();
	52	+
	53	+
	54	+ return std::nullopt;
	55	+}
	56	+
	57	+bool hasConstZero(const OpFoldResult ofr) {
	58	+ auto intAttr = getIntAttr(ofr);
	59	+ if (intAttr.has_value()) {
	60	+ if (intAttr.value() == 0) {
	61	+ return true;
	62	+ }
	63	+ return false;
	64	+ }
	65	+
	66	+ auto val = dyn_cast<Value>(ofr);
	67	+ assert(val);
	68	+ auto constOp = val.getDefiningOp<arith::ConstantOp>();
	69	+ if (!constOp)
	70	+ return false;
	71	+
	72	+ intAttr = getIntAttr(constOp.getValue());
	73	+ if (intAttr.has_value()) {
	74	+ if (intAttr.value() == 0) {
	75	+ return true;
	76	+ }
	77	+ return false;
	78	+ }
	79	+
	80	+ return false;
	81	+}
	82	+
	83	+Value ofrToIndexValue(const OpFoldResult ofr, const Location loc,
	84	+ OpBuilder &b) {
	85	+ if (Value val = dyn_cast<Value>(ofr)) {
	86	+ assert(val.getType().isIntOrIndex());
	87	+ if (!val.getType().isIndex()) {
	88	+ val = b.create<arith::IndexCastOp>(loc, b.getIndexType(), val);
	89	+ }
	90	+ return val;
	91	+ }
	92	+
	93	+ auto intVal = getIntAttr(ofr);
	94	+ if (intVal.has_value()) {
	95	+ return b.create<arith::ConstantOp>(loc, b.getIndexAttr(intVal.value()));
	96	+ }
	97	+ llvm_unreachable("Unexpected OpFoldResult state");
	98	+ return nullptr;
	99	+}
	100	+
	101	+SmallVector<Value> ofrsToIndexValues(ArrayRef<OpFoldResult> ofrs,
	102	+ const Location loc, OpBuilder &b) {
	103	+ return llvm::to_vector<4>(
	104	+ llvm::map_range(ofrs, [&](OpFoldResult ofr) -> Value {
	105	+ return ofrToIndexValue(ofr, loc, b);
	106	+ }));
	107	+}
	108	+
	109	+OpFoldResult addOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	110	+ const Location loc, OpBuilder &b) {
	111	+ auto lhsIntAttr = getIntAttr(lhs);
	112	+ auto rhsIntAttr = getIntAttr(rhs);
	113	+
	114	+ // shortcut for special cases
	115	+ if (!lhsIntAttr && rhsIntAttr && rhsIntAttr.value() == 0)
	116	+ return lhs;
	117	+ if (!rhsIntAttr && lhsIntAttr && lhsIntAttr.value() == 0)
	118	+ return rhs;
	119	+
	120	+ // both lhs and rhs are constants, return result directly
	121	+ if (lhsIntAttr && rhsIntAttr)
	122	+ return b.getIndexAttr(lhsIntAttr.value() + rhsIntAttr.value());
	123	+
	124	+ auto indextype = b.getIndexType();
	125	+ // otherwise, need to create instructions to calculate new attribute value
	126	+ auto lhsValue = dyn_cast<Value>(lhs);
	127	+ if (lhsIntAttr) {
	128	+ auto lhsOp =
	129	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(lhsIntAttr.value()));
	130	+ lhsValue = lhsOp.getResult();
	131	+ } else {
	132	+ if(isa<IntegerType>(lhsValue.getType())){
	133	+ lhsValue = b.create<arith::IndexCastOp>(loc, indextype, lhsValue).getResult();
	134	+ }
	135	+ assert(isa<IndexType>(lhsValue.getType()));
	136	+ }
	137	+
	138	+ auto rhsValue = dyn_cast<Value>(rhs);
	139	+ if (rhsIntAttr) {
	140	+ auto rhsOp =
	141	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(rhsIntAttr.value()));
	142	+ rhsValue = rhsOp.getResult();
	143	+ } else {
	144	+ if(isa<IntegerType>(rhsValue.getType())){
	145	+ rhsValue = b.create<arith::IndexCastOp>(loc, indextype, rhsValue).getResult();
	146	+ }
	147	+ assert(isa<IndexType>(rhsValue.getType()));
	148	+ }
	149	+
	150	+ return b.create<arith::AddIOp>(loc, lhsValue, rhsValue).getResult();
	151	+}
	152	+
	153	+OpFoldResult subOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	154	+ const Location loc, OpBuilder &b) {
	155	+ auto lhsIntAttr = getIntAttr(lhs);
	156	+ auto rhsIntAttr = getIntAttr(rhs);
	157	+
	158	+ // shortcut for special cases
	159	+ if (!lhsIntAttr && rhsIntAttr && rhsIntAttr.value() == 0)
	160	+ return lhs;
	161	+
	162	+ // both lhs and rhs are constants, return result directly
	163	+ if (lhsIntAttr && rhsIntAttr)
	164	+ return b.getIndexAttr(lhsIntAttr.value() - rhsIntAttr.value());
	165	+
	166	+ // otherwise, need to create instructions to calculate new attribute value
	167	+ auto lhsValue = dyn_cast<Value>(lhs);
	168	+ if (lhsIntAttr) {
	169	+ auto lhsOp =
	170	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(lhsIntAttr.value()));
	171	+ lhsValue = lhsOp.getResult();
	172	+ }
	173	+
	174	+ auto rhsValue = dyn_cast<Value>(rhs);
	175	+ if (rhsIntAttr) {
	176	+ auto rhsOp =
	177	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(rhsIntAttr.value()));
	178	+ rhsValue = rhsOp.getResult();
	179	+ }
	180	+
	181	+ auto sumOp = b.create<arith::SubIOp>(loc, lhsValue, rhsValue);
	182	+ return sumOp.getResult();
	183	+}
	184	+
	185	+OpFoldResult mulOFRValue(const OpFoldResult lhs, const Value rhs,
	186	+ const Location loc, OpBuilder &b) {
	187	+ auto lhsIntAttr = getIntAttr(lhs);
	188	+
	189	+ auto rhsIsConst = false;
	190	+ // if rhs is not a const, use max value since min is used to represent
	191	+ // dynamic size or stride
	192	+ auto rhsConstValue = std::numeric_limits<int64_t>::max();
	193	+ auto rhsOp = rhs.getDefiningOp<arith::ConstantOp>();
	194	+ if (rhsOp) {
	195	+ rhsIsConst = true;
	196	+ rhsConstValue = cast<IntegerAttr>(rhsOp.getValue()).getInt();
	197	+ }
	198	+
	199	+ // 0. both lhs and rhs are constants
	200	+ if (lhsIntAttr && rhsIsConst)
	201	+ return b.getIndexAttr(lhsIntAttr.value() * rhsConstValue);
	202	+
	203	+ // shortcuts for special cases
	204	+ if (lhsIntAttr) {
	205	+ if (lhsIntAttr.value() == 0)
	206	+ return lhs;
	207	+ if (lhsIntAttr.value() == 1)
	208	+ return rhs;
	209	+ }
	210	+ if (rhsIsConst) {
	211	+ if (rhsConstValue == 0)
	212	+ return rhsOp.getResult();
	213	+ if (rhsConstValue == 1)
	214	+ return lhs;
	215	+ }
	216	+
	217	+
	218	+ // 1. if lhs is constant but rhs is not
	219	+ if (lhsIntAttr && !rhsIsConst) {
	220	+ auto lhsConstOp =
	221	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(lhsIntAttr.value()));
	222	+ auto mulOp = b.create<arith::MulIOp>(loc, lhsConstOp.getResult(), rhs);
	223	+ return mulOp.getResult();
	224	+ }
	225	+
	226	+ // 2. if lhs is not constant
	227	+ assert(!lhsIntAttr);
	228	+ auto mulOp = b.create<arith::MulIOp>(loc, lhs.get<Value>(), rhs);
	229	+ return mulOp.getResult();
	230	+}
	231	+
	232	+OpFoldResult divOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	233	+ const Location loc, OpBuilder &b) {
	234	+ auto lhsIntAttr = getIntAttr(lhs);
	235	+ auto rhsIntAttr = getIntAttr(rhs);
	236	+
	237	+ // both lhs and rhs are constants, return result directly
	238	+ if (lhsIntAttr && rhsIntAttr)
	239	+ return b.getIndexAttr(lhsIntAttr.value() / rhsIntAttr.value());
	240	+
	241	+ // shortcut for special cases
	242	+ if (rhsIntAttr && rhsIntAttr.value() == 1)
	243	+ return lhs;
	244	+
	245	+ // otherwise, need to create instructions to calculate new attribute value
	246	+ auto lhsValue = lhs.dyn_cast<Value>();
	247	+ if (lhsIntAttr) {
	248	+ auto lhsOp =
	249	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(lhsIntAttr.value()));
	250	+ lhsValue = lhsOp.getResult();
	251	+ }
	252	+
	253	+ auto rhsValue = rhs.dyn_cast<Value>();
	254	+ if (rhsIntAttr) {
	255	+ auto rhsOp =
	256	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(rhsIntAttr.value()));
	257	+ rhsValue = rhsOp.getResult();
	258	+ }
	259	+
	260	+ auto divOp = b.create<arith::DivSIOp>(loc, lhsValue, rhsValue);
	261	+ return divOp.getResult();
	262	+}
	263	+
	264	+OpFoldResult remOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	265	+ const Location loc, OpBuilder &b) {
	266	+ auto lhsIntAttr = getIntAttr(lhs);
	267	+ auto rhsIntAttr = getIntAttr(rhs);
	268	+
	269	+ // both lhs and rhs are constants, return result directly
	270	+ if (lhsIntAttr && rhsIntAttr)
	271	+ return b.getIndexAttr(lhsIntAttr.value() % rhsIntAttr.value());
	272	+
	273	+ // shortcut for special cases
	274	+ if (rhsIntAttr && rhsIntAttr.value() == 1)
	275	+ return b.getIndexAttr(0);
	276	+
	277	+ // otherwise, need to create instructions to calculate new attribute value
	278	+ auto lhsValue = lhs.dyn_cast<Value>();
	279	+ if (lhsIntAttr) {
	280	+ auto lhsOp =
	281	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(lhsIntAttr.value()));
	282	+ lhsValue = lhsOp.getResult();
	283	+ }
	284	+
	285	+ auto rhsValue = rhs.dyn_cast<Value>();
	286	+ if (rhsIntAttr) {
	287	+ auto rhsOp =
	288	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(rhsIntAttr.value()));
	289	+ rhsValue = rhsOp.getResult();
	290	+ }
	291	+
	292	+ auto remOp = b.create<arith::RemSIOp>(loc, lhsValue, rhsValue);
	293	+ return remOp.getResult();
	294	+}
	295	+
	296	+OpFoldResult mulOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	297	+ const Location loc, OpBuilder &b) {
	298	+ auto lhsIntAttr = getIntAttr(lhs);
	299	+ auto rhsIntAttr = getIntAttr(rhs);
	300	+
	301	+ // both lhs and rhs are constants, return result directly
	302	+ if (lhsIntAttr && rhsIntAttr)
	303	+ return b.getIndexAttr(lhsIntAttr.value() * rhsIntAttr.value());
	304	+
	305	+ // shortcuts for special cases
	306	+ if (lhsIntAttr) {
	307	+ if (lhsIntAttr.value() == 0)
	308	+ return lhs;
	309	+ if (lhsIntAttr.value() == 1)
	310	+ return rhs;
	311	+ }
	312	+ if (rhsIntAttr) {
	313	+ if (rhsIntAttr.value() == 0)
	314	+ return rhs;
	315	+ if (rhsIntAttr.value() == 1)
	316	+ return lhs;
	317	+ }
	318	+
	319	+
	320	+ // otherwise, need to create instructions to calculate new attribute value
	321	+ auto lhsValue = dyn_cast<Value>(lhs);
	322	+ if (lhsIntAttr) {
	323	+ auto lhsOp =
	324	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(lhsIntAttr.value()));
	325	+ lhsValue = lhsOp.getResult();
	326	+ }
	327	+
	328	+ auto rhsValue = dyn_cast<Value>(rhs);
	329	+ if (rhsIntAttr) {
	330	+ auto rhsOp =
	331	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(rhsIntAttr.value()));
	332	+ rhsValue = rhsOp.getResult();
	333	+ }
	334	+
	335	+ auto mulOp = b.create<arith::MulIOp>(loc, lhsValue, rhsValue);
	336	+ return mulOp.getResult();
	337	+}
	338	+
	339	+OpFoldResult minOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	340	+ const Location loc, OpBuilder &b) {
	341	+ auto lhsIntAttr = getIntAttr(lhs);
	342	+ auto rhsIntAttr = getIntAttr(rhs);
	343	+
	344	+ // both lhs and rhs are constants, return result directly
	345	+ if (lhsIntAttr && rhsIntAttr)
	346	+ return b.getIndexAttr(std::min(lhsIntAttr.value(), rhsIntAttr.value()));
	347	+
	348	+ // otherwise, need to create instructions to calculate new attribute value
	349	+ auto lhsValue = dyn_cast<Value>(lhs);
	350	+ if (lhsIntAttr) {
	351	+ auto lhsOp =
	352	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(lhsIntAttr.value()));
	353	+ lhsValue = lhsOp.getResult();
	354	+ }
	355	+
	356	+ auto rhsValue = dyn_cast<Value>(rhs);
	357	+ if (rhsIntAttr) {
	358	+ auto rhsOp =
	359	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(rhsIntAttr.value()));
	360	+ rhsValue = rhsOp.getResult();
	361	+ }
	362	+
	363	+ auto minOp = b.create<arith::MinSIOp>(loc, lhsValue, rhsValue);
	364	+ return minOp.getResult();
	365	+}
	366	+
	367	+OpFoldResult maxOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
	368	+ const Location loc, OpBuilder &b) {
	369	+ auto lhsIntAttr = getIntAttr(lhs);
	370	+ auto rhsIntAttr = getIntAttr(rhs);
	371	+
	372	+ // both lhs and rhs are constants, return result directly
	373	+ if (lhsIntAttr && rhsIntAttr)
	374	+ return b.getIndexAttr(std::max(lhsIntAttr.value(), rhsIntAttr.value()));
	375	+
	376	+ // otherwise, need to create instructions to calculate new attribute value
	377	+ auto lhsValue = dyn_cast<Value>(lhs);
	378	+ if (lhsIntAttr) {
	379	+ auto lhsOp =
	380	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(lhsIntAttr.value()));
	381	+ lhsValue = lhsOp.getResult();
	382	+ }
	383	+
	384	+ auto rhsValue = dyn_cast<Value>(rhs);
	385	+ if (rhsIntAttr) {
	386	+ auto rhsOp =
	387	+ b.create<arith::ConstantOp>(loc, b.getIndexAttr(rhsIntAttr.value()));
	388	+ rhsValue = rhsOp.getResult();
	389	+ }
	390	+
	391	+ auto maxOp = b.create<arith::MaxSIOp>(loc, lhsValue, rhsValue);
	392	+ return maxOp.getResult();
	393	+}
	394	+
	395	+} // namespace mlir

Aascend/triton-adapter/lib/TritonLinearize/PtrAnalysis.cpp+3623-0

Aascend/triton-adapter/lib/TritonLinearize/TritonLinearize.cpp+49-0

@@ -0,0 +1,49 @@
	1	+/*
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	3	+ *
	4	+ * Permission is hereby granted, free of charge, to any person obtaining a copy
	5	+ * of this software and associated documentation files (the "Software"), to deal
	6	+ * in the Software without restriction, including without limitation the rights
	7	+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	8	+ * copies of the Software, and to permit persons to whom the Software is
	9	+ * furnished to do so, subject to the following conditions:
	10	+ *
	11	+ * The above copyright notice and this permission notice shall be included in
	12	+ * all copies or substantial portions of the Software.
	13	+ *
	14	+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	15	+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	16	+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	17	+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	18	+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	19	+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
	20	+ * THE SOFTWARE.
	21	+ */
	22	+
	23	+#include "Utils/Utils.h"
	24	+#include "triton/Dialect/Triton/IR/Dialect.h"
	25	+#include "TritonLinearize/TritonLinearize.h"
	26	+
	27	+#include "llvm/ADT/SmallVectorExtras.h"
	28	+#include "llvm/ADT/TypeSwitch.h"
	29	+#include "llvm/Support/ErrorHandling.h"
	30	+#include "llvm/Support/LogicalResult.h"
	31	+#include "llvm/Support/raw_ostream.h"
	32	+#include "llvm/Support/FormatVariadic.h"
	33	+#include <utility>
	34	+
	35	+#include "mlir/Dialect/Arith/IR/Arith.h"
	36	+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
	37	+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
	38	+#include "mlir/Dialect/Linalg/IR/Linalg.h"
	39	+#include "mlir/Dialect/MemRef/IR/MemRef.h"
	40	+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
	41	+#include "mlir/IR/OpDefinition.h"
	42	+#include "mlir/IR/ValueRange.h"
	43	+#include "TritonLinearize/OpFoldResultUtils.h"
	44	+
	45	+namespace mlir{
	46	+namespace triton{
	47	+
	48	+}
	49	+}

Aascend/triton-adapter/lib/TritonLinearize/TritonLinearizePass.cpp+348-0

@@ -0,0 +1,348 @@
	1	+/*
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	3	+ *
	4	+ * Permission is hereby granted, free of charge, to any person obtaining a copy
	5	+ * of this software and associated documentation files (the "Software"), to deal
	6	+ * in the Software without restriction, including without limitation the rights
	7	+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	8	+ * copies of the Software, and to permit persons to whom the Software is
	9	+ * furnished to do so, subject to the following conditions:
	10	+ *
	11	+ * The above copyright notice and this permission notice shall be included in
	12	+ * all copies or substantial portions of the Software.
	13	+ *
	14	+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	15	+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	16	+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	17	+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	18	+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	19	+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
	20	+ * THE SOFTWARE.
	21	+ */
	22	+
	23	+#include "mlir/Dialect/Linalg/IR/Linalg.h"
	24	+#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
	25	+#include "mlir/Dialect/SCF/Transforms/Patterns.h"
	26	+#include "mlir/IR/BuiltinAttributes.h"
	27	+#include "mlir/IR/BuiltinOps.h"
	28	+#include "mlir/IR/BuiltinTypes.h"
	29	+#include "mlir/IR/MLIRContext.h"
	30	+#include "mlir/IR/TypeRange.h"
	31	+#include "mlir/IR/Types.h"
	32	+#include "mlir/IR/ValueRange.h"
	33	+#include "mlir/Support/LogicalResult.h"
	34	+
	35	+#include "TritonLinearize/OpFoldResultUtils.h"
	36	+#include "TritonLinearize/MaskAnalysis.h"
	37	+#include "TritonLinearize/PtrAnalysis.h"
	38	+#include "TritonLinearize/TritonLinearize.h"
	39	+#include "Dialect/TritonStructured/IR/TritonStructuredDialect.h"
	40	+
	41	+#include "triton/Dialect/Triton/IR/Dialect.h"
	42	+#include "mlir/Dialect/Affine/IR/AffineOps.h"
	43	+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
	44	+#include "mlir/Pass/PassManager.h"
	45	+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
	46	+#include "mlir/Transforms/OneToNTypeConversion.h"
	47	+#include "mlir/Transforms/Passes.h"
	48	+#include "triton/Dialect/Triton/IR/Types.h"
	49	+
	50	+#include "llvm/ADT/STLExtras.h"
	51	+#include "llvm/ADT/SmallVector.h"
	52	+#include "llvm/Support/Casting.h"
	53	+#include "llvm/Support/Debug.h"
	54	+#include "llvm/Support/LogicalResult.h"
	55	+#include <cassert>
	56	+#include <optional>
	57	+
	58	+#define DEBUG_TYPE "triton-to-structured"
	59	+
	60	+namespace mlir {
	61	+namespace triton {
	62	+
	63	+#define GEN_PASS_CLASSES
	64	+#include "TritonLinearize/Passes.h.inc"
	65	+
	66	+//namespace {
	67	+
	68	+class TritonLinearizePass
	69	+ : public TritonLinearizeBase<TritonLinearizePass>
	70	+{
	71	+private:
	72	+ static TupleType getStructuredStateTupleType(MLIRContext *context, Type t) {
	73	+ SmallVector<Type> tupleTypes{t};
	74	+ auto [offsetTypes, strideTypes] =
	75	+ *tts::GetStructuredStateOp::getOffsetAndStrideTypes(context, t);
	76	+ tupleTypes.append(offsetTypes);
	77	+ tupleTypes.append(strideTypes);
	78	+ return TupleType::get(context, tupleTypes);
	79	+ }
	80	+
	81	+public:
	82	+ void getDependentDialects(DialectRegistry &registry) const override {
	83	+ registry
	84	+ .insert<arith::ArithDialect, math::MathDialect, affine::AffineDialect,
	85	+ scf::SCFDialect, tensor::TensorDialect, triton::TritonDialect,
	86	+ linalg::LinalgDialect, tts::TritonStructuredDialect>();
	87	+ }
	88	+
	89	+ LogicalResult convertToPointerTupleWithOffsetsAndStrides() {
	90	+ auto moduleOp = getOperation();
	91	+
	92	+ RewritePatternSet patterns(&getContext());
	93	+
	94	+ auto context = &getContext();
	95	+ OneToNTypeConverter converter;
	96	+ converter.addConversion([](Type type) { return type; });
	97	+
	98	+ // We are doing a 1->1 type conversion here, where a triton pointer type
	99	+ // maps to a tuple of {pointer, offset_0, offset_1,..., stride_0,
	100	+ // stride_1,...} type.
	101	+ //
	102	+ // Case 1: Unstructured pointers (tensor<!tt.ptr<type>>)
	103	+ converter.addConversion([context](RankedTensorType tensorType,
	104	+ SmallVectorImpl<Type> &types)
	105	+ -> std::optional<LogicalResult> {
	106	+ // Important note:
	107	+ // We only care about tensor of index / int (in addition to pointer type)
	108	+ // because only values of int and index type can potentially be part of a
	109	+ // pointer arithmetic sequence.
	110	+ if (!isa<triton::PointerType>(tensorType.getElementType()) &&
	111	+ !tensorType.getElementType().isIntOrIndex()) {
	112	+ return std::nullopt;
	113	+ }
	114	+ types =
	115	+ SmallVector<Type>{getStructuredStateTupleType(context, tensorType)};
	116	+ return success();
	117	+ });
	118	+
	119	+ // Case 2: Block pointers (!tt.ptr<tensor<type>> or !tt.ptr<type>)
	120	+ converter.addConversion([context](triton::PointerType ptrType,
	121	+ SmallVectorImpl<Type> &types)
	122	+ -> std::optional<LogicalResult> {
	123	+ types = SmallVector<Type>{getStructuredStateTupleType(context, ptrType)};
	124	+ return success();
	125	+ });
	126	+
	127	+ // Hooks to compute the correct materialization, "argument" and "source"
	128	+ // materialization are used when we need to convert the tuple type back to
	129	+ // the original triton pointer type. These are used when there are ops that
	130	+ // still need to use the original pointer type. For instance, we convert the
	131	+ // result of tt.addptr from tt.ptr type to a tuple, but the original ptr
	132	+ // result is still being used by another tt.load or tt.store.
	133	+ auto materialize = [](OpBuilder &builder, Type resultType,
	134	+ ValueRange inputs, Location loc) {
	135	+ return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
	136	+ .getResult(0);
	137	+ };
	138	+
	139	+ converter.addArgumentMaterialization(materialize);
	140	+ converter.addSourceMaterialization(materialize);
	141	+
	142	+ // Compute the target materialization, given a value with the pointer type,
	143	+ // convert that value to a tuple type.
	144	+ converter.addTargetMaterialization(
	145	+ [](OpBuilder &builder, TypeRange resultTypes, Value input,
	146	+ Location loc) -> std::optional<SmallVector<Value>> {
	147	+ return builder
	148	+ .create<UnrealizedConversionCastOp>(loc, resultTypes, input)
	149	+ ->getResults();
	150	+ });
	151	+
	152	+ scf::populateSCFStructuralOneToNTypeConversions(converter, patterns);
	153	+
	154	+ if (failed(applyPartialOneToNConversion(getOperation(), converter,
	155	+ std::move(patterns)))) {
	156	+ return failure();
	157	+ }
	158	+
	159	+ PassManager pm(&getContext(), moduleOp.getOperationName());
	160	+ pm.addPass(createCanonicalizerPass());
	161	+ if (failed(runPipeline(pm, getOperation()))) {
	162	+ return failure();
	163	+ }
	164	+
	165	+ return success();
	166	+ }
	167	+
	168	+ LogicalResult decomposePointerTuple() {
	169	+ auto moduleOp = getOperation();
	170	+
	171	+ auto context = &getContext();
	172	+ OneToNTypeConverter converter;
	173	+ converter.addConversion([](Type type) { return type; });
	174	+
	175	+ // We are doing a 1->N type conversion here, where a pointer tuple type
	176	+ // maps to a sequence of {pointer, offset_0, offset_1,..., stride_0,
	177	+ // stride_1,...}
	178	+ converter.addConversion(
	179	+ [context](TupleType tupleType, SmallVectorImpl<Type> &types)
	180	+ -> std::optional<LogicalResult> {
	181	+ tupleType.getFlattenedTypes(types);
	182	+ return success();
	183	+ });
	184	+
	185	+ // Hooks to compute the correct materialization, "argument" and "source"
	186	+ // materialization are used when we need to convert a series of {pointer,
	187	+ // offset_0, offset_1,..., stride_0, stride_1,...} type back to the "pointer
	188	+ // tuple type".
	189	+ //
	190	+ // Because we actually want to get rid of the tuple type, return `inputs[0]`
	191	+ // which corresponds to a "triton pointer type". This approach will work as
	192	+ // intended because the ops that currently take "pointer tuple type" are
	193	+ // `unrealized_conversion_cast` ops which will get removed below during
	194	+ // reconcile-unrealized-conversion-casts.
	195	+ auto materialize = [](OpBuilder &builder, Type resultType,
	196	+ ValueRange inputs,
	197	+ Location loc) { return inputs[0]; };
	198	+ converter.addArgumentMaterialization(materialize);
	199	+ converter.addSourceMaterialization(materialize);
	200	+
	201	+ // For each value of "pointer tuple type" that gets decomposed into a
	202	+ // sequence of {pointer, offset_0, offset_1,..., stride_0, stride_1,...},
	203	+ // create a `tts.get_structured_state` op that serves as a placeholder.
	204	+ // The return values for this op will be used as the init-args for scf.for.
	205	+ // At the end of pointer analysis, we will use the PtrState to create the
	206	+ // correct offsets, strides, and remove these ops.
	207	+ converter.addTargetMaterialization([](OpBuilder &builder,
	208	+ TypeRange resultTypes, Value input,
	209	+ Location loc) {
	210	+ auto placeholder = builder.create<tts::GetStructuredStateOp>(
	211	+ loc, input.getDefiningOp()->getOperand(0));
	212	+ assert(llvm::equal(placeholder.getResultTypes(), resultTypes));
	213	+ return placeholder.getResults();
	214	+ });
	215	+
	216	+ RewritePatternSet patterns(&getContext());
	217	+ scf::populateSCFStructuralOneToNTypeConversions(converter, patterns);
	218	+ if (failed(applyPartialOneToNConversion(getOperation(), converter,
	219	+ std::move(patterns)))) {
	220	+ return failure();
	221	+ }
	222	+
	223	+ // Note:
	224	+ // Be careful not to run canonicalization here, because the
	225	+ // tts.get_structured_state ops created above are just placeholders and
	226	+ // don't have any effects. Canonicalization will remove them altogether.
	227	+ PassManager pm(&getContext(), moduleOp.getOperationName());
	228	+ pm.addPass(mlir::createReconcileUnrealizedCastsPass());
	229	+ if (failed(runPipeline(pm, getOperation()))) {
	230	+ signalPassFailure();
	231	+ }
	232	+
	233	+ return success();
	234	+ }
	235	+
	236	+ // Prepass that inserts `tts.get_structured_state` ops. These ops are used as
	237	+ // placeholders to make passing structured pointer state into scf.for loop's
	238	+ // init args easier, especially with multiple levels of loops.
	239	+ //
	240	+ // Background:
	241	+ //
	242	+ // PtrAnalysis computes a PtrState for every operand (or triton value)
	243	+ // involved in a sequence of pointer arithmetic; some examples include: triton
	244	+ // pointer, offsets (which could be a tensor of indices or just a simple index
	245	+ // value).
	246	+ //
	247	+ // If a triton value is updated and returned in a scf.for op, it means
	248	+ // that we have to carry its offsets and strides in the scf.for's iterargs.
	249	+ //
	250	+ // Previously, we have to manually rewrite the loops to include the
	251	+ // relevant information from a PtrState which was rather involved and
	252	+ // error-prone; this was also hard to scale up to multiple level of loops
	253	+ // because there are several book-keeping data structures that we have to
	254	+ // maintain.
	255	+ //
	256	+ // With the introduction of the prepass that inserts
	257	+ // `tts.get_structured_state`. The return values of these ops, which include a
	258	+ // triton value with its original result type and its corresponding offsets
	259	+ // and strides, will be used as "placeholders" into the scf.for's init-args.
	260	+ // We leverage standard MLIR infrastructure 1->N conversion to perform this
	261	+ // rewrite, which helps simplify the logic significantly.
	262	+ //
	263	+ // After PtrAnalysis finishes, the return values of these
	264	+ // `tts.get_structured_state` ops will be remapped to the correct
	265	+ // initialization of the value's offsets and strides through the value's
	266	+ // computed PtrState.
	267	+ //
	268	+ // Implementation details:
	269	+ // In essence, what we really want to do in the prepass is, for every value
	270	+ // of triton-pointer-like type (tt.ptr or tensor<tt.ptr<>>) and tensor of
	271	+ // indices (tensor<i32>) which might be used in a sequence of pointer
	272	+ // arithmetic, we want to create an op `tts.get_structured_state` that takes
	273	+ // in the original triton value and returns a series of values:
	274	+ //
	275	+ // {triton_value, offset_0, offset_1, ..., stride_0, stride_1,...}
	276	+ //
	277	+ // Applying the above conversion will also mean that any structural ops such
	278	+ // as scf.for and scf.yield that originally takes the triton pointer will
	279	+ // then take {triton_value, offset_0, offset_1, ..., stride_0, stride_1,...}.
	280	+ //
	281	+ // The 1->N type conversion is a perfect fit for this transformation.
	282	+ // Unfortunately, we cannot do this is one pass, because the current 1->N
	283	+ // type conversion implementation for scf.for ops doesn't provide us with a
	284	+ // way to detect that a type conversion is recursive. So a triton_value type
	285	+ // that gets converted to a {triton_value, offset_0, offset_1, ..., stride_0,
	286	+ // stride_1,...} will recursively trigger other conversions.
	287	+ //
	288	+ // To fix this issue, we have to first convert triton_value to
	289	+ // tuple<triton_value, offset_0, offset_1, ..., stride_0, stride_1,...>.
	290	+ // Finally, we decompose these tuples into the desired sequence.
	291	+ //
	292	+ // Note that even though the type conversion happens for every integer tensor
	293	+ // appearing in loops' iter-args, this conversion is reversible. If the
	294	+ // integer tensor isn't used in a pointer arithmetic sequence,
	295	+ // canonicalization will remove all the `tts.get_structured_state` ops and
	296	+ // revert the IR back to its original form.
	297	+ LogicalResult runTritonToStructuredPrepass() {
	298	+ if (failed(convertToPointerTupleWithOffsetsAndStrides())) {
	299	+ return failure();
	300	+ }
	301	+
	302	+ if( failed( decomposePointerTuple())){
	303	+ return failure() ;
	304	+ }
	305	+
	306	+ return success() ;
	307	+
	308	+ }
	309	+
	310	+ void runOnOperation() override {
	311	+ if (failed(runTritonToStructuredPrepass())) {
	312	+ signalPassFailure();
	313	+ return;
	314	+ }
	315	+
	316	+ PtrAnalysis ptrAnalysis;
	317	+ auto moduleOp = getOperation();
	318	+ ptrAnalysis.initializeMaybeStructuredArgs(moduleOp);
	319	+ if (failed(ptrAnalysis.rewriteOp(moduleOp))) {
	320	+ moduleOp->emitWarning("PtrAnalysis failed");
	321	+ signalPassFailure();
	322	+ }
	323	+
	324	+ // Now that all the PtrStates have been populated, we can wire up the states
	325	+ // with the tts.get_structured_state ops inserted in the prepass.
	326	+ moduleOp.walk([&ptrAnalysis](tts::GetStructuredStateOp op) {
	327	+ if (failed(ptrAnalysis.rewriteGetStructuredStateOp(op))) {
	328	+ op->emitError("Rewriting GetStructuredStateOp failed.");
	329	+ }
	330	+ });
	331	+
	332	+ PassManager pm(&getContext(), moduleOp.getOperationName());
	333	+ pm.addPass(createCSEPass());
	334	+ pm.addPass(createCanonicalizerPass());
	335	+ if (failed(runPipeline(pm, getOperation()))) {
	336	+ signalPassFailure();
	337	+ }
	338	+ }
	339	+};
	340	+
	341	+//} // namespace
	342	+
	343	+std::unique_ptr<OperationPass<ModuleOp>> createTritonLinearizePass()
	344	+{
	345	+ return std::make_unique<TritonLinearizePass>();
	346	+}
	347	+
	348	+}}

Mascend/triton-adapter/lib/TritonToLinalg/BlockPtrAnalysis.cpp+11-3

	@@ -257,9 +257,17 @@ void BlockData::mulBlock(BlockData &lBlock, BlockData &rBlock, Location loc,
257		257
258	assert(!(lBlock.hasSource() && rBlock.hasSource()));	258	assert(!(lBlock.hasSource() && rBlock.hasSource()));
259		259
260	- assert(	260	+ if (lBlock.isScalar() && rBlock.isScalar()) {
261	- (lBlock.isScalar() ^ rBlock.isScalar()) &&	261	+ LLVM_DEBUG({llvm::dbgs() << "lBlock.scalar:" << lBlock.getScalar()
262	- "Currently only support one and only one scalar in function mulBlock()");	262	+ << " rBlbock.scalar:" << rBlock.getScalar() << "\n"; });
		263	+
		264	+ auto scalar = mulOpFoldResult(lBlock.getScalar(), rBlock.getScalar(), loc, rewriter);
		265	+ this->scalar = scalar;
		266	+ }
		267	+
		268	+ // assert(
		269	+ // (lBlock.isScalar() ^ rBlock.isScalar()) &&
		270	+ // "Currently only support one and only one scalar in function mulBlock()");
263		271
264	BlockData *lb = &lBlock;	272	BlockData *lb = &lBlock;
265	BlockData *rb = &rBlock;	273	BlockData *rb = &rBlock;

Mascend/triton-adapter/tools/triton-adapter-opt/CMakeLists.txt+1-0

	@@ -14,6 +14,7 @@ target_link_libraries(triton-adapter-opt PRIVATE
14	TritonToLLVM	14	TritonToLLVM
15	TritonToUnstructure	15	TritonToUnstructure
16	TritonTransforms	16	TritonTransforms
		17	+ TritonLinearize
17	${dialect_libs}	18	${dialect_libs}
18	${conversion_libs}	19	${conversion_libs}
19	TritonGPUTransforms	20	TritonGPUTransforms

Mascend/triton-adapter/tools/triton-adapter-opt/triton-adapter-opt.cpp+2-0

	@@ -27,6 +27,7 @@
27	#include "TritonToLinalg/Passes.h"	27	#include "TritonToLinalg/Passes.h"
28	#include "TritonToLLVM/Passes.h"	28	#include "TritonToLLVM/Passes.h"
29	#include "TritonToUnstructure/Passes.h"	29	#include "TritonToUnstructure/Passes.h"
		30	+#include "TritonLinearize/Passes.h"
30	#include "bishengir/InitAllDialects.h"	31	#include "bishengir/InitAllDialects.h"
31	#include "mlir/Dialect/Bufferization/IR/Bufferization.h"	32	#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
32	#include "mlir/Dialect/Func/IR/FuncOps.h"	33	#include "mlir/Dialect/Func/IR/FuncOps.h"
	@@ -52,6 +53,7 @@ int main(int argc, char **argv) {
52	bishengir::registerAllDialects(registry);	53	bishengir::registerAllDialects(registry);
53		54
54	// Register all passes.	55	// Register all passes.
		56	+ mlir::triton::registerTritonLinearizePass();
55	mlir::triton::registerTritonToLinalgPass();	57	mlir::triton::registerTritonToLinalgPass();
56	mlir::triton::registerTritonToLLVMPass();	58	mlir::triton::registerTritonToLLVMPass();
57	mlir::triton::registerTritonToAnnotationPass();	59	mlir::triton::registerTritonToAnnotationPass();