Star977
421
代码
代码
Issues405
Pull Requests1030
项目讨论
Wiki
分析
Star977
421
已合并
[AscendNPU IR][SIMT] Added SIMT support for non power of two tensors #1357
Max_Wang_Huawei创建于 7月4日
[AscendNPU IR][SIMT] Added SIMT support for non power of two tensors #1357
已合并
Max_Wang_Huawei创建于 7月4日
从Max_Wang_Huawei/AscendNPU-IRfeature-convert-non-power-two-tensors合入到Ascend/AscendNPU-IRdev/regbase
共 9 个文件变更+3991-48
Mbishengir/include/bishengir/Dialect/Triton/Transforms/Passes.h+4-0
@@ -20,6 +20,10 @@ namespace bishengir {
20 20 
21namespace triton {21namespace triton {
22 22 
23+/// Creates a pass that converts tensors with non power of 2 dimensions
24+/// into tensors with power of 2 dimensions
25+std::unique_ptr<mlir::Pass> createConvertNonPowerTwoTensorsPass();
26+ 
23/// Creates wrappers and attributes for SIMT functions27/// Creates wrappers and attributes for SIMT functions
24std::unique_ptr<mlir::Pass>28std::unique_ptr<mlir::Pass>
25createAdaptGPUKernelPass(TritonRemapOptions options = {});29createAdaptGPUKernelPass(TritonRemapOptions options = {});
Mbishengir/include/bishengir/Dialect/Triton/Transforms/Passes.td+11-0
@@ -20,6 +20,17 @@ def DecomposeFRem : Pass<"decompose-frem", "ModuleOp"> {
20 20 
21#ifdef BISHENGIR_ENABLE_TRITON_COMPILE21#ifdef BISHENGIR_ENABLE_TRITON_COMPILE
22 22 
23+def ConvertNonPowerTwoTensors
24+ : Pass<"convert-non-power-two-tensors", "mlir::triton::FuncOp"> {
25+ let summary = "Converts tensors with dimension sizes that are not powers of 2 into tensors with powers of 2";
26+ let constructor = "bishengir::triton::createConvertNonPowerTwoTensorsPass()";
27+ let dependentDialects = [
28+ "mlir::arith::ArithDialect",
29+ "mlir::triton::TritonDialect",
30+ "mlir::tensor::TensorDialect",
31+ ];
32+}
33+ 
23def SetBishengirSimtOptAttr34def SetBishengirSimtOptAttr
24 : Pass<"set-bishengir-simt-opt-attr", "mlir::ModuleOp"> {35 : Pass<"set-bishengir-simt-opt-attr", "mlir::ModuleOp"> {
25 let summary = "Set BishengIR SIMT optimization module attribute";36 let summary = "Set BishengIR SIMT optimization module attribute";
Mbishengir/lib/Dialect/Triton/Pipelines/TritonPipelines.cpp+1-0
@@ -101,6 +101,7 @@ void buildLowerTritonPipeline(OpPassManager &pm,
101 bishengir::SetBishengirSimtOptAttrOptions optionsSimtOpt;101 bishengir::SetBishengirSimtOptAttrOptions optionsSimtOpt;
102 optionsSimtOpt.enableBishengirSimtOptimization =102 optionsSimtOpt.enableBishengirSimtOptimization =
103 options.enableBishengirSimtOptimization;103 options.enableBishengirSimtOptimization;
104+ pm.addNestedPass<mlir::triton::FuncOp>(createConvertNonPowerTwoTensorsPass());
104 pm.addPass(105 pm.addPass(
105 bishengir::triton::createSetBishengirSimtOptAttrPass(optionsSimtOpt));106 bishengir::triton::createSetBishengirSimtOptAttrPass(optionsSimtOpt));
106 AdaptTritonIRKernelOptions adaptOpt;107 AdaptTritonIRKernelOptions adaptOpt;
Mbishengir/lib/Dialect/Triton/Transforms/CMakeLists.txt+1-0
@@ -3,6 +3,7 @@ if(BISHENGIR_ENABLE_TRITON_COMPILE)
3 list(APPEND BS_FEAT_SOURCES3 list(APPEND BS_FEAT_SOURCES
4 AdaptTritonIRKernel.cpp4 AdaptTritonIRKernel.cpp
5 AdaptGPUKernel.cpp5 AdaptGPUKernel.cpp
6+ ConvertNonPowerTwoTensors.cpp
6 DecomposeReduction.cpp7 DecomposeReduction.cpp
7 DumpFractalLayout.cpp8 DumpFractalLayout.cpp
8 FixFusedCatPass.cpp9 FixFusedCatPass.cpp
Abishengir/lib/Dialect/Triton/Transforms/ConvertNonPowerTwoTensors.cpp+2440-0
@@ -0,0 +1,2440 @@
1+//===- ConvertNonPowerTwoTensors.cpp -----------------------------*- C++-*-===//
yue-xy
yue-xyyue-xy7月23日

try rebase and use the latest llvm? I checked and other commits do not have the werror problems that leads to SC-FAIL

likedislike
2+//
3+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+// See https://llvm.org/LICENSE.txt for license information.
5+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+//
7+//===----------------------------------------------------------------------===//
8+// This pass converts tensors with non-power of 2 sizes to tensors of power of 2
9+// sizes so triton accepts them. This pass pads tensors until their dimensions
10+// are powers of 2
11+// Ignores operations that only act on power of two tensors/tensor pointers (as
12+// well as non tensor/tensor pointer operations)
13+//
14+// Important data structures:
15+// LocalTensorShapeData: stores the data shape (initial, actual shape of the
16+// tensor) and the virtual shape (the padded shape of the tensor)
17+// DenseMap<Value, LocalTensorShapeData>: maps operation tensor results/operands
18+// to their shape data (only for non power of 2 tensors) DenseMap<Value,
19+// PotentialPaddingRequirements>: For each non power of 2 tensor value, keeps
20+// track of:
21+// - The Value's immediate users that have a padding requirement (and which
22+// operands have which padding requirement)
23+// - The Value's downstream users that have a padding requirement (and which
24+// operands have which padding requirement)
25+// DenseMap<Value, SmallVector<std::pair<TypedAttr, SmallPtrSet<OpOperand*,
26+// 2>>>>:
27+// For each non power of 2 tensor value, keeps track of:
28+// - The Value's required paddings that have been chosen to be set at this
29+// value in the final codegen stage (TypedAttr)
30+// - The Operands that require the specific padding for each padding type
31+// (SmallPtrSet<OpOperand*, 2>)
32+//
33+// Overall Steps:
34+// 1: Scan for non power of 2 tensor ops that are not supported by this pass
35+// 2: For each tt.load/tt.store/tt.atomic_rmw op: if the pointer operand is of
36+// type tensor<shape x !tt.ptr<type>> then add a default mask filled with
37+// true, other filled with 0 (for load)
38+// Otherwise if ptr operand is of type !tt.ptr<tensor<shape x elemment type>,
39+// sets boundary check attr
40+// 3: When slice ops can go out of bounds due to expansion of the non power of 2
41+// dim, or when tensor.insert_slice ops that have a non power of 2 dim along
42+// the slice axis and expanding them would override previous existing data,
43+// split up the tensor.insert_slice/tensor.extract_slice op into power
44+// of 2 dimensions along the slice axis
45+// 4: Create a DenseMap<Value, LocalTensorShapeData> that tracks the data shape
46+// and virtual shape of values
47+// 5: Create a DenseMap<Value, PotentialPaddingRequirements> that tracks the
48+// padding requirement of its usage chain
49+// 6: Create a DenseMap<Value, SmallVector<std::pair<TypedAttr,
50+// SmallPtrSet<OpOperand *, 2>>>> (finalize padding choices) based on
51+// DenseMap<Value, PotentialPaddingRequirements>
52+// - Currently uses a very naive/basic approach, could be optimized
53+// 7: Replace non power of 2 tensors via preorder traversal (for nested ops,
54+// visit outside wrapping op before visiting the inside), top to bottom
55+// This guarantees we visit an op's sources before we visit the op itself
56+//
57+// Notable implementation details:
58+// triton::ReduceOp:
59+// - No change if the reduce axis was a power of 2
60+// - For each input tensor (or 'lane') tries to find an identity for the math
61+// calculation
62+// For example, if one tensor gets reduced by arith.addf, finds the identity
63+// 0.0
64+// - Limitations: Only supports reduction algorithms that have one
65+// operation (and have an identity), and must only access its own
66+// accumulator and value argument
67+// - If able to find an identity, requests that the tensor corresponding to
68+// that reduction algorithm is padded with the identity element
69+// - Otherwise, adds the mask tensor as an argument and uses it to mask out
70+// unwanted elements
71+//
72+// triton::ScanOp:
73+// - Very similar implementation to triton::ReduceOp, except:
74+// Only care about padding/masking if reverse=true currently
75+// - If reverse=false, then the elements at the end are padding and don't
76+// affect calculations
77+//
78+// triton::LoadOp/triton::StoreOp/triton::AtomicRMWOp:
79+// - For load/store op's that accept a tensor<sizex!tt.ptr<>> argument, we add
80+// an initial mask (all true)
81+// - Note that masks are NOT added if the load/store op accepts a
82+// !tt.ptr<tensor> arg
83+// - No other changes, these are treated as 'general tensor ops'
84+//
85+// triton::ReshapeOp:
86+// - If just updating the operand and result shape results in the data being
87+// in the wrong locations,
88+// instead reshapes to 1D and uses slice ops to move around data before
89+// reshaping to result shape
90+//
91+// tensor::ExtractSliceOp/tensor::InsertSliceOp:
92+// - Supports cases where only the tensors are only sliced along one dimension
93+// - Needed for tensor::InsertSliceOp correctness, RewriteSliceOpToTriton.cpp
94+// also has this restriction
95+// - Sometimes needs to split the insert_slice/extract_slice into multiple
96+// slice ops which have powers of two dimensions along the slice axis
97+// - slice ops are split when the offset is dynamic or the offset + expanded
98+// result axis dim > expanded source axis dim
99+// - insert_slice ops are also split when the slice axis of the insert
100+// tensor is not a power of 2, and expanding would override original data
101+//
102+// arith::ConstantOp:
103+// - Note that constant ops in the form `arith.constant dense<[1, 2, 3]>` for
104+// example with non power of two dimensions are not supported
105+// - This is because of a restriction that requires arith.constant dense ops
106+// declared this way to have num elements == num threads per warp
107+//
108+// triton::MakeTensorPtrOp:
109+// - Just updates the size of the output tensor
110+// - When this tensor is loaded, relies on MakeTensorPtrOp's lowering to
111+// provide masking/valid ptrs
112+//
113+// General Tensor Ops: (see isGeneralTensorOp(Operation* op) for which ops are
114+// categorized as a general tensor op)
115+// - Note that some ops that implement the InferTypeOpInterface are not
116+// supported yet
117+// - These are ops that implement the InferTypeOpInterface/have the
118+// elementwise trait and some manually added ones
119+// - For these ops, we just need to update some operand/return types
120+// - Ex: elementwise ops, triton::MakeTensorPtrOp, scf::ForOp, triton::LoadOp,
121+// triton::StoreOp, etc
122+//
123+// - Sometimes used as a cleanup step to finish modifying more complex ops
124+// (Ex: tt.reduce, tt.scan)
125+//
126+// Other notes:
127+// data shape refers to the shape of tensors before the pass
128+// virtual shape refers to the padded shape of tensors after the pass
129+// check isSupportedOp(Operation* op) to view which ops are
130+// supported/unsupported
131+//===----------------------------------------------------------------------===//
132+ 
133+#include "bishengir/Dialect/Triton/Transforms/Passes.h"
134+#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
135+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
136+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
137+#include "mlir/Analysis/DataFlowFramework.h"
138+#include "mlir/Dialect/Arith/IR/Arith.h"
139+#include "mlir/Dialect/SCF/IR/SCF.h"
140+#include "mlir/Dialect/Tensor/IR/Tensor.h"
141+#include "mlir/IR/BuiltinAttributeInterfaces.h"
142+#include "mlir/IR/BuiltinAttributes.h"
143+#include "mlir/IR/BuiltinTypes.h"
144+#include "mlir/IR/MLIRContext.h"
145+#include "mlir/IR/OpDefinition.h"
146+#include "mlir/IR/Operation.h"
147+#include "mlir/IR/PatternMatch.h"
148+#include "mlir/IR/SymbolTable.h"
149+#include "mlir/IR/Value.h"
150+#include "mlir/Interfaces/CallInterfaces.h"
151+#include "mlir/Interfaces/CastInterfaces.h"
152+#include "mlir/Interfaces/ControlFlowInterfaces.h"
153+#include "mlir/Interfaces/InferTypeOpInterface.h"
154+#include "mlir/Pass/Pass.h"
155+#include "mlir/Support/LLVM.h"
156+#include "mlir/Support/TypeID.h"
157+#include "mlir/Transforms/DialectConversion.h"
158+#include "triton/Dialect/Triton/IR/Dialect.h"
159+#include "triton/Dialect/Triton/IR/Types.h"
160+#include "llvm/ADT/ArrayRef.h"
161+#include "llvm/ADT/STLExtras.h"
162+#include "llvm/ADT/SmallPtrSet.h"
163+#include "llvm/ADT/SmallVector.h"
164+#include "llvm/ADT/TypeSwitch.h"
165+#include "llvm/Support/MathExtras.h"
166+#include <optional>
167+#include <queue>
168+ 
169+namespace bishengir::triton {
170+#define GEN_PASS_DEF_CONVERTNONPOWERTWOTENSORS
171+#include "bishengir/Dialect/Triton/Transforms/Passes.h.inc"
172+ 
173+namespace {
174+ 
175+using namespace mlir;
176+using namespace mlir::triton;
177+ 
178+// Given a data shape, calculate virtual shape and virtual shape size
179+std::pair<SmallVector<int64_t>, uint64_t>
180+expandShape(ArrayRef<int64_t> dataShape) {
181+ SmallVector<int64_t> dest;
182+ dest.reserve(dataShape.size());
183+ 
184+ uint64_t totalSize = 1;
185+ 
186+ for (int64_t dim : dataShape) {
187+ uint64_t expandedDim = llvm::PowerOf2Ceil(dim);
188+ totalSize *= expandedDim;
189+ 
190+ dest.push_back(expandedDim);
191+ }
192+ 
193+ return std::make_pair(std::move(dest), totalSize);
194+}
195+ 
196+// For a non power of two tensor Value, stores:
197+// non power of two size (data shape),
198+// padded size (localVirtualShape),
199+// num elements in padded tensor (localVirtualShapeSize)
200+struct LocalTensorShapeData {
201+ SmallVector<int64_t> localVirtualShape;
202+ uint64_t localVirtualShapeSize = 0;
203+ ArrayRef<int64_t> dataShape;
204+ 
205+ LocalTensorShapeData() = default;
206+ 
207+ explicit LocalTensorShapeData(ArrayRef<int64_t> dataShape)
208+ : dataShape(dataShape) {
209+ initializeShape(dataShape);
210+ }
211+ 
212+ void initializeShape(const ArrayRef<int64_t> initialDataShape) {
213+ localVirtualShape.clear();
214+ auto [shape, size] = expandShape(initialDataShape);
215+ localVirtualShape = std::move(shape);
216+ localVirtualShapeSize = size;
217+ }
218+};
219+ 
220+// Stores the padding requirements of a non power of two tensor Value and its
221+// chain of users
222+struct PotentialPaddingRequirements {
223+ // Counts the number of select ops that would be saved by having this op's
224+ // upstream sources padded
225+ DenseMap<TypedAttr, SmallPtrSet<OpOperand *, 2>> looseReqCounts;
226+ 
227+ // Lists strict padding requirements that MUST be handled at this Value, and
228+ // for which ops
229+ DenseMap<TypedAttr, SmallPtrSet<OpOperand *, 2>> strictReqs;
230+ 
231+ TypedAttr popMostFrequentLooseReq() {
232+ if (looseReqCounts.size() == 0) {
233+ return nullptr;
234+ }
235+ 
236+ TypedAttr mostFrequentReq;
237+ unsigned int mostFrequentReqUses = 0;
238+ 
239+ for (const auto &count : looseReqCounts) {
240+ if (count.second.size() > mostFrequentReqUses) {
241+ mostFrequentReq = count.first;
242+ mostFrequentReqUses = count.second.size();
243+ }
244+ }
245+ 
246+ removeLooseReq(mostFrequentReq);
247+ return mostFrequentReq;
248+ }
249+ 
250+ void removeLooseReq(TypedAttr req) { looseReqCounts.erase(req); }
251+ 
252+ void addLoose(TypedAttr padding, OpOperand *operand) {
253+ looseReqCounts[padding].insert(operand);
254+ }
255+ 
256+ void addStrict(TypedAttr padding, OpOperand *operand) {
257+ // Add to padding count
258+ addLoose(padding, operand);
259+ 
260+ strictReqs[padding].insert(operand);
261+ }
262+};
263+ 
264+// Stores reqs for a Value during padding backwards dataflow analysis
265+struct AllPaddingRequirements {
266+ SmallVector<TypedAttr, 2> reqs;
267+ 
268+ bool initialized = false;
269+ 
270+ ChangeResult setToDefault() {
271+ if (!initialized) {
272+ initialized = true;
273+ return ChangeResult::Change;
274+ }
275+ if (reqs.size() == 0) {
276+ return ChangeResult::NoChange;
277+ }
278+ 
279+ reqs.clear();
280+ return ChangeResult::Change;
281+ }
282+ 
283+ bool has(TypedAttr attr) {
284+ for (TypedAttr reqAttr : reqs) {
285+ if (reqAttr == attr) {
286+ return true;
287+ }
288+ }
289+ 
290+ return false;
291+ }
292+ 
293+ // returns ChangeResult::Change if a new attr was inserted,
294+ // ChangeResult::NoChange otherwise (attr already in list)
295+ ChangeResult add(TypedAttr attr) {
296+ if (has(attr)) {
297+ return ChangeResult::NoChange;
298+ }
299+ reqs.push_back(attr);
300+ return ChangeResult::Change;
301+ }
302+ 
303+ void print(raw_ostream &os) const {}
304+ 
305+ ChangeResult meet(const AllPaddingRequirements &other) {
306+ if (!other.initialized) {
307+ return ChangeResult::NoChange;
308+ }
309+ if (!initialized) {
310+ initialized = true;
311+ reqs = other.reqs;
312+ return ChangeResult::Change;
313+ }
314+ 
315+ ChangeResult changed = ChangeResult::NoChange;
316+ 
317+ for (TypedAttr attr : other.reqs) {
318+ if (add(attr) == ChangeResult::Change) {
319+ changed = ChangeResult::Change;
320+ }
321+ }
322+ 
323+ return changed;
324+ }
325+ 
326+ static AllPaddingRequirements join(const AllPaddingRequirements &lhs,
327+ const AllPaddingRequirements &rhs) {
328+ AllPaddingRequirements newReqs;
329+ newReqs.reqs = lhs.reqs;
330+ for (TypedAttr attr : rhs.reqs) {
331+ newReqs.reqs.push_back(attr);
332+ }
333+ newReqs.initialized = lhs.initialized || rhs.initialized;
334+ return newReqs;
335+ }
336+ 
337+ // does not get run
338+ static AllPaddingRequirements meet(const AllPaddingRequirements &lhs,
339+ const AllPaddingRequirements &rhs) {
340+ return AllPaddingRequirements();
341+ }
342+ 
343+ bool operator==(const AllPaddingRequirements &other) {
344+ if (!initialized) {
345+ return !other.initialized;
346+ }
347+ if (!other.initialized) {
348+ return false;
349+ }
350+ return reqs == other.reqs;
351+ }
352+};
353+ 
354+// Stores data about a tensor::InsertSliceOp or tensor::ExtractSliceOp
355+struct SliceOpData {
356+ int axis = -1;
357+ bool isOffsetStatic = false;
358+ int64_t offsetVal = -1;
359+ OpFoldResult offset;
360+ uint64_t largeAxisDim = 0;
361+ uint64_t smallAxisDim = 0;
362+ size_t rank;
363+ 
364+ SliceOpData() {}
365+ SliceOpData(Operation *op, ArrayRef<int64_t> smallDataShape,
366+ ArrayRef<int64_t> largeDataShape)
367+ : rank(largeDataShape.size()) {
368+ ArrayRef<int64_t> offsets;
369+ axis = -1;
370+ offsetVal = -1;
atomgit-bot
atomgit-botatomgit-bot7月4日

🟡 Medium Priority

ConvertNonPowerTwoTensors.cpp 第307-310行，SliceOpData(Operation*, ArrayRef<int64_t>, ArrayRef<int64_t>) 构造函数中包含 llvm::errs() 调试输出：

这些输出在每次构建 SliceOpData 时都会执行，包括：

  • splitSliceOps（第686行、第692行）
  • updateGeneralTensorOp（第1171行、第1177行）

影响：在生产环境中，每个包含 slice 操作的 kernel 编译都会向 stderr 输出大量调试信息，包括完整 op->dump() 的输出，干扰正常的诊断输出并降低编译性能。

建议：删除或使用 LLVM_DEBUG 宏包裹这些调试输出，使其仅在 debug 构建中生效。

改动建议
370
+ LLVM_DEBUG({
371
+ llvm::errs() << "initialize slice op data with op: ";
370
- offsetVal = -1;
372
+ op->dump();
373
+ llvm::errs() << "large rank: " << largeDataShape.size() << "\n";
374
+ llvm::errs() << "small rank: " << smallDataShape.size() << "\n";
375
+ });
应用建议
likedislike
不准确？
M
Max_Wang_Huawei
7月14日 评论：
371+ 
372+ for (size_t i = 0; i < rank; i++) {
373+ if (largeDataShape[i] != smallDataShape[i]) {
374+ axis = static_cast<int>(i);
375+ }
376+ }
377+ 
378+ if (axis == -1) {
379+ return;
380+ }
381+ 
382+ smallAxisDim = static_cast<uint64_t>(smallDataShape[axis]);
383+ largeAxisDim = static_cast<uint64_t>(largeDataShape[axis]);
384+ 
385+ if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
386+ offsets = extractSliceOp.getStaticOffsets();
387+ offset = extractSliceOp.getMixedOffsets()[axis];
388+ } else if (auto insertSliceOp = dyn_cast<tensor::InsertSliceOp>(op)) {
389+ offsets = insertSliceOp.getStaticOffsets();
390+ offset = insertSliceOp.getMixedOffsets()[axis];
391+ }
392+ 
393+ if (!ShapedType::isDynamic(offsets[axis])) {
394+ offsetVal = offsets[axis];
395+ isOffsetStatic = true;
396+ }
397+ }
398+ SliceOpData(int axis, int64_t offsetVal, OpFoldResult offset,
399+ int64_t largeAxisDim, int64_t smallAxisDim, size_t rank)
400+ : axis(axis), offsetVal(offsetVal), offset(offset),
401+ largeAxisDim(largeAxisDim), smallAxisDim(smallAxisDim), rank(rank) {
402+ if (offsetVal == -1) {
403+ isOffsetStatic = false;
404+ }
atomgit-bot
atomgit-botatomgit-bot7月4日

🟡 Medium Priority

SliceOpData 的第一个构造函数（第 305-340 行）在偏移为动态时无法正确设置 isOffsetStatic = false。该字段初始化为 true（第 297 行），且从未被改为 false。当偏移为动态时，offsetVal 保持为 -1，但 isOffsetStatic 仍为 true。

在 splitSliceOps（第 718 行）中，isOffsetStatic 用于决定 isInitialExtractOffsetStatic 和 isInitialInsertOffsetStatic。由于它们始终为 true，静态偏移路径会被执行，并使用 initialExtractOffsetVal + curOffset（即 -1 + curOffset）作为偏移量——这显然是错误的。

例如，对于一个动态偏移为 X、位于非 2 次幂轴上的 extract_slice：

修复方法：在构造函数中，当 ShapedType::isDynamic(offsets[axis]) 为 true 时，设置 isOffsetStatic = false。

改动建议
404
+ SliceOpData(Operation* op, ArrayRef<int64_t> smallDataShape, ArrayRef<int64_t> largeDataShape) {
405
+ ArrayRef<int64_t> offsets;
406
+ llvm::errs() << "initialize slice op data with op: ";
407
+ op->dump();
408
+ llvm::errs() << "large rank: " << largeDataShape.size() << "\n";
409
+ llvm::errs() << "small rank: " << smallDataShape.size() << "\n";
410
+
411
+ rank = largeDataShape.size();
412
+ axis = -1;
413
+ offsetVal = -1;
414
+
415
+ for (size_t i = 0; i < rank; i++) {
416
+ if (largeDataShape[i] != smallDataShape[i]) {
417
+ axis = i;
404
- }
418
+ }
419
+ }
420
+
421
+ if (axis == -1) {
422
+ return;
423
+ }
424
+
425
+ smallAxisDim = static_cast<uint64_t>(smallDataShape[axis]);
426
+ largeAxisDim = static_cast<uint64_t>(largeDataShape[axis]);
427
+
428
+ if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
429
+ offsets = extractSliceOp.getStaticOffsets();
430
+ offset = extractSliceOp.getMixedOffsets()[axis];
431
+ } else if (auto insertSliceOp = dyn_cast<tensor::InsertSliceOp>(op)) {
432
+ offsets = insertSliceOp.getStaticOffsets();
433
+ offset = insertSliceOp.getMixedOffsets()[axis];
434
+ }
435
+
436
+ if (!ShapedType::isDynamic(offsets[axis])) {
437
+ offsetVal = offsets[axis];
438
+ } else {
439
+ isOffsetStatic = false;
440
+ }
441
+ }
应用建议
likedislike
不准确？
M
Max_Wang_Huawei
7月14日 评论：
405+ }
406+};
407+ 
408+bool isTensorProducer(Operation *op) {
409+ return isa<triton::SplatOp, triton::MakeRangeOp, arith::ConstantOp>(op);
410+}
411+ 
412+bool isGenericTensorOp(Operation *op) {
413+ if (isa<triton::ReduceOp, triton::ScanOp>(op)) {
414+ return false;
415+ }
416+ 
417+ if (op->hasTrait<OpTrait::Elementwise>()) {
418+ return true;
419+ }
420+ 
421+ if (isa<triton::AdvanceOp, triton::AtomicRMWOp, triton::BroadcastOp,
422+ triton::MakeTensorPtrOp>(op)) {
423+ return true;
424+ }
425+ 
426+ if (isa<scf::ForOp, scf::IfOp>(op)) {
427+ return true;
428+ }
429+ 
430+ if (isa<tensor::ExtractSliceOp, tensor::InsertSliceOp>(op)) {
431+ return true;
432+ }
433+ 
434+ if (isa<InferTypeOpInterface>(op)) {
435+ return true;
436+ }
437+ 
438+ return false;
439+}
440+ 
441+bool hasTensorArgs(Operation *op) {
442+ for (Value operand : op->getOperands()) {
443+ if (isa<RankedTensorType>(operand.getType())) {
444+ return true;
445+ }
446+ }
447+ 
448+ return false;
449+}
450+ 
451+bool hasNonPowerTwoDim(ArrayRef<int64_t> shape) {
452+ for (int64_t dim : shape) {
453+ if (!llvm::isPowerOf2_64(dim)) {
454+ return true;
455+ }
456+ }
457+ 
458+ return false;
459+}
460+ 
461+RankedTensorType getNonPowerTwoTensorType(Type type) {
462+ RankedTensorType tensorType = dyn_cast<RankedTensorType>(type);
463+ if (tensorType && hasNonPowerTwoDim(tensorType.getShape())) {
464+ return tensorType;
465+ }
466+ return nullptr;
467+}
468+ 
469+RankedTensorType getNestedTensorType(Type type) {
470+ RankedTensorType tensorType = nullptr;
471+ if (auto pointerType = dyn_cast<triton::PointerType>(type)) {
472+ if (auto pointeeType =
473+ dyn_cast<RankedTensorType>(pointerType.getPointeeType())) {
474+ tensorType = pointeeType;
475+ }
476+ } else if (auto rankedTensorType = dyn_cast<RankedTensorType>(type)) {
477+ tensorType = rankedTensorType;
478+ }
479+ return tensorType;
480+}
481+ 
482+RankedTensorType getNonPowerTwoNestedTensorType(Type type) {
483+ RankedTensorType tensorType = getNestedTensorType(type);
484+ if (tensorType && hasNonPowerTwoDim(tensorType.getShape())) {
485+ return tensorType;
486+ }
487+ return nullptr;
488+}
489+ 
490+bool isNestedTensorType(Type type) {
491+ return getNestedTensorType(type) != nullptr;
492+}
493+ 
494+bool isNonPowerTwoTensorOrTensorPtr(Type type) {
495+ return getNonPowerTwoNestedTensorType(type) != nullptr;
496+}
497+ 
498+bool isTensorPointerLoadStoreOperation(Operation *op) {
499+ return isa<triton::LoadOp, triton::StoreOp, triton::AtomicRMWOp>(op) &&
500+ isa<triton::PointerType>(op->getOperandTypes()[0]);
501+}
502+ 
503+bool hasTensorRes(Operation *op) {
504+ for (Value operand : op->getResults()) {
505+ if (isNestedTensorType(operand.getType())) {
506+ return true;
507+ }
508+ }
509+ 
510+ return false;
511+}
512+ 
513+bool isNonPow2TensorOperation(Operation *op) {
514+ for (Value val : op->getOperands()) {
515+ if (isNonPowerTwoTensorOrTensorPtr(val.getType())) {
516+ return true;
517+ }
518+ }
519+ 
520+ for (Value val : op->getResults()) {
521+ if (isNonPowerTwoTensorOrTensorPtr(val.getType())) {
522+ return true;
523+ }
524+ }
525+ 
526+ if (auto scfForOp = dyn_cast<scf::ForOp>(op)) {
527+ for (BlockArgument regionArg : scfForOp.getRegionIterArgs()) {
528+ if (isNonPowerTwoTensorOrTensorPtr(regionArg.getType())) {
529+ return true;
530+ }
531+ }
532+ }
533+ 
534+ return false;
535+}
536+ 
537+Value createTensor(OpBuilder &builder, Location loc,
538+ RankedTensorType tensorType,
539+ Attribute paddingVal = nullptr) {
540+ Value tensor;
541+ Type elementType = tensorType.getElementType();
542+ if (auto ptrType = dyn_cast<triton::PointerType>(elementType)) {
543+ Value scalarZero =
544+ builder.create<arith::ConstantOp>(loc, builder.getI64IntegerAttr(0));
545+ Value tritonNullptr =
546+ builder.create<triton::IntToPtrOp>(loc, elementType, scalarZero);
547+ tensor = builder.create<triton::SplatOp>(loc, tensorType, tritonNullptr);
548+ } else {
549+ DenseElementsAttr constAttr;
550+ if (paddingVal) {
551+ constAttr = DenseElementsAttr::get(tensorType, paddingVal);
552+ } else {
553+ constAttr =
554+ DenseElementsAttr::get(tensorType, builder.getZeroAttr(elementType));
555+ }
556+ tensor = builder.create<arith::ConstantOp>(loc, constAttr);
557+ }
558+ 
559+ return tensor;
560+}
561+ 
562+// Given the data (unpadded) shape of a tensor and the size of its virtual
563+// (padded) shape, Returns a vector<bool> that is true where there is data when
564+// flattened, false otherwise
565+std::vector<bool> getFlattenedDataLocs(ArrayRef<int64_t> dataShape,
566+ uint64_t virtualShapeSize) {
567+ std::vector<bool> res;
568+ res.reserve(virtualShapeSize);
569+ 
570+ for (uint64_t i = 0; i < virtualShapeSize; i++) {
571+ bool inShape = true;
572+ uint64_t rem = i;
573+ for (int64_t r = static_cast<int64_t>(dataShape.size() - 1); r >= 0; r--) {
574+ uint64_t fullDim = llvm::PowerOf2Ceil(dataShape[r]);
575+ uint64_t coord = rem % fullDim;
576+ rem /= fullDim;
577+ 
578+ if (coord >= static_cast<uint64_t>(dataShape[r])) {
579+ inShape = false;
580+ break;
581+ }
582+ }
583+ res.push_back(inShape);
584+ }
585+ 
586+ return res;
587+}
588+ 
589+// If the triton::ReshapeOp is complex, returns the data locations of a
590+// flattened operand and result, otherwise return std::nullopt
591+std::optional<std::pair<std::vector<bool>, std::vector<bool>>>
592+getComplexReshapeDataLocs(const LocalTensorShapeData &operandData,
593+ const LocalTensorShapeData &resultData) {
594+ size_t operandIdx = 0;
595+ size_t resultIdx = 0;
596+ ArrayRef<int64_t> operandShape = operandData.localVirtualShape;
597+ ArrayRef<int64_t> resultShape = resultData.localVirtualShape;
598+ std::vector<bool> sourceDataLocs = getFlattenedDataLocs(
599+ operandData.dataShape, operandData.localVirtualShapeSize);
600+ std::vector<bool> resultDataLocs = getFlattenedDataLocs(
601+ resultData.dataShape, resultData.localVirtualShapeSize);
602+ 
603+ std::pair<std::vector<bool>, std::vector<bool>> result{sourceDataLocs,
604+ resultDataLocs};
605+ 
606+ // Checks the dimensions are equal while skipping over dimensions of size 1
607+ while (resultIdx < resultShape.size() && operandIdx < operandShape.size()) {
608+ while (operandIdx < operandShape.size() && operandShape[operandIdx] == 1) {
609+ operandIdx += 1;
610+ }
611+ 
612+ while (resultIdx < resultShape.size() && resultShape[resultIdx] == 1) {
613+ resultIdx += 1;
614+ }
615+ 
616+ if (operandIdx == operandShape.size() && resultIdx == resultShape.size()) {
617+ return std::nullopt;
618+ }
619+ 
620+ if (operandShape[operandIdx] != resultShape[resultIdx]) {
621+ return result;
622+ }
623+ 
624+ operandIdx += 1;
625+ resultIdx += 1;
626+ }
627+ 
628+ return result;
629+}
630+ 
631+bool isSupportedSliceOp(Operation *op) {
632+ ArrayRef<int64_t> largeShape;
633+ ArrayRef<int64_t> smallShape;
634+ ArrayRef<int64_t> offsets;
635+ ArrayRef<int64_t> strides;
636+ 
637+ if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
638+ largeShape = extractSliceOp.getSource().getType().getShape();
639+ smallShape = extractSliceOp.getResult().getType().getShape();
640+ offsets = extractSliceOp.getStaticOffsets();
641+ strides = extractSliceOp.getStaticStrides();
642+ } else if (auto insertSliceOp = dyn_cast<tensor::InsertSliceOp>(op)) {
643+ largeShape = insertSliceOp.getDest().getType().getShape();
644+ smallShape = insertSliceOp.getSource().getType().getShape();
645+ offsets = insertSliceOp.getStaticOffsets();
646+ strides = insertSliceOp.getStaticStrides();
647+ }
648+ 
649+ int rank = static_cast<int>(largeShape.size());
650+ int axis = -1;
651+ 
652+ if (largeShape.size() < 1 || largeShape.size() != smallShape.size()) {
653+ op->emitError("large and small tensors must have matching rank >= 1; "
654+ "got large ")
655+ << largeShape << " vs small " << smallShape;
656+ return false;
657+ }
658+ 
659+ for (int i = 0; i < rank; i++) {
660+ if (largeShape[i] != smallShape[i]) {
661+ if (axis != -1) {
662+ op->emitError("only single-axis slicing is supported; both axis ")
663+ << axis << " and axis " << i << " differ between large and small";
664+ return false;
665+ }
666+ axis = i;
667+ }
668+ }
669+ 
670+ for (int i = 0; i < rank; ++i) {
671+ if (strides[i] != 1) {
672+ op->emitError("strides must all be 1; got non-unit stride at axis ") << i;
673+ return false;
674+ }
675+ }
676+ 
677+ for (int i = 0; i < rank; ++i) {
678+ if (i != axis && offsets[i] != 0) {
679+ return false;
680+ }
681+ }
682+ 
683+ return true;
684+}
685+ 
686+// Determines is an op is supported by this pass
687+bool isSupportedOp(Operation *op) {
688+ if (!isNonPow2TensorOperation(op) || !hasTensorRes(op)) {
689+ return true;
690+ }
691+ 
692+ if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
693+ if (!cast<DenseElementsAttr>(constantOp.getValue()).isSplat()) {
694+ op->emitError("Non splat tensor constant ops must satisfy num elements "
695+ "== threads per warp");
696+ return false;
697+ }
698+ return true;
699+ }
700+ 
701+ if (isa<tensor::InsertOp, tensor::ExtractOp>(op)) {
702+ return true;
703+ }
704+ 
705+ if (isa<tensor::InsertSliceOp, tensor::ExtractSliceOp>(op)) {
706+ return isSupportedSliceOp(op);
707+ }
708+ 
709+ if (isa<scf::IfOp, scf::ForOp>(op)) {
710+ return true;
711+ }
712+ 
713+ if (isa<triton::AdvanceOp, triton::AtomicRMWOp, triton::BroadcastOp,
714+ triton::DotOp, triton::ExpandDimsOp, triton::GatherOp,
715+ triton::HistogramOp, triton::LoadOp, triton::MakeRangeOp,
716+ triton::MakeTensorPtrOp, triton::ReduceOp, triton::ReshapeOp,
717+ triton::ScanOp, triton::StoreOp, triton::SplatOp, triton::TransOp>(
718+ op)) {
719+ return true;
720+ }
721+ 
722+ if (isa<triton::SplitOp, triton::JoinOp>(op)) {
723+ return false;
724+ }
725+ 
726+ if (op->hasTrait<OpTrait::Elementwise>()) {
727+ return true;
728+ }
729+ 
730+ return false;
731+}
732+ 
733+// Makes sure all ops are supported by this pass
734+LogicalResult verifyOps(FuncOp &mod) {
735+ bool allOpsSupported = true;
736+ mod->walk([&](Operation *op) {
737+ if (!isSupportedOp(op)) {
738+ op->emitError(
739+ "Non power of two tensor operation is currently unsupported");
740+ allOpsSupported = false;
741+ }
742+ });
743+ 
744+ if (allOpsSupported) {
745+ return success();
746+ } else {
747+ return failure();
748+ }
749+}
750+ 
751+// Searches for triton::LoadOp, triton::StoreOp, and triton::AtmoicRMW ops that
752+// take in a tensor<!tt.ptr<>> for their pointer operand.
753+// For load/store ops that have !tt.ptr<tensor<>> pointers, sets boundary check
754+// attr instead
755+void addMasksToLoadAndStores(FuncOp &mod) {
atomgit-bot
atomgit-botatomgit-bot7月4日

🟡 Medium Priority

verifyOps 函数（第 595-603 行）始终返回 success()，即使在其中检测到不支持的操作并调用了 emitError。这意味着第 1771 行的 failed(verifyOps(module)) 检查永远不会为 true，即使有不支持的操作，pass 也会继续运行而不会提前中止。

虽然 emitError 仍会在操作上设置错误状态，但 pass 会继续执行后续步骤（如 addMasksToLoadAndStores、splitSliceOps 等），可能导致对不支持的 IR 进行错误转换。修复方法是在有任何操作发出错误时返回 failure()。

建议：在 walk lambda 中跟踪是否发生了任何错误，当检测到不支持的操作时返回 failure()。

改动建议
755
- void addMasksToLoadAndStores(FuncOp &mod) {
755
+ LogicalResult verifyOps(FuncOp &mod) {
756
+ bool hasError = false;
757
+ mod->walk([&](Operation* op) {
758
+ if (!isSupportedOp(op)) {
759
+ op->emitError("Operation is currently unsupported");
760
+ hasError = true;
761
+ }
762
+ });
763
+ if (hasError) {
764
+ return failure();
765
+ }
766
+ return success();
767
+ }
应用建议
likedislike
不准确？
M
Max_Wang_Huawei
7月14日 评论：
756+ std::queue<Operation *> workQueue;
757+ std::queue<std::pair<Operation *, RankedTensorType>> tensorPtrWorkQueue;
758+ 
759+ mod.walk([&](Operation *op) {
760+ Type ptrArgType = nullptr;
761+ if (auto loadOp = dyn_cast<triton::LoadOp>(op)) {
762+ if (loadOp.getMask()) {
763+ return;
764+ }
765+ ptrArgType = loadOp.getPtr().getType();
766+ } else if (auto storeOp = dyn_cast<triton::StoreOp>(op)) {
767+ if (storeOp.getMask()) {
768+ return;
769+ }
770+ ptrArgType = storeOp.getPtr().getType();
771+ } else if (auto atomicRmwOp = dyn_cast<triton::AtomicRMWOp>(op)) {
772+ if (atomicRmwOp.getMask()) {
773+ return;
774+ }
775+ ptrArgType = atomicRmwOp.getPtr().getType();
776+ }
777+ if (ptrArgType) {
778+ if (RankedTensorType type = getNonPowerTwoNestedTensorType(ptrArgType)) {
779+ if (isa<triton::PointerType>(ptrArgType)) {
780+ tensorPtrWorkQueue.emplace(op, type);
781+ } else {
782+ workQueue.push(op);
783+ }
784+ }
785+ }
786+ });
787+ 
788+ while (!workQueue.empty()) {
789+ Operation *cur = workQueue.front();
790+ workQueue.pop();
791+ IRRewriter builder(cur);
792+ Location loc = cur->getLoc();
793+ Type i1Type = builder.getI1Type();
794+ TypedAttr trueAttr = builder.getOneAttr(i1Type);
795+ 
796+ // create true mask
797+ RankedTensorType shapeType =
798+ cast<RankedTensorType>(cur->getOperand(0).getType());
799+ 
800+ RankedTensorType resType =
801+ RankedTensorType::get(shapeType.getShape(), i1Type);
802+ Value maskOp = createTensor(builder, loc, resType, trueAttr);
803+ 
804+ // Replace op
805+ if (auto loadOp = dyn_cast<triton::LoadOp>(cur)) {
806+ // Create other value mask
807+ Value ptr = loadOp.getPtr();
808+ auto newOp = builder.create<triton::LoadOp>(
809+ loc, ptr, maskOp, loadOp.getCache(), loadOp.getEvict(),
810+ loadOp.getIsVolatile());
811+ builder.replaceOp(cur, newOp);
812+ } else if (auto storeOp = dyn_cast<triton::StoreOp>(cur)) {
813+ Value ptr = storeOp.getPtr();
814+ Value val = storeOp.getValue();
815+ builder.create<triton::StoreOp>(loc, ptr, val, maskOp,
816+ storeOp.getBoundaryCheck(),
817+ storeOp.getCache(), storeOp.getEvict());
818+ storeOp->erase();
819+ } else if (auto atomicRmwOp = dyn_cast<triton::AtomicRMWOp>(cur)) {
820+ Type type = atomicRmwOp.getType();
821+ RMWOp modifyOp = atomicRmwOp.getAtomicRmwOp();
822+ MemSemantic semantic = atomicRmwOp.getSem();
823+ MemSyncScope scope = atomicRmwOp.getScope();
824+ Value ptr = atomicRmwOp.getPtr();
825+ Value val = atomicRmwOp.getVal();
826+ builder.create<triton::AtomicRMWOp>(loc, type, modifyOp, ptr, val, maskOp,
827+ semantic, scope);
828+ atomicRmwOp->erase();
829+ }
830+ }
831+ 
832+ // For tt.load/store ops that take a !tt.ptr<tensor>, only set boundary check
833+ // attrs instead of masking
834+ while (!tensorPtrWorkQueue.empty()) {
835+ auto [cur, type] = tensorPtrWorkQueue.front();
836+ tensorPtrWorkQueue.pop();
837+ IRRewriter builder(cur);
838+ 
839+ SmallVector<int32_t> boundDims;
840+ for (int64_t i = 0; i < type.getRank(); i++) {
841+ if (!llvm::isPowerOf2_64(i)) {
842+ boundDims.push_back(i);
843+ }
844+ }
845+ DenseI32ArrayAttr boundaryCheckAttr =
846+ DenseI32ArrayAttr::get(mod->getContext(), boundDims);
847+ 
848+ if (auto loadOp = dyn_cast<triton::LoadOp>(cur)) {
849+ loadOp.setBoundaryCheckAttr(boundaryCheckAttr);
850+ } else if (auto storeOp = dyn_cast<triton::StoreOp>(cur)) {
851+ storeOp.setBoundaryCheckAttr(boundaryCheckAttr);
852+ }
853+ }
854+}
855+ 
856+// In the case that either:
857+// - A tensor::InsertSliceOp has a non power of 2 dimension on the slice axis
858+// - Required to ensure correctness as otherwise extra data will be written
859+// - A tensor::ExtractSliceOp has a dynamic offset or the new slice size once
860+// padded + offset is larger than the source tensor axis size
861+// - Required to prevent out of bound errors
862+// Split the insert_slice/extract_slice op into power of two
863+// insert_slice/extract_slice op's (along the slice axis)
864+void splitSliceOps(FuncOp &mod) {
865+ std::queue<Operation *> workQueue;
866+ DenseMap<Operation *, SliceOpData> sliceData;
867+ 
868+ mod.walk([&](Operation *op) {
869+ if (!isNonPow2TensorOperation(op)) {
870+ return;
871+ }
872+ if (auto insertSliceOp = dyn_cast<tensor::InsertSliceOp>(op)) {
873+ SliceOpData data(op, insertSliceOp.getSourceType().getShape(),
874+ insertSliceOp.getDestType().getShape());
875+ // if not a no op and small axis dim size is not a power of two, and
876+ // offset is dynamic, override existing (non padding) data, or insert out
877+ // of bounds, split up into power of two slices
878+ if (data.axis != -1 && !llvm::isPowerOf2_64(data.smallAxisDim)) {
879+ if (!data.isOffsetStatic ||
880+ data.offsetVal + data.smallAxisDim < data.largeAxisDim ||
881+ data.offsetVal + llvm::PowerOf2Ceil(data.smallAxisDim) >
882+ llvm::PowerOf2Ceil(data.largeAxisDim)) {
883+ workQueue.push(op);
884+ sliceData[op] = data;
885+ }
886+ }
887+ } else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
888+ SliceOpData data(op, extractSliceOp.getResultType().getShape(),
889+ extractSliceOp.getSourceType().getShape());
890+ // if not a no op and offset dynamic and small axis dim size not power of
891+ // two
892+ // or offset static and would result in extracting out of bounds,
893+ // split up into power of two slices
894+ if (data.axis != -1 && !llvm::isPowerOf2_64(data.smallAxisDim)) {
895+ if (extractSliceOp.isDynamicOffset(data.axis)) {
896+ workQueue.push(op);
897+ sliceData[op] = data;
898+ } else {
899+ int64_t offset = extractSliceOp.getStaticOffset(data.axis);
900+ if (offset + llvm::PowerOf2Ceil(data.smallAxisDim) >
901+ llvm::PowerOf2Ceil(data.largeAxisDim)) {
902+ workQueue.push(op);
903+ sliceData[op] = data;
904+ }
905+ }
906+ }
907+ }
908+ });
909+ 
910+ while (!workQueue.empty()) {
911+ Operation *cur = workQueue.front();
912+ workQueue.pop();
913+ IRRewriter rewriter(cur);
914+ 
915+ Location loc = cur->getLoc();
916+ 
917+ const SliceOpData &sliceOpData = sliceData.at(cur);
918+ const size_t rank = sliceOpData.rank;
919+ const int axis = sliceOpData.axis;
920+ const bool isOffsetStatic = sliceOpData.isOffsetStatic;
921+ const int64_t offsetVal = sliceOpData.offsetVal;
922+ const OpFoldResult &offset = sliceOpData.offset;
923+ 
924+ // Need to split this tensor into smaller tensors along this axis
925+ SmallVector<int64_t> dimSizes;
926+ uint64_t rowsLeft = sliceOpData.smallAxisDim;
927+ 
928+ // Calculate the sizes of each slice
929+ while (rowsLeft > 0) {
930+ uint64_t nextSliceSize;
931+ if (llvm::isPowerOf2_64(rowsLeft)) {
932+ nextSliceSize = rowsLeft;
933+ } else {
934+ nextSliceSize = llvm::PowerOf2Ceil(rowsLeft) /
935+ 2; // largest power of two <= rowsLeft
936+ }
937+ dimSizes.push_back(nextSliceSize);
938+ rowsLeft -= nextSliceSize;
939+ }
940+ 
941+ SmallVector<Value> slices;
942+ slices.reserve(dimSizes.size());
943+ 
944+ int64_t curOffset = 0;
945+ 
946+ Value source;
947+ SmallVector<OpFoldResult> newOffsets(rank, rewriter.getIndexAttr(0));
948+ SmallVector<OpFoldResult> newSizes;
949+ SmallVector<OpFoldResult> newStrides(rank, rewriter.getIndexAttr(1));
950+ OpFoldResult initialExtractOffset;
951+ int64_t initialExtractOffsetVal;
952+ bool isInitialExtractOffsetStatic;
953+ 
954+ OpFoldResult initialInsertOffset;
955+ int64_t initialInsertOffsetVal;
956+ bool isInitialInsertOffsetStatic;
957+ 
958+ Value dest;
959+ 
960+ if (auto insertSliceOp = dyn_cast<tensor::InsertSliceOp>(cur)) {
961+ source = insertSliceOp.getSource();
962+ newSizes = insertSliceOp.getMixedSizes();
963+ initialExtractOffsetVal = 0;
964+ isInitialExtractOffsetStatic = true;
965+ 
966+ dest = insertSliceOp.getDest();
967+ initialInsertOffset = offset;
968+ initialInsertOffsetVal = offsetVal;
969+ isInitialInsertOffsetStatic = isOffsetStatic;
970+ } else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(cur)) {
971+ source = extractSliceOp.getSource();
972+ newSizes = extractSliceOp.getMixedSizes();
973+ initialExtractOffset = offset;
974+ initialExtractOffsetVal = offsetVal;
975+ isInitialExtractOffsetStatic = isOffsetStatic;
976+ 
977+ RankedTensorType resultType =
978+ cast<RankedTensorType>(cur->getResult(0).getType());
979+ dest = createTensor(rewriter, loc, resultType);
980+ 
981+ initialInsertOffsetVal = 0;
982+ isInitialInsertOffsetStatic = true;
983+ }
984+ 
985+ // Interleave extract slice -> insert slice pattern to reduce register
986+ // pressure
987+ for (int64_t sliceSize : dimSizes) {
988+ newSizes[axis] = rewriter.getIndexAttr(sliceSize);
989+ 
990+ // Extract slice
991+ if (isInitialExtractOffsetStatic) {
992+ newOffsets[axis] =
993+ rewriter.getIndexAttr(initialExtractOffsetVal + curOffset);
994+ } else {
995+ Value curConstOffset = rewriter.create<arith::ConstantOp>(
996+ loc, rewriter.getIndexAttr(curOffset));
997+ Value newOffset = rewriter.create<arith::AddIOp>(
998+ loc, curConstOffset, initialExtractOffset.get<Value>());
999+ newOffsets[axis] = newOffset;
1000+ }
atomgit-bot
atomgit-botatomgit-bot7月4日

🟠 High Priority

在 ConvertNonPowerTwoTensors.cpp 的 splitSliceOps 函数中，当将非2的幂次维度的 insert_slice/extract_slice 拆分为多个2的幂次切片时，动态偏移路径在两个位置错误地使用了 sliceSize（当前切片的大小）而非 curOffset（累积偏移量）。

问题位置1（第790行，extract 路径）： 本应为 rewriter.getIndexAttr(curOffset)。

问题位置2（第800行，insert 路径）： 同上。

触发条件：当 tensor::ExtractSliceOp 具有动态偏移（!isOffsetStatic）且非2的幂次 slice 轴维度需要拆分时（例如将 size=5 的切片拆分为 [4, 1]），或者 tensor::InsertSliceOp 具有非2的幂次 slice 轴且动态偏移时。

失败模式：以将 size=5 拆分为 [4, 1] 为例：

导致子切片被放置在错误的位置，引发数据损坏或越界访问。

建议：将 rewriter.getIndexAttr(sliceSize) 改为 rewriter.getIndexAttr(curOffset)。同时建议将变量名从 constOffset 改为 curOffsetConst 以明确语义。第800行需做相同修改。

改动建议
1000
- }
1000
+ Value curOffsetConst = rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(curOffset));
1001
+ Value offsetUsed = rewriter.create<arith::AddIOp>(loc, curOffsetConst, initialExtractOffset.get<Value>());
应用建议
likedislike
不准确？
M
Max_Wang_Huawei
7月14日 评论：
1001+ Value extractedSlice = rewriter.create<tensor::ExtractSliceOp>(
1002+ loc, source, newOffsets, newSizes, newStrides);
1003+ 
1004+ // Insert slice
1005+ if (isInitialInsertOffsetStatic) {
1006+ newOffsets[axis] =
1007+ rewriter.getIndexAttr(initialInsertOffsetVal + curOffset);
1008+ } else {
1009+ Value curConstOffset = rewriter.create<arith::ConstantOp>(
1010+ loc, rewriter.getIndexAttr(curOffset));
1011+ Value newOffset = rewriter.create<arith::AddIOp>(
1012+ loc, curConstOffset, initialInsertOffset.get<Value>());
1013+ newOffsets[axis] = newOffset;
1014+ }
1015+ 
1016+ dest = rewriter.create<tensor::InsertSliceOp>(
1017+ loc, extractedSlice, dest, newOffsets, newSizes, newStrides);
1018+ 
1019+ curOffset += sliceSize;
1020+ }
1021+ 
1022+ rewriter.replaceOp(cur, dest);
1023+ }
1024+}
1025+ 
1026+// Traverses through all ops, checks their operands and results for non power of
1027+// two tensors For each non power of two tensor, calculates its virtual shape
1028+DenseMap<Value, LocalTensorShapeData> populateLocalShapeData(FuncOp &mod) {
1029+ DenseMap<Value, LocalTensorShapeData> shapeMap;
1030+ 
1031+ mod.walk([&](Operation *op) {
1032+ bool calculateAll = false;
1033+ if (isa<tensor::InsertSliceOp, tensor::ExtractSliceOp, triton::HistogramOp>(
1034+ op)) {
1035+ calculateAll = true;
1036+ }
1037+ 
1038+ for (Value val : op->getOperands()) {
1039+ if (RankedTensorType type = getNestedTensorType(val.getType())) {
1040+ if (calculateAll || hasNonPowerTwoDim(type.getShape())) {
1041+ shapeMap[val] = LocalTensorShapeData(type.getShape());
1042+ }
1043+ }
1044+ }
1045+ 
1046+ for (Value val : op->getResults()) {
1047+ if (RankedTensorType type = getNestedTensorType(val.getType())) {
1048+ if (calculateAll || hasNonPowerTwoDim(type.getShape())) {
1049+ shapeMap[val] = LocalTensorShapeData(type.getShape());
1050+ }
1051+ }
1052+ }
1053+ });
1054+ 
1055+ return shapeMap;
1056+}
1057+ 
1058+// Looks at one tensor input in a triton::ReduceOp/triton::ScanOp and its
1059+// corresponding arguments in the reduce block Tries to examine its block to see
1060+// if it has an identity Can find an identity for cases where we have just one
1061+// op acting on the accumulator and next element (adding, multiplying, etc)
1062+// return nullptr if an identity was not found
1063+TypedAttr getReduceOrScanLaneIdentity(Operation *op, uint32_t laneIdx,
1064+ uint32_t numLanes) {
1065+ BlockArgument acc;
1066+ BlockArgument next;
1067+ Value laneRes;
1068+ TypeSwitch<Operation *, void>(op)
1069+ .Case<triton::ReduceOp>([&](triton::ReduceOp reduceOp) {
1070+ Block &block = reduceOp.getCombineOp().front();
1071+ acc = block.getArgument(laneIdx);
1072+ next = block.getArgument(laneIdx + numLanes);
1073+ 
1074+ auto returnOp = cast<triton::ReduceReturnOp>(block.getTerminator());
1075+ laneRes = returnOp->getOperand(laneIdx);
1076+ })
1077+ .Case<triton::ScanOp>([&](triton::ScanOp scanOp) {
1078+ Block &block = scanOp.getCombineOp().front();
1079+ acc = block.getArgument(laneIdx);
1080+ next = block.getArgument(laneIdx + numLanes);
1081+ 
1082+ auto returnOp = cast<triton::ScanReturnOp>(block.getTerminator());
1083+ laneRes = returnOp->getOperand(laneIdx);
1084+ })
1085+ .Default([](Operation *) {
1086+ llvm_unreachable(
1087+ "Non tt.reduce/tt.scan op passed to function "
1088+ "getReduceOrScanLaneIdentity in ConvertNonPowerTwoTensors.cpp");
1089+ });
1090+ 
1091+ Operation *foldOp = laneRes.getDefiningOp();
1092+ if ((!foldOp) || foldOp->getNumOperands() != 2) {
1093+ return nullptr;
1094+ }
1095+ 
1096+ bool withinLane =
1097+ (foldOp->getOperand(0) == acc && foldOp->getOperand(1) == next) ||
1098+ (foldOp->getOperand(1) == acc && foldOp->getOperand(0) == next);
1099+ if ((!withinLane) || (!isa<BlockArgument>(foldOp->getOperand(0))) ||
1100+ (!isa<BlockArgument>(foldOp->getOperand(1)))) {
1101+ // using other vars, or we didnt return `arg0 [op] arg1`, abort!
1102+ return nullptr;
1103+ }
1104+ 
1105+ std::optional<TypedAttr> potentialIdentity = arith::getNeutralElement(foldOp);
1106+ if (potentialIdentity) {
1107+ TypedAttr identity = *potentialIdentity;
1108+ if (auto floatIdentity = dyn_cast<FloatAttr>(identity)) {
1109+ if (floatIdentity.getValue().isNaN()) {
1110+ // Replace NAN identities with positive/negative INF to prevent NAN
1111+ // propagation
1112+ OpBuilder builder(op);
1113+ FloatType elementType = cast<FloatType>(floatIdentity.getType());
1114+ if (floatIdentity.getValue().isNegative()) {
1115+ return builder.getFloatAttr(
1116+ elementType, APFloat::getInf(elementType.getFloatSemantics(),
1117+ /*Negative=*/true));
1118+ } else {
1119+ return builder.getFloatAttr(
1120+ elementType, APFloat::getInf(elementType.getFloatSemantics(),
1121+ /*Negative=*/false));
1122+ }
1123+ }
1124+ }
1125+ 
1126+ return identity;
1127+ }
1128+ return nullptr;
1129+}
1130+ 
1131+// Returns a SmallVector<TypedAttr> containing the identity for the algorithm
1132+// associated with that tensor First tensor's identity is the first entry, etc.
1133+// In the return result, nullptr entries indicate no identity found
1134+SmallVector<TypedAttr> getReduceOrScanOpIdentities(Operation *op) {
1135+ unsigned int numOperands = op->getNumOperands();
1136+ 
1137+ SmallVector<TypedAttr> cur;
1138+ cur.reserve(numOperands);
1139+ 
1140+ for (unsigned int i = 0; i < numOperands; i++) {
1141+ cur.push_back(getReduceOrScanLaneIdentity(op, i, op->getNumOperands()));
1142+ }
1143+ 
1144+ return cur;
1145+}
1146+ 
1147+void expandToSize(SmallVector<TypedAttr> &vec, size_t size) {
1148+ while (vec.size() < size) {
1149+ vec.push_back(nullptr);
1150+ }
1151+}
1152+ 
1153+// Returns a SmallVector of TypedAttr padding requirements, one for each operand
1154+// (non tensor operands will have a nullptr TypedAttr) Operations that do not
1155+// require any padding will just return an empty SmallVector
1156+SmallVector<TypedAttr>
1157+getPaddingRequirements(OpBuilder &builder, Operation *op,
1158+ const DenseMap<Value, LocalTensorShapeData> &shapeData) {
1159+ SmallVector<TypedAttr> res(op->getNumOperands(), nullptr);
1160+ TypeSwitch<Operation *>(op)
1161+ .Case<triton::ReduceOp>([&](triton::ReduceOp reduceOp) {
1162+ uint32_t reduceAxis = reduceOp.getAxis();
1163+ if (llvm::isPowerOf2_64(shapeData.find(reduceOp->getOperand(0))
1164+ ->getSecond()
1165+ .dataShape[reduceAxis])) {
1166+ return;
1167+ }
1168+ SmallVector<TypedAttr> paddingAttrs =
1169+ getReduceOrScanOpIdentities(reduceOp);
1170+ 
1171+ if (paddingAttrs.size() != 0) {
1172+ for (size_t i = 0; i < paddingAttrs.size(); i++) {
1173+ if (paddingAttrs[i]) {
1174+ res[i] = paddingAttrs[i];
1175+ }
1176+ }
1177+ }
1178+ })
1179+ .Case<triton::ScanOp>([&](triton::ScanOp scanOp) {
1180+ if (!scanOp.getReverse()) {
1181+ // No padding needed on forward scan op as padding is in the backmost
1182+ // positions
1183+ // TODO - If padding locations are changed, this needs to be removed
1184+ res.clear();
1185+ return;
1186+ }
1187+ uint32_t reduceAxis = scanOp.getAxis();
1188+ if (llvm::isPowerOf2_64(shapeData.find(scanOp->getOperand(0))
1189+ ->getSecond()
1190+ .dataShape[reduceAxis])) {
1191+ return;
1192+ }
1193+ SmallVector<TypedAttr> paddingAttrs =
1194+ getReduceOrScanOpIdentities(scanOp);
1195+ 
1196+ if (paddingAttrs.size() != 0) {
1197+ for (size_t i = 0; i < paddingAttrs.size(); i++) {
1198+ if (paddingAttrs[i]) {
1199+ res[i] = paddingAttrs[i];
1200+ }
1201+ }
1202+ }
1203+ })
1204+ .Case<triton::LoadOp>([&](triton::LoadOp loadOp) {
1205+ // tt.load mask tensor should be false padded
1206+ expandToSize(res, 3);
1207+ res[1] = builder.getZeroAttr(builder.getI1Type());
1208+ })
1209+ .Case<triton::StoreOp>([&](triton::StoreOp storeOp) {
1210+ // tt.store mask tensor should be false padded
1211+ expandToSize(res, 3);
1212+ res[2] = builder.getZeroAttr(builder.getI1Type());
1213+ })
1214+ .Case<triton::AtomicRMWOp>([&](triton::AtomicRMWOp storeOp) {
1215+ // tt.store mask tensor should be false padded
1216+ expandToSize(res, 3);
1217+ res[2] = builder.getZeroAttr(builder.getI1Type());
1218+ })
1219+ .Case<triton::DotOp>([&](triton::DotOp dotOp) {
1220+ // tt.dot should be zero padded
1221+ TypedAttr paddingAttr = builder.getZeroAttr(
1222+ cast<RankedTensorType>(dotOp.getOperand(0).getType())
1223+ .getElementType());
1224+ for (size_t i = 0; i < res.size() - 1; i++) {
1225+ res[i] = paddingAttr;
1226+ }
1227+ })
1228+ .Default([&](Operation *op) {
1229+ // No padding needed, empty vector
1230+ res.clear();
1231+ });
1232+ 
1233+ return res;
1234+}
1235+ 
1236+// Assumes builder's insertion point is already set
1237+// Given the data shape and virtual shape of an op's result, creates a mask
1238+// which is true where the original elements are in the padded shape Is used
1239+// with arith::SelectOp to add/swap padding
1240+Value createPaddingMask(OpBuilder &builder, Location loc,
1241+ ArrayRef<int64_t> dataShape,
1242+ ArrayRef<int64_t> paddedShape) {
1243+ Type i32Type = builder.getI32Type();
1244+ Type i1Type = builder.getI1Type();
1245+ RankedTensorType largeType = RankedTensorType::get(paddedShape, i1Type);
1246+ 
1247+ size_t rank = dataShape.size();
1248+ SmallVector<Value> dimMasks;
1249+ dimMasks.reserve(rank);
1250+ 
1251+ for (size_t i = 0; i < rank; i++) {
1252+ SmallVector<int64_t> curShape = {paddedShape[i]};
1253+ RankedTensorType curType = RankedTensorType::get(curShape, i32Type);
1254+ Value cur =
1255+ builder.create<triton::MakeRangeOp>(loc, curType, 0, paddedShape[i]);
1256+ Value constant = builder.create<arith::ConstantOp>(
1257+ loc, i32Type, builder.getI32IntegerAttr(dataShape[i]));
1258+ Value splatted = builder.create<triton::SplatOp>(loc, curType, constant);
1259+ 
1260+ cur = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, cur,
1261+ splatted);
1262+ 
1263+ // We now expand out tensor to be tensor<1x1x...xNx1x1x...x1xT>
1264+ bool insertAtEnd = false;
1265+ for (size_t j = 0; j < rank; j++) {
1266+ if (j == i) {
1267+ insertAtEnd = true;
1268+ continue;
1269+ }
1270+ 
1271+ if (insertAtEnd) {
1272+ curShape.push_back(1);
1273+ curType = RankedTensorType::get(curShape, i1Type);
1274+ cur = builder.create<triton::ExpandDimsOp>(loc, curType, cur,
1275+ curShape.size() - 1);
1276+ } else {
1277+ curShape.insert(curShape.begin(), 1);
1278+ curType = RankedTensorType::get(curShape, i1Type);
1279+ cur = builder.create<triton::ExpandDimsOp>(loc, curType, cur, 0);
1280+ }
1281+ }
1282+ 
1283+ // Just neeed to broadcast now
1284+ cur = builder.create<triton::BroadcastOp>(loc, largeType, cur);
1285+ dimMasks.push_back(cur);
1286+ }
1287+ 
1288+ Value res = dimMasks[0];
1289+ for (size_t i = 1; i < rank; i++) {
1290+ res = builder.create<arith::AndIOp>(loc, res, dimMasks[i]);
1291+ }
1292+ return res;
1293+}
1294+ 
1295+// Assumes sets builder's insertion point before creating mask
1296+// Given the data shape and virtual shape of an op's result, creates a mask
1297+// which is true where the original elements are in the padded shape Is used
1298+// with arith::SelectOp to add/swap padding
1299+Value createPaddingMask(OpBuilder &builder, Operation *op,
1300+ ArrayRef<int64_t> dataShape,
1301+ ArrayRef<int64_t> paddedShape,
1302+ bool setInsertionPointAfter = false) {
1303+ if (setInsertionPointAfter) {
1304+ builder.setInsertionPointAfter(op);
1305+ } else {
1306+ builder.setInsertionPoint(op);
1307+ }
1308+ Location loc = op->getLoc();
1309+ 
1310+ return createPaddingMask(builder, loc, dataShape, paddedShape);
1311+}
1312+ 
1313+// Updates the return shape and iter_args of a scf ForOp
1314+void updateForOp(IRRewriter &rewriter, scf::ForOp op,
1315+ const DenseMap<Value, LocalTensorShapeData> &shapeMap) {
1316+ ValueRange initialVals = op.getInitArgs();
1317+ rewriter.modifyOpInPlace(op, [&]() {
1318+ Block &block = op.getRegion().front();
1319+ for (size_t i = 1; i < block.getNumArguments(); i++) {
1320+ block.getArgument(i).setType(initialVals[i - 1].getType());
1321+ }
1322+ 
1323+ auto yieldOp = cast<scf::YieldOp>(op.getBody()->getTerminator());
1324+ 
1325+ for (size_t i = 0; i < op.getNumResults(); i++) {
1326+ Value res = op.getResult(i);
1327+ if (RankedTensorType type =
1328+ getNonPowerTwoNestedTensorType(res.getType())) {
1329+ Type elementType = type.getElementType();
1330+ Value yieldSrc = yieldOp.getOperand(i);
1331+ ArrayRef<int64_t> newShape = shapeMap.at(yieldSrc).localVirtualShape;
1332+ 
1333+ RankedTensorType newTensorType =
1334+ RankedTensorType::get(newShape, elementType);
1335+ 
1336+ if (auto ptrType = dyn_cast<triton::PointerType>(res.getType())) {
1337+ triton::PointerType newPointerType = triton::PointerType::get(
1338+ newTensorType, ptrType.getAddressSpace());
1339+ res.setType(newPointerType);
1340+ } else {
1341+ res.setType(newTensorType);
1342+ }
1343+ }
1344+ }
1345+ });
1346+ return;
1347+}
1348+ 
1349+// Updates the return shape of a scf IfOp
1350+void updateIfOp(IRRewriter &rewriter, scf::IfOp op,
1351+ const DenseMap<Value, LocalTensorShapeData> &shapeMap) {
1352+ rewriter.modifyOpInPlace(op, [&]() {
1353+ auto yieldOp = cast<scf::YieldOp>(op.getBody()->getTerminator());
1354+ 
1355+ for (size_t i = 0; i < op.getNumResults(); i++) {
1356+ Value res = op.getResult(i);
1357+ if (RankedTensorType type = getNonPowerTwoTensorType(res.getType())) {
1358+ Type elementType = type.getElementType();
1359+ Value yieldSrc = yieldOp.getOperand(i);
1360+ ArrayRef<int64_t> newShape = shapeMap.at(yieldSrc).localVirtualShape;
1361+ 
1362+ RankedTensorType newTensorType =
1363+ RankedTensorType::get(newShape, elementType);
1364+ 
1365+ if (auto ptrType = dyn_cast<triton::PointerType>(res.getType())) {
1366+ triton::PointerType newPointerType = triton::PointerType::get(
1367+ newTensorType, ptrType.getAddressSpace());
1368+ res.setType(newPointerType);
1369+ } else {
1370+ res.setType(newTensorType);
1371+ }
1372+ }
1373+ }
1374+ });
1375+ return;
1376+}
1377+ 
1378+// Updates the return shape and modifies sizes data of a tensor::ExtractSliceOp
1379+void updateExtractSliceOp(IRRewriter &rewriter, tensor::ExtractSliceOp op,
1380+ const SliceOpData &sliceData,
1381+ const LocalTensorShapeData &result) {
1382+ Type elementType = op.getResultType().getElementType();
1383+ RankedTensorType resultType =
1384+ RankedTensorType::get(result.localVirtualShape, elementType);
1385+ rewriter.modifyOpInPlace(op, [&]() {
1386+ op.setStaticSizes(result.localVirtualShape);
1387+ op.getResult().setType(resultType);
1388+ });
1389+}
1390+ 
1391+// Updates the return shape and modifies sizes data of a tensor::InsertSliceOp
1392+void updateInsertSliceOp(IRRewriter &rewriter, tensor::InsertSliceOp op,
1393+ const SliceOpData &sliceData,
1394+ const LocalTensorShapeData &source,
1395+ const LocalTensorShapeData &result) {
1396+ Type elementType = op.getDestType().getElementType();
1397+ RankedTensorType destType =
1398+ RankedTensorType::get(result.localVirtualShape, elementType);
1399+ rewriter.modifyOpInPlace(op, [&]() {
1400+ op.setStaticSizes(source.localVirtualShape);
1401+ op.getResult().setType(destType);
1402+ });
1403+}
1404+ 
1405+// Updates the return type of a triton::AdvanceOp to be the same as the input
1406+// tensor ptr
1407+void updateAdvanceOp(IRRewriter &rewriter, triton::AdvanceOp op) {
1408+ rewriter.modifyOpInPlace(
1409+ op, [&]() { op.getResult().setType(op.getPtr().getType()); });
1410+}
1411+ 
1412+// Updates the return shape of a triton BroadcastOp
1413+void updateBroadcastOp(IRRewriter &rewriter, triton::BroadcastOp op,
1414+ const LocalTensorShapeData &data) {
1415+ Type elementType =
1416+ cast<RankedTensorType>(op.getResult().getType()).getElementType();
1417+ RankedTensorType newType =
1418+ RankedTensorType::get(data.localVirtualShape, elementType);
1419+ 
1420+ rewriter.modifyOpInPlace(op, [&]() { op.getResult().setType(newType); });
1421+}
1422+ 
1423+// Updates the return shape of a triton MakeTensorPtrOp
1424+void updateMakeTensorPtrOp(IRRewriter &rewriter, triton::MakeTensorPtrOp op,
1425+ const LocalTensorShapeData &data) {
1426+ triton::PointerType oldPtrType = cast<triton::PointerType>(op.getType());
1427+ RankedTensorType oldPtrTensorType =
1428+ cast<RankedTensorType>(oldPtrType.getPointeeType());
1429+ Type elementType = oldPtrTensorType.getElementType();
1430+ 
1431+ RankedTensorType newPtrTensorType =
1432+ RankedTensorType::get(data.localVirtualShape, elementType);
1433+ triton::PointerType newPtrType =
1434+ triton::PointerType::get(newPtrTensorType, oldPtrType.getAddressSpace());
1435+ 
1436+ rewriter.modifyOpInPlace(op, [&]() { op.getResult().setType(newPtrType); });
1437+}
1438+ 
1439+// Updates the return shape of an InferTypeOpInterface op
1440+void updateInferTypeOp(IRRewriter &rewriter, InferTypeOpInterface op) {
1441+ SmallVector<Type> inferredRetTypes;
1442+ inferredRetTypes.reserve(op->getNumResults());
1443+ 
1444+ if (succeeded(op.inferReturnTypes(op->getContext(), op->getLoc(),
1445+ op->getOperands(), op->getAttrDictionary(),
1446+ op->getPropertiesStorage(),
1447+ op->getRegions(), inferredRetTypes))) {
1448+ 
1449+ rewriter.modifyOpInPlace(op, [&]() {
1450+ for (size_t i = 0; i < op->getNumResults(); i++) {
1451+ op->getResult(i).setType(inferredRetTypes[i]);
1452+ }
1453+ });
1454+ }
1455+}
1456+ 
1457+// Updates the return shapes of elementwise ops
1458+void updateElementwiseOp(IRRewriter &rewriter, Operation *op) {
1459+ ArrayRef<int64_t> resShape;
1460+ for (Value operand : op->getOperands()) {
1461+ if (auto operandType = dyn_cast<RankedTensorType>(operand.getType())) {
1462+ resShape = operandType.getShape();
1463+ break;
1464+ }
1465+ }
1466+ 
1467+ rewriter.modifyOpInPlace(op, [&]() {
1468+ for (Value result : op->getResults()) {
1469+ ShapedType prevType = cast<ShapedType>(result.getType());
1470+ ShapedType newType =
1471+ prevType.cloneWith(resShape, prevType.getElementType());
1472+ result.setType(newType);
1473+ }
1474+ });
1475+}
1476+ 
1477+// Used for ops whose only change is updating the return type
1478+// Mostly ops that implement InferTypeOpInterface, but also some manually added
1479+// ops function isGenericTensorOp is used to determine what is a 'general tensor
1480+// op'
1481+void updateGeneralTensorOp(
1482+ IRRewriter &rewriter, Operation *op,
1483+ const DenseMap<Value, LocalTensorShapeData> &shapeMap) {
1484+ TypeSwitch<Operation *, void>(op)
1485+ .Case<scf::ForOp>(
1486+ [&](scf::ForOp forOp) { updateForOp(rewriter, forOp, shapeMap); })
1487+ .Case<scf::IfOp>(
1488+ [&](scf::IfOp op) { updateIfOp(rewriter, op, shapeMap); })
1489+ .Case<tensor::ExtractSliceOp>([&](tensor::ExtractSliceOp extractSliceOp) {
1490+ const LocalTensorShapeData &source =
1491+ shapeMap.at(extractSliceOp.getSource());
1492+ const LocalTensorShapeData &result =
1493+ shapeMap.at(extractSliceOp.getResult());
1494+ SliceOpData sliceData(op, result.dataShape, source.dataShape);
1495+ updateExtractSliceOp(rewriter, extractSliceOp, sliceData, result);
1496+ })
1497+ .Case<tensor::InsertSliceOp>([&](tensor::InsertSliceOp insertSliceOp) {
1498+ const LocalTensorShapeData &source =
1499+ shapeMap.at(insertSliceOp.getSource());
1500+ const LocalTensorShapeData &result =
1501+ shapeMap.at(insertSliceOp.getResult());
1502+ SliceOpData sliceData(op, source.dataShape, result.dataShape);
1503+ updateInsertSliceOp(rewriter, insertSliceOp, sliceData, source, result);
1504+ })
1505+ .Case<triton::AdvanceOp>([&](triton::AdvanceOp advanceOp) {
1506+ updateAdvanceOp(rewriter, advanceOp);
1507+ })
1508+ .Case<triton::AtomicRMWOp>([&](triton::AtomicRMWOp atomicRMWOp) {
1509+ updateElementwiseOp(rewriter, op);
1510+ })
1511+ .Case<triton::BroadcastOp>([&](triton::BroadcastOp broadcastOp) {
1512+ const LocalTensorShapeData &data = shapeMap.at(broadcastOp.getResult());
1513+ updateBroadcastOp(rewriter, broadcastOp, data);
1514+ })
1515+ .Case<triton::MakeTensorPtrOp>(
1516+ [&](triton::MakeTensorPtrOp makeTensorPtrOp) {
1517+ const LocalTensorShapeData &data =
1518+ shapeMap.at(makeTensorPtrOp.getResult());
1519+ updateMakeTensorPtrOp(rewriter, makeTensorPtrOp, data);
1520+ })
1521+ .Case<InferTypeOpInterface>([&](InferTypeOpInterface inferOp) {
1522+ updateInferTypeOp(rewriter, inferOp);
1523+ })
1524+ .Default([&](Operation *) {
1525+ if (op->hasTrait<OpTrait::Elementwise>()) {
1526+ updateElementwiseOp(rewriter, op);
1527+ }
1528+ });
1529+}
1530+ 
1531+// Returns true if the reduce op reduction algorithms all have identities
1532+bool reduceOrScanOpNeedsMask(Operation *op) {
1533+ SmallVector<TypedAttr> identities = getReduceOrScanOpIdentities(op);
1534+ for (TypedAttr attr : identities) {
1535+ if (!attr) {
1536+ return true;
1537+ }
1538+ }
1539+ 
1540+ return false;
1541+}
1542+ 
1543+// Use only when it is known that at least one reduce op tensor calculation does
1544+// not have a simple identity If all tensor operand reduction algorithms have a
1545+// simple identity use updateGeneralTensorOp For each tensor operand whose
1546+// reduction algorithm does not have a simple identity, uses a mask to ignore
1547+// those operands
1548+triton::ReduceOp getReplacementReduceOp(IRRewriter &rewriter,
1549+ triton::ReduceOp reduceOp, Value mask) {
1550+ SmallVector<TypedAttr> identities = getReduceOrScanOpIdentities(reduceOp);
1551+ // Need to add a mask as another argument to ignore padding values
1552+ 
1553+ rewriter.setInsertionPoint(reduceOp);
1554+ Location loc = reduceOp.getLoc();
1555+ 
1556+ size_t axis = reduceOp.getAxis();
1557+ SmallVector<Value> newOperands(reduceOp->getOperands().begin(),
1558+ reduceOp->getOperands().end());
1559+ newOperands.push_back(mask);
1560+ 
1561+ auto newOp = rewriter.create<triton::ReduceOp>(loc, newOperands, axis);
1562+ 
1563+ {
1564+ Block *block = rewriter.createBlock(&(newOp.getCombineOp()));
1565+ 
1566+ IRRewriter::InsertionGuard guard(rewriter);
1567+ rewriter.setInsertionPointToStart(block);
1568+ 
1569+ unsigned numOriginalOperands = reduceOp.getNumOperands();
1570+ 
1571+ // Setting up reduce block arguments
1572+ SmallVector<Type> argElementTypes;
1573+ argElementTypes.reserve(numOriginalOperands + 1);
1574+ 
1575+ for (Value operand : reduceOp.getOperands()) {
1576+ argElementTypes.push_back(
1577+ cast<RankedTensorType>(operand.getType()).getElementType());
1578+ }
1579+ 
1580+ // Element type of mask
1581+ argElementTypes.push_back(rewriter.getI1Type());
1582+ 
1583+ SmallVector<Location> locations;
1584+ locations.assign(argElementTypes.size(), loc);
1585+ 
1586+ // add args twice for both the current value and accumulator args
1587+ block->addArguments(argElementTypes, locations);
1588+ block->addArguments(argElementTypes, locations);
1589+ 
1590+ // map block arguments in the old reduceOp to the new block arguments
1591+ IRMapping mapping;
1592+ Block &oldBlock = reduceOp.getRegion().front();
1593+ for (unsigned i = 0; i < numOriginalOperands; i++) {
1594+ mapping.map(oldBlock.getArgument(i), block->getArgument(i));
1595+ }
1596+ 
1597+ for (unsigned i = numOriginalOperands; i < 2 * numOriginalOperands; i++) {
1598+ mapping.map(
1599+ oldBlock.getArgument(i),
1600+ block->getArgument(i + 1)); // add one to skip the mask operand
1601+ }
1602+ 
1603+ // copy ops using mapping
1604+ for (auto &op : oldBlock.without_terminator()) {
1605+ rewriter.clone(op, mapping);
1606+ }
1607+ 
1608+ // Use mask now
1609+ Value curMaskVal = block->getArgument(numOriginalOperands);
1610+ Value curMaskAcc = block->getArgument(2 * numOriginalOperands + 1);
1611+ 
1612+ auto oldReturn = cast<triton::ReduceReturnOp>(oldBlock.getTerminator());
1613+ SmallVector<Value> outputs;
1614+ outputs.reserve(numOriginalOperands + 1);
1615+ for (unsigned i = 0; i < numOriginalOperands; i++) {
1616+ TypedAttr identity = identities[i];
1617+ Value laneCombinedRes = mapping.lookup(oldReturn->getOperand(i));
1618+ 
1619+ if (identity) {
1620+ // This operand has an identity, no need to mask
1621+ outputs.push_back(laneCombinedRes);
1622+ continue;
1623+ }
1624+ Value curLaneVal = block->getArgument(i);
1625+ Value curLaneAcc = block->getArgument(numOriginalOperands + 1 + i);
1626+ 
1627+ // This is the value to return if the mask is true at this spot.
1628+ // Using a select op here in case we have only seen padding so far (which
1629+ // means that laneCombinedRes is the result of f(padding, val))
1630+ Value accVal = rewriter.create<arith::SelectOp>(
1631+ loc, curMaskAcc, laneCombinedRes, curLaneVal);
1632+ 
1633+ // Final value
1634+ Value chosenVal =
1635+ rewriter.create<arith::SelectOp>(loc, curMaskVal, accVal, curLaneAcc);
1636+ outputs.push_back(chosenVal);
1637+ }
1638+ 
1639+ // accumulating with or (accumulated mask is true if we have seen a
1640+ // non-padding element so far)
1641+ Value chosenMask =
1642+ rewriter.create<arith::OrIOp>(loc, curMaskVal, curMaskAcc);
1643+ outputs.push_back(chosenMask);
1644+ 
1645+ rewriter.create<triton::ReduceReturnOp>(loc, outputs);
1646+ }
1647+ 
1648+ // triton::ReduceOp implements the InferTypeOpInterface (we use it to update
1649+ // the result type(s))
1650+ updateInferTypeOp(rewriter, newOp);
1651+ 
1652+ return newOp;
1653+}
1654+ 
1655+// Use only when the scan op has reverse=true and at least one scan op tensor
1656+// calculation does not have a simple identity If all tensor operand reduction
1657+// algorithms have a simple identity use updateGeneralTensorOp For each tensor
1658+// operand whose reduction algorithm does not have a simple identity, uses a
1659+// mask to ignore those operands
1660+triton::ScanOp getReplacementScanOp(IRRewriter &rewriter, triton::ScanOp scanOp,
1661+ Value mask) {
1662+ SmallVector<TypedAttr> identities = getReduceOrScanOpIdentities(scanOp);
1663+ // Need to add a mask as another argument to ignore padding values
1664+ 
1665+ rewriter.setInsertionPoint(scanOp);
1666+ Location loc = scanOp.getLoc();
1667+ 
1668+ SmallVector<Value> newOperands(scanOp->getOperands().begin(),
1669+ scanOp->getOperands().end());
1670+ newOperands.push_back(mask);
1671+ 
1672+ auto newOp = rewriter.create<triton::ScanOp>(
1673+ loc, newOperands, scanOp.getAxis(), scanOp.getReverse());
1674+ 
1675+ {
1676+ Block *block = rewriter.createBlock(&(newOp.getCombineOp()));
1677+ 
1678+ IRRewriter::InsertionGuard guard(rewriter);
1679+ rewriter.setInsertionPointToStart(block);
1680+ 
1681+ unsigned numOriginalOperands = scanOp.getNumOperands();
1682+ 
1683+ // Setting up reduce block arguments
1684+ SmallVector<Type> argElementTypes;
1685+ argElementTypes.reserve(numOriginalOperands + 1);
1686+ 
1687+ for (Value operand : scanOp.getOperands()) {
1688+ argElementTypes.push_back(
1689+ cast<RankedTensorType>(operand.getType()).getElementType());
1690+ }
1691+ 
1692+ // Element type of mask
1693+ argElementTypes.push_back(rewriter.getI1Type());
1694+ 
1695+ SmallVector<Location> locations;
1696+ locations.assign(argElementTypes.size(), loc);
1697+ 
1698+ // add args twice for both the current value and accumulator args
1699+ block->addArguments(argElementTypes, locations);
1700+ block->addArguments(argElementTypes, locations);
1701+ 
1702+ // map block arguments in the old scanOp to the new block arguments
1703+ IRMapping mapping;
1704+ Block &oldBlock = scanOp.getRegion().front();
1705+ for (unsigned i = 0; i < numOriginalOperands; i++) {
1706+ mapping.map(oldBlock.getArgument(i), block->getArgument(i));
1707+ }
1708+ 
1709+ for (unsigned i = numOriginalOperands; i < 2 * numOriginalOperands; i++) {
1710+ mapping.map(
1711+ oldBlock.getArgument(i),
1712+ block->getArgument(i + 1)); // add one to skip the mask operand
1713+ }
1714+ 
1715+ // copy ops using mapping
1716+ for (auto &op : oldBlock.without_terminator()) {
1717+ rewriter.clone(op, mapping);
1718+ }
1719+ 
1720+ // Use mask now
1721+ Value curMaskVal = block->getArgument(numOriginalOperands);
1722+ Value curMaskAcc = block->getArgument(2 * numOriginalOperands + 1);
1723+ 
1724+ auto oldReturn = cast<triton::ScanReturnOp>(oldBlock.getTerminator());
1725+ SmallVector<Value> outputs;
1726+ outputs.reserve(numOriginalOperands + 1);
1727+ for (unsigned i = 0; i < numOriginalOperands; i++) {
1728+ TypedAttr identity = identities[i];
1729+ Value laneCombinedRes = mapping.lookup(oldReturn->getOperand(i));
1730+ 
1731+ if (identity) {
1732+ // This operand has an identity, no need to mask
1733+ outputs.push_back(laneCombinedRes);
1734+ continue;
1735+ }
1736+ Value curLaneVal = block->getArgument(i);
1737+ Value curLaneAcc = block->getArgument(numOriginalOperands + 1 + i);
1738+ 
1739+ // This is the value to return if the mask is true at this spot.
1740+ // Using a select op here in case we have only seen padding so far (which
1741+ // means that laneCombinedRes is the result of f(padding, val))
1742+ Value accVal = rewriter.create<arith::SelectOp>(
1743+ loc, curMaskAcc, laneCombinedRes, curLaneVal);
1744+ 
1745+ // Final value
1746+ Value chosenVal =
1747+ rewriter.create<arith::SelectOp>(loc, curMaskVal, accVal, curLaneAcc);
1748+ outputs.push_back(chosenVal);
1749+ }
1750+ 
1751+ // accumulating with or (accumulated mask is true if we have seen a
1752+ // non-padding element so far)
1753+ Value chosenMask =
1754+ rewriter.create<arith::OrIOp>(loc, curMaskVal, curMaskAcc);
1755+ outputs.push_back(chosenMask);
1756+ 
1757+ rewriter.create<triton::ScanReturnOp>(loc, outputs);
1758+ }
1759+ 
1760+ // triton::ScanOp implements the InferTypeOpInterface (we use it to update the
1761+ // result type(s))
1762+ updateInferTypeOp(rewriter, newOp);
1763+ 
1764+ return newOp;
1765+}
1766+ 
1767+// Creates a replacement tensor creation op (arith::ConstantOp, triton::SplatOp,
1768+// triton::MakeRangeOp) without any specific padding
1769+Value getReplacementCreationOp(OpBuilder &builder, Operation *tensorCreatorOp,
1770+ const LocalTensorShapeData &data) {
1771+ Value res;
1772+ Location loc = tensorCreatorOp->getLoc();
1773+ 
1774+ builder.setInsertionPointAfter(tensorCreatorOp);
1775+ TypeSwitch<Operation *, void>(tensorCreatorOp)
1776+ .Case<triton::SplatOp>([&](triton::SplatOp splatOp) {
1777+ Type elementType = splatOp.getType().getElementType();
1778+ res = builder.create<triton::SplatOp>(
1779+ loc, RankedTensorType::get(data.localVirtualShape, elementType),
1780+ splatOp.getSrc());
1781+ })
1782+ .Case<arith::ConstantOp>([&](arith::ConstantOp constOp) {
1783+ RankedTensorType dataType = cast<RankedTensorType>(constOp.getType());
1784+ Type elementType = dataType.getElementType();
1785+ auto origAttr = cast<DenseElementsAttr>(constOp.getValue());
1786+ RankedTensorType virtualShapeType =
1787+ RankedTensorType::get(data.localVirtualShape, elementType);
1788+ // The constant op must be a splat constant op (otherwise the tensor
1789+ // would not have the same number of elements as number of threads per
1790+ // warp which is a strict requirement)
1791+ Attribute attr = origAttr.getSplatValue<Attribute>();
1792+ res = createTensor(builder, loc, virtualShapeType, attr);
1793+ })
1794+ .Case<triton::MakeRangeOp>([&](triton::MakeRangeOp makeRangeOp) {
1795+ Type elementType = makeRangeOp.getType().getElementType();
1796+ int start = makeRangeOp.getStart();
1797+ res = builder.create<triton::MakeRangeOp>(
1798+ loc, RankedTensorType::get(data.localVirtualShape, elementType),
1799+ start, start + data.localVirtualShapeSize);
1800+ });
1801+ 
1802+ return res;
1803+}
1804+ 
1805+// Updates a simple triton ReshapeOp, which is a reshape op that just
1806+// adds/removes dimensions of size 1
1807+void updateSimpleReshapeOp(IRRewriter &rewriter, triton::ReshapeOp reshapeOp,
1808+ const LocalTensorShapeData &resultShape) {
1809+ Type elementType = reshapeOp.getResult().getType().getElementType();
1810+ RankedTensorType newType =
1811+ RankedTensorType::get(resultShape.localVirtualShape, elementType);
1812+ rewriter.modifyOpInPlace(reshapeOp,
1813+ [&]() { reshapeOp.getResult().setType(newType); });
1814+}
1815+ 
1816+// Updates a complex triton ReshapeOp, which is a reshape op that does not just
1817+// add/remove dimensions of size 1 Flattens the operand tensor, then uses
1818+// tensor.extract/insert slice to reposition the data elements Then reshapes
1819+// this tensor to the final shape Very inefficient, so its best to avoid using
1820+// this if possible
1821+Value getReplacementComplexReshapeOp(OpBuilder &builder,
1822+ triton::ReshapeOp reshapeOp,
1823+ const LocalTensorShapeData &resultData,
1824+ const std::vector<bool> &srcValueLocs,
1825+ const std::vector<bool> &resValueLocs) {
1826+ Location loc = reshapeOp->getLoc();
1827+ 
1828+ RankedTensorType srcType = reshapeOp.getSrc().getType();
1829+ int64_t srcNumEls = static_cast<int64_t>(srcValueLocs.size());
1830+ int64_t resNumEls = static_cast<int64_t>(resValueLocs.size());
1831+ Type elementType = srcType.getElementType();
1832+ builder.setInsertionPoint(reshapeOp);
1833+ 
1834+ // Reshape source tensor to 1D
1835+ RankedTensorType flattenedSrcType =
1836+ RankedTensorType::get({srcNumEls}, elementType);
1837+ Value flattenedSrc;
1838+ if (srcType.getRank() == 1) {
1839+ flattenedSrc = reshapeOp.getSrc();
1840+ } else {
1841+ flattenedSrc = builder.create<triton::ReshapeOp>(loc, flattenedSrcType,
1842+ reshapeOp.getSrc());
1843+ }
1844+ 
1845+ // Create 1D destination tensor
1846+ RankedTensorType flattenedResType =
1847+ RankedTensorType::get({resNumEls}, elementType);
1848+ Value dest = createTensor(builder, loc, flattenedResType);
1849+ 
1850+ // Extract and insert slices
1851+ int64_t srcIdx = 0;
1852+ int64_t dstIdx = 0;
1853+ 
1854+ SmallVector<OpFoldResult> strides(1, builder.getIndexAttr(1));
1855+ while (srcIdx < srcNumEls && dstIdx < resNumEls) {
1856+ if (!srcValueLocs[srcIdx]) {
1857+ srcIdx += 1;
1858+ }
1859+ if (!resValueLocs[dstIdx]) {
1860+ dstIdx += 1;
1861+ }
1862+ int64_t srcStart = srcIdx;
1863+ int64_t dstStart = dstIdx;
1864+ 
1865+ while (srcIdx < srcNumEls && dstIdx < resNumEls && srcValueLocs[srcIdx] &&
1866+ resValueLocs[dstIdx]) {
1867+ srcIdx += 1;
1868+ dstIdx += 1;
1869+ }
1870+ 
1871+ int64_t sliceSize = srcIdx - srcStart;
1872+ if (sliceSize > 0) {
1873+ int64_t paddedSliceSize =
1874+ static_cast<int64_t>(llvm::PowerOf2Ceil(sliceSize));
1875+ if (dstStart + paddedSliceSize <= resNumEls &&
1876+ srcStart + paddedSliceSize <= srcNumEls) {
1877+ // If we have enough space, just extract and insert a large slice
1878+ // extra elements will be overriden anyways later
1879+ SmallVector<OpFoldResult> offsets(1, builder.getIndexAttr(srcStart));
1880+ SmallVector<OpFoldResult> sizes(1,
1881+ builder.getIndexAttr(paddedSliceSize));
1882+ 
1883+ Value extractedSlice = builder.create<tensor::ExtractSliceOp>(
1884+ loc, flattenedSrc, offsets, sizes, strides);
1885+ offsets[0] = builder.getIndexAttr(dstStart);
1886+ dest = builder.create<tensor::InsertSliceOp>(loc, extractedSlice, dest,
1887+ offsets, sizes, strides);
1888+ } else {
1889+ // Not enough space to insert a large slice
1890+ // extract + insert small power of two slices
1891+ SmallVector<Value> extractedSlices;
1892+ int64_t srcOffset = srcStart;
1893+ int64_t dstOffset = dstStart;
1894+ 
1895+ while (sliceSize > 0) {
1896+ int64_t chosenSlice =
1897+ static_cast<int64_t>(llvm::PowerOf2Ceil(sliceSize));
1898+ if (chosenSlice > sliceSize) {
1899+ chosenSlice /= 2;
1900+ }
1901+ sliceSize -= chosenSlice;
1902+ 
1903+ SmallVector<OpFoldResult> offsets(1, builder.getIndexAttr(srcOffset));
1904+ SmallVector<OpFoldResult> sizes(1, builder.getIndexAttr(chosenSlice));
1905+ Value slice = builder.create<tensor::ExtractSliceOp>(
1906+ loc, flattenedSrc, offsets, sizes, strides);
1907+ offsets[0] = builder.getIndexAttr(dstOffset);
1908+ dest = builder.create<tensor::InsertSliceOp>(loc, slice, dest,
1909+ offsets, sizes, strides);
1910+ 
1911+ srcOffset += chosenSlice;
1912+ dstOffset += chosenSlice;
1913+ }
1914+ }
1915+ }
1916+ }
1917+ RankedTensorType resType =
1918+ RankedTensorType::get(resultData.localVirtualShape, elementType);
1919+ Value res = builder.create<triton::ReshapeOp>(loc, resType, dest);
1920+ 
1921+ return res;
1922+}
1923+ 
1924+Value getReplacementHistogramOp(OpBuilder &builder, triton::HistogramOp op,
1925+ const LocalTensorShapeData &operandData,
1926+ const LocalTensorShapeData &resultData) {
1927+ Value source = op.getSrc();
1928+ builder.setInsertionPoint(op);
1929+ Location loc = op->getLoc();
1930+ 
1931+ Type elementType = op.getResult().getType().getElementType();
1932+ RankedTensorType resType =
1933+ RankedTensorType::get(resultData.localVirtualShape, elementType);
1934+ Value res;
1935+ 
1936+ if (hasNonPowerTwoDim(operandData.dataShape)) {
1937+ Value mask = op.getMask();
1938+ if (mask) {
1939+ Value paddingMask = createPaddingMask(builder, op, operandData.dataShape,
1940+ operandData.localVirtualShape);
1941+ mask = builder.create<arith::AndIOp>(loc, mask, paddingMask);
1942+ } else {
1943+ mask = createPaddingMask(builder, op, operandData.dataShape,
1944+ operandData.localVirtualShape);
1945+ }
1946+ res = builder.create<triton::HistogramOp>(loc, resType, source, mask);
1947+ } else {
1948+ res = builder.create<triton::HistogramOp>(loc, resType, source);
1949+ }
1950+ 
1951+ return res;
1952+}
1953+ 
1954+struct PaddingLattice : public dataflow::Lattice<AllPaddingRequirements> {
1955+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PaddingLattice)
1956+ using Lattice::Lattice;
1957+};
1958+ 
1959+// Pushes padding requirements upstream
1960+class BackwardPaddingPopulationAnalysis
1961+ : public dataflow::SparseBackwardDataFlowAnalysis<PaddingLattice> {
1962+public:
1963+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
1964+ BackwardPaddingPopulationAnalysis)
1965+ BackwardPaddingPopulationAnalysis(
1966+ DataFlowSolver &solver, SymbolTableCollection &symbolTable,
1967+ const DenseMap<Value, LocalTensorShapeData> &shapeData)
1968+ : SparseBackwardDataFlowAnalysis(solver, symbolTable),
1969+ shapeData(shapeData) {}
1970+ 
1971+ DenseMap<Value, PotentialPaddingRequirements> paddingMap;
1972+ const DenseMap<Value, LocalTensorShapeData> &shapeData;
1973+ 
1974+ LogicalResult
1975+ visitOperation(Operation *op, ArrayRef<PaddingLattice *> operands,
1976+ ArrayRef<const PaddingLattice *> results) override {
1977+ OpBuilder builder(op);
1978+ 
1979+ if (!hasTensorArgs(op)) {
1980+ return success();
1981+ }
1982+ 
1983+ // Power of two tensor/non tensor ops do not need padding, we cannot add
1984+ // padding to a tensor pointer
1985+ if ((!isNonPow2TensorOperation(op)) ||
1986+ isTensorPointerLoadStoreOperation(op)) {
1987+ return success();
1988+ }
1989+ 
1990+ // In general, tries to propagate padding requirements up by pushing the
1991+ // requirements of its results up to tensor operands that share the same
1992+ // element type For some ops though, want to push it to certain operands
1993+ // while ignoring other operands For example for arith::SelectOp, when its
1994+ // operands are boolean tensors, the padding reqs of its result could be
1995+ // applied to the mask operand which is not desired
1996+ TypeSwitch<Operation *>(op)
1997+ .Case<arith::SelectOp>([&](arith::SelectOp selectOp) {
1998+ Type elementType =
1999+ cast<RankedTensorType>(selectOp.getType()).getElementType();
2000+ AllPaddingRequirements &trueTensorLattice = operands[1]->getValue();
2001+ Value trueTensorValue = operands[1]->getAnchor();
2002+ OpOperand *trueTensorOperand = &selectOp->getOpOperand(1);
2003+ AllPaddingRequirements &falseTensorLattice = operands[2]->getValue();
2004+ Value falseTensorValue = operands[2]->getAnchor();
2005+ OpOperand *falseTensorOperand = &selectOp->getOpOperand(2);
2006+ 
2007+ ChangeResult trueLatticeChanged = ChangeResult::NoChange;
2008+ ChangeResult falseLatticeChanged = ChangeResult::NoChange;
2009+ 
2010+ for (const TypedAttr &req : results[0]->getValue().reqs) {
2011+ if (req.getType() == elementType) {
2012+ trueLatticeChanged |= trueTensorLattice.add(req);
2013+ paddingMap[trueTensorValue].addLoose(req, trueTensorOperand);
2014+ falseLatticeChanged |= falseTensorLattice.add(req);
2015+ paddingMap[falseTensorValue].addLoose(req, falseTensorOperand);
2016+ }
2017+ }
2018+ 
2019+ propagateIfChanged(operands[1], trueLatticeChanged);
2020+ propagateIfChanged(operands[2], falseLatticeChanged);
2021+ })
2022+ .Default([&](Operation *op) {
2023+ for (size_t i = 0; i < operands.size(); i++) {
2024+ PaddingLattice *operandLattice = operands[i];
2025+ Value val = operandLattice->getAnchor();
2026+ AllPaddingRequirements &operandReqs = operandLattice->getValue();
2027+ ChangeResult changed = ChangeResult::NoChange;
2028+ 
2029+ for (const PaddingLattice *resLattice : results) {
2030+ if (RankedTensorType type = dyn_cast<RankedTensorType>(
2031+ operandLattice->getAnchor().getType())) {
2032+ for (const TypedAttr &req : resLattice->getValue().reqs) {
2033+ if (req.getType() == type.getElementType()) {
2034+ paddingMap[val].addLoose(req, &op->getOpOperand(i));
2035+ changed |= operandReqs.add(req);
2036+ }
2037+ }
2038+ }
2039+ }
2040+ 
2041+ propagateIfChanged(operandLattice, changed);
2042+ }
2043+ });
2044+ SmallVector<TypedAttr> opReqs =
2045+ getPaddingRequirements(builder, op, shapeData);
2046+ 
2047+ if (opReqs.size() == 0) {
2048+ return success();
2049+ }
2050+ 
2051+ // Add this op's required padding
2052+ for (size_t i = 0; i < opReqs.size(); i++) {
2053+ if (opReqs[i]) {
2054+ Value val = op->getOperand(i);
2055+ paddingMap[val].addStrict(opReqs[i], &op->getOpOperand(i));
2056+ 
2057+ ChangeResult changed = operands[i]->getValue().add(opReqs[i]);
2058+ propagateIfChanged(operands[i], changed);
2059+ }
2060+ }
2061+ 
2062+ // Update padding map
2063+ for (const auto &count : paddingMap) {
2064+ paddingMap[count.getFirst()] = count.getSecond();
2065+ }
2066+ 
2067+ return success();
2068+ }
2069+ 
2070+ void visitBranchOperand(OpOperand &operand) override {
2071+ auto branchOp = dyn_cast<BranchOpInterface>(operand.getOwner());
2072+ if (!branchOp) {
2073+ return;
2074+ }
2075+ 
2076+ std::optional<BlockArgument> successorBlockArg =
2077+ branchOp.getSuccessorBlockArgument(operand.getOperandNumber());
2078+ if (!successorBlockArg) {
2079+ return;
2080+ }
2081+ 
2082+ BlockArgument blockArg = *successorBlockArg;
2083+ const PaddingLattice *blockLattice = getLatticeElement(blockArg);
2084+ if ((!blockLattice) || (!blockLattice->getValue().initialized)) {
2085+ return;
2086+ }
2087+ 
2088+ PaddingLattice *srcLattice = getLatticeElement(operand.get());
2089+ 
2090+ ChangeResult changed =
2091+ srcLattice->getValue().meet(blockLattice->getValue());
2092+ 
2093+ propagateIfChanged(srcLattice, changed);
2094+ }
2095+ 
2096+ void visitCallOperand(OpOperand &operand) override {
2097+ auto callOp = dyn_cast<CallOpInterface>(operand.getOwner());
2098+ if (!callOp) {
2099+ return;
2100+ }
2101+ Operation *callable = callOp.resolveCallable();
2102+ if (!callable) {
2103+ return;
2104+ }
2105+ 
2106+ auto callableOp = cast<CallableOpInterface>(callable);
2107+ Region *region = callableOp.getCallableRegion();
2108+ if ((!region) || region->empty()) {
2109+ return;
2110+ }
2111+ 
2112+ unsigned int operandNum = operand.getOperandNumber();
2113+ BlockArgument arg = region->front().getArgument(operandNum);
2114+ 
2115+ const PaddingLattice *argLattice = getLatticeElement(arg);
2116+ if ((!argLattice) || (!argLattice->getValue().initialized)) {
2117+ return;
2118+ }
2119+ 
2120+ PaddingLattice *callerLattice = getLatticeElement(operand.get());
2121+ 
2122+ ChangeResult changed =
2123+ callerLattice->getValue().meet(argLattice->getValue());
2124+ 
2125+ propagateIfChanged(callerLattice, changed);
2126+ }
2127+ 
2128+ void setToExitState(PaddingLattice *lattice) override {
2129+ ChangeResult changed = lattice->getValue().setToDefault();
2130+ propagateIfChanged(lattice, changed);
2131+ }
2132+};
2133+ 
2134+void addOpOperands(
2135+ SmallVector<std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>>> &vec,
2136+ const std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>> &element) {
2137+ for (auto &existingPair : vec) {
2138+ if (existingPair.first == element.first) {
2139+ for (OpOperand *operand : element.second) {
2140+ existingPair.second.insert(operand);
2141+ }
2142+ return;
2143+ }
2144+ }
2145+ 
2146+ vec.push_back(element);
2147+}
2148+ 
2149+// Given the mapping of values to all padding requirements, choose which Value's
2150+// will get padded
2151+// TODO - Currently, the approach replaces the padding with the requirement
2152+// whenever it requires certain padding, regardless of if padding occured
2153+// upstream, etc.
2154+// Could try to improve
2155+// For example, if a tensor requires 0 padding (and is 0 padded), and one of
2156+// its users requires 0 padding, its user will add the 0 padding again In some
2157+// special cases, the padding might still be preserved (for example adding two
2158+// zero padded tensors gives a zero padded tensor)
2159+DenseMap<Value, SmallVector<std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>>>>
2160+choosePadding(const DenseMap<Value, PotentialPaddingRequirements> &reqs) {
2161+ DenseMap<Value,
2162+ SmallVector<std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>>>>
2163+ curPadding;
2164+ for (const auto &req : reqs) {
2165+ Value cur = req.getFirst();
2166+ const DenseMap<TypedAttr, SmallPtrSet<OpOperand *, 2>> &strictReqs =
2167+ req.getSecond().strictReqs;
2168+ for (const auto &strictReq : strictReqs) {
2169+ const std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>> &p = strictReq;
2170+ addOpOperands(curPadding[cur], p);
2171+ }
2172+ }
2173+ 
2174+ return curPadding;
2175+}
2176+ 
2177+// Given a Value and a padding, return a new Value that is the original Value
2178+// padded with the desired padding Sets insertion point to after mask
2179+Value setPaddingForValue(OpBuilder &builder, Location loc, Value val,
2180+ TypedAttr padding, Value paddingMask) {
2181+ RankedTensorType type = cast<RankedTensorType>(val.getType());
2182+ builder.setInsertionPointAfterValue(paddingMask);
2183+ 
2184+ Value tensorPadding = createTensor(builder, loc, type, padding);
2185+ 
2186+ Value selectOp =
2187+ builder.create<arith::SelectOp>(loc, paddingMask, val, tensorPadding);
2188+ return selectOp;
2189+}
2190+ 
2191+// Before replacing the op with the replacement Values, updates the shapeMap and
2192+// chosenPadding map to ensure these keys exist
2193+void replaceOpSafely(
2194+ IRRewriter &rewriter, Operation *op, const ValueRange &replacements,
2195+ DenseMap<Value, LocalTensorShapeData> &shapeMap,
2196+ DenseMap<Value,
2197+ SmallVector<std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>>>>
2198+ &chosenPadding,
2199+ Operation *&newOpPtr) {
2200+ for (size_t i = 0; i < op->getNumResults(); i++) {
2201+ Value res = op->getResult(i);
2202+ Value replacement = replacements[i];
2203+ if (shapeMap.contains(res)) {
2204+ const LocalTensorShapeData shapeData = shapeMap.at(res);
2205+ shapeMap.erase(res);
2206+ shapeMap[replacement] = shapeData;
2207+ }
2208+ if (chosenPadding.contains(res)) {
2209+ const auto paddingData = chosenPadding.at(res);
2210+ chosenPadding.erase(res);
2211+ chosenPadding[replacement] = paddingData;
2212+ }
2213+ }
2214+ 
2215+ rewriter.replaceOp(op, replacements);
2216+ newOpPtr = replacements[0].getDefiningOp();
2217+}
2218+ 
2219+// final ir replacement/rewrite step
2220+void finalCodegen(
2221+ FuncOp &mod, DenseMap<Value, LocalTensorShapeData> &shapeMap,
2222+ DenseMap<Value,
2223+ SmallVector<std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>>>>
2224+ &chosenPadding) {
2225+ std::queue<Operation *> workQueue;
2226+ 
2227+ mod.walk<WalkOrder::PreOrder>([&](Operation *op) {
2228+ if (isNonPow2TensorOperation(op)) {
2229+ workQueue.push(op);
2230+ }
2231+ });
2232+ 
2233+ IRRewriter rewriter(mod);
2234+ 
2235+ while (!workQueue.empty()) {
2236+ Operation *op = workQueue.front();
2237+ workQueue.pop();
2238+ 
2239+ // Need to keep track of results for iteration later
2240+ // original Operation* op ptr may not be valid after using replaceOp
2241+ Operation *resOp = op;
2242+ if (isTensorProducer(op)) {
2243+ Value val = op->getResult(0);
2244+ const LocalTensorShapeData &shapeData = shapeMap.at(val);
2245+ Value replacement = getReplacementCreationOp(rewriter, op, shapeData);
2246+ 
2247+ replaceOpSafely(rewriter, op, replacement, shapeMap, chosenPadding,
2248+ resOp);
2249+ } else if (auto reduceOp = dyn_cast<triton::ReduceOp>(op)) {
2250+ // Getting mask
2251+ // all operands are the same shape before (and thus same shape after)
2252+ Value firstOperand = reduceOp.getOperand(0);
2253+ const LocalTensorShapeData &data = shapeMap.at(firstOperand);
2254+ 
2255+ if (!reduceOrScanOpNeedsMask(reduceOp)) {
2256+ updateGeneralTensorOp(rewriter, op, shapeMap);
2257+ } else {
2258+ Value mask = createPaddingMask(rewriter, op, data.dataShape,
2259+ data.localVirtualShape);
2260+ triton::ReduceOp res = getReplacementReduceOp(rewriter, reduceOp, mask);
2261+ ValueRange replacements = res->getResults().drop_back();
2262+ replaceOpSafely(rewriter, op, replacements, shapeMap, chosenPadding,
2263+ resOp);
2264+ }
2265+ } else if (auto scanOp = dyn_cast<triton::ScanOp>(op)) {
2266+ Value firstOperand = scanOp.getOperand(0);
2267+ if (!scanOp.getReverse() || !reduceOrScanOpNeedsMask(scanOp)) {
2268+ updateGeneralTensorOp(rewriter, op, shapeMap);
2269+ } else {
2270+ const LocalTensorShapeData &data = shapeMap.at(firstOperand);
2271+ Value mask = createPaddingMask(rewriter, op, data.dataShape,
2272+ data.localVirtualShape);
2273+ triton::ScanOp res = getReplacementScanOp(rewriter, scanOp, mask);
2274+ ValueRange replacements = res->getResults().drop_back();
2275+ replaceOpSafely(rewriter, op, replacements, shapeMap, chosenPadding,
2276+ resOp);
2277+ }
2278+ } else if (auto reshapeOp = dyn_cast<triton::ReshapeOp>(op)) {
2279+ const LocalTensorShapeData &operandData = shapeMap.at(reshapeOp.getSrc());
2280+ const LocalTensorShapeData &resultData =
2281+ shapeMap.at(reshapeOp.getResult());
2282+ std::optional<std::pair<std::vector<bool>, std::vector<bool>>>
2283+ potentialDataLocs =
2284+ getComplexReshapeDataLocs(operandData, resultData);
2285+ if (!potentialDataLocs) {
2286+ updateSimpleReshapeOp(rewriter, reshapeOp, resultData);
2287+ } else {
2288+ std::pair<std::vector<bool>, std::vector<bool>> dataLocs =
2289+ *potentialDataLocs;
2290+ Value replacement = getReplacementComplexReshapeOp(
2291+ rewriter, reshapeOp, resultData, dataLocs.first, dataLocs.second);
2292+ replaceOpSafely(rewriter, op, replacement, shapeMap, chosenPadding,
2293+ resOp);
2294+ }
2295+ } else if (auto histogramOp = dyn_cast<triton::HistogramOp>(op)) {
2296+ const LocalTensorShapeData &operandData =
2297+ shapeMap.at(histogramOp.getSrc());
2298+ const LocalTensorShapeData &resultData =
2299+ shapeMap.at(histogramOp.getResult());
2300+ 
2301+ Value replacement = getReplacementHistogramOp(rewriter, histogramOp,
2302+ operandData, resultData);
2303+ ValueRange replacements = replacement;
2304+ replaceOpSafely(rewriter, op, replacements, shapeMap, chosenPadding,
2305+ resOp);
2306+ } else if (isGenericTensorOp(op)) {
2307+ updateGeneralTensorOp(rewriter, op, shapeMap);
2308+ }
2309+ 
2310+ ValueRange results = resOp->getResults();
2311+ ValueRange blockArgs;
2312+ if (!resOp->getRegions().empty() && !resOp->getRegion(0).empty()) {
2313+ blockArgs = resOp->getRegion(0).front().getArguments();
2314+ }
2315+ 
2316+ // TODO - support other ops here or in isGenericTensorOp +
2317+ // updateGeneralTensorOp (add before the isGenericTensorOp check)
2318+ 
2319+ // If any results require padding then apply the padding
2320+ for (Value res : results) {
2321+ if (chosenPadding.contains(res)) {
2322+ const SmallVector<std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>>>
2323+ &allPadding = chosenPadding.at(res);
2324+ const LocalTensorShapeData shapeData = shapeMap.at(res);
2325+ Value mask = createPaddingMask(rewriter, res.getDefiningOp(),
2326+ shapeData.dataShape,
2327+ shapeData.localVirtualShape, true);
2328+ // Go through pairs of (Padding value, Set of Users)
2329+ for (const auto &paddingUsers : allPadding) {
2330+ // Apply padding
2331+ TypedAttr padding = paddingUsers.first;
2332+ Value replacement = setPaddingForValue(
2333+ rewriter, res.getDefiningOp()->getLoc(), res, padding, mask);
2334+ shapeMap[replacement] = shapeData;
2335+ 
2336+ // Set padding for users
2337+ for (OpOperand *operand : paddingUsers.second) {
2338+ rewriter.modifyOpInPlace(operand->getOwner(),
2339+ [&]() { operand->set(replacement); });
2340+ }
2341+ }
2342+ }
2343+ }
2344+ 
2345+ // If any block args require padding then apply the padding
2346+ for (Value arg : blockArgs) {
2347+ if (chosenPadding.contains(arg)) {
2348+ const SmallVector<std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>>>
2349+ &allPadding = chosenPadding.at(arg);
2350+ const LocalTensorShapeData shapeData = shapeMap.at(arg);
2351+ 
2352+ // Insert at start of block instead of before op
2353+ Block *block = arg.getParentBlock();
2354+ Location loc = block->getParentOp()->getLoc();
2355+ rewriter.setInsertionPointToStart(block);
2356+ Value mask = createPaddingMask(rewriter, loc, shapeData.dataShape,
2357+ shapeData.localVirtualShape);
2358+ // Go through pairs of (Padding value, Set of Users)
2359+ for (const auto &paddingUsers : allPadding) {
2360+ // Apply padding
2361+ TypedAttr padding = paddingUsers.first;
2362+ Value replacement =
2363+ setPaddingForValue(rewriter, loc, arg, padding, mask);
2364+ shapeMap[replacement] = shapeData;
2365+ 
2366+ // Set padding for users
2367+ for (OpOperand *operand : paddingUsers.second) {
2368+ rewriter.modifyOpInPlace(operand->getOwner(),
2369+ [&]() { operand->set(replacement); });
2370+ }
2371+ }
2372+ }
2373+ }
2374+ }
2375+}
2376+ 
2377+class ConvertNonPowerTwoTensorsPass
2378+ : public impl::ConvertNonPowerTwoTensorsBase<
2379+ ConvertNonPowerTwoTensorsPass> {
2380+public:
2381+ using ConvertNonPowerTwoTensorsBase::ConvertNonPowerTwoTensorsBase;
2382+ void runOnOperation() override {
2383+ FuncOp module = getOperation();
2384+ 
2385+ // Make sure that the non power of two tensor ops are supported by this pass
2386+ if (failed(verifyOps(module))) {
2387+ module.emitError("Unsupported non power of two tensor operations found");
2388+ signalPassFailure();
2389+ return;
2390+ }
2391+ 
2392+ // For tt.load, tt.store, tt.atomic_rmw ops that accept a tensor of
2393+ // pointers, add a default mask and other tensor if missing, or set the
2394+ // boundary check attr for load/stores on tensor pointers
2395+ addMasksToLoadAndStores(module);
2396+ 
2397+ // For tensor.insert_slice ops that have a non power of 2 dim on the slice
2398+ // axis, or tensor.extract_slice ops that have a dynamic offset, or offset +
2399+ // expanded slice axis > source tensor slice axis, splits the
2400+ // insert_slice/extract_slice into power of 2 slices along the slice axis
2401+ splitSliceOps(module);
2402+ 
2403+ // For each non power of two tensor operand/result, calculate shape data
2404+ DenseMap<Value, LocalTensorShapeData> shapeMap =
2405+ populateLocalShapeData(module);
2406+ 
2407+ // Determine padding requirements and push them up to their sources
2408+ DataFlowSolver solver;
2409+ SymbolTableCollection symbolTable;
2410+ solver.load<dataflow::DeadCodeAnalysis>();
2411+ solver.load<dataflow::SparseConstantPropagation>();
2412+ auto *paddingPropagationAnalysis =
2413+ solver.load<BackwardPaddingPopulationAnalysis>(symbolTable, shapeMap);
2414+ 
2415+ if (failed(solver.initializeAndRun(module))) {
2416+ module.emitError("Error occured trying to perform padding analysis");
2417+ signalPassFailure();
2418+ return;
2419+ }
2420+ 
2421+ DenseMap<Value, PotentialPaddingRequirements> potentialPadding =
2422+ paddingPropagationAnalysis->paddingMap;
2423+ 
2424+ // Choose what padding(s) to use for each Value which needs to have padding
2425+ // TODO - Currently implemented in an inefficient way
2426+ DenseMap<Value,
2427+ SmallVector<std::pair<TypedAttr, SmallPtrSet<OpOperand *, 2>>>>
2428+ chosenPadding = choosePadding(potentialPadding);
2429+ // Replace and generate instructions
2430+ finalCodegen(module, shapeMap, chosenPadding);
2431+ }
2432+};
2433+ 
2434+} // namespace
2435+ 
2436+std::unique_ptr<mlir::Pass> createConvertNonPowerTwoTensorsPass() {
2437+ return std::make_unique<ConvertNonPowerTwoTensorsPass>();
2438+}
2439+ 
2440+} // namespace bishengir::triton
Mbishengir/lib/Dialect/Triton/Transforms/RewriteSliceOpToTriton.cpp+54-32
@@ -302,26 +302,15 @@ planSliceRewrite(Operation *op, RankedTensorType large, RankedTensorType small,
302 return failure();302 return failure();
303 }303 }
304 304 
305- if (plan.r % plan.S != 0) {305+ if (plan.isOffsetStatic && plan.r % plan.S != 0) {
306- op->emitError("offset at indexed axis ")306+ OpBuilder builder(op);
307- << plan.axis << " must be a multiple of size " << plan.S << "; got "307+ plan.isOffsetStatic = false;
308- << plan.r;308+ plan.dynamicR = builder.create<arith::ConstantOp>(op->getLoc(), builder.getIndexAttr(plan.r));
309- return failure();
310 }309 }
311 310 
312 plan.m = plan.r / plan.S;311 plan.m = plan.r / plan.S;
313 }312 }
314 313 
315- // If we have a `tensor.insert_slice` op, we must check that N is divisible by
316- // S as we are not guaranteed N and S are multiples of 2
317- bool isInsertSliceOp = isa<tensor::InsertSliceOp>(op);
318- if (isInsertSliceOp && plan.N % plan.S != 0) {
319- op->emitError(
320- "dynamic insert_slice rewrite requires slice size of small tensor ")
321- << plan.S << " to divide the dest dim " << plan.N;
322- return failure();
323- }
324- 
325 plan.k = log2Pow2(plan.N / plan.S);314 plan.k = log2Pow2(plan.N / plan.S);
326 return plan;315 return plan;
327}316}
@@ -447,35 +436,68 @@ static Value useMaskToInsert(OpBuilder &builder, Location loc, Value large,
447 Value mask, Value tensorToInsert,436 Value mask, Value tensorToInsert,
448 const SlicePlan &plan) {437 const SlicePlan &plan) {
449 RankedTensorType largeType = cast<RankedTensorType>(large.getType());438 RankedTensorType largeType = cast<RankedTensorType>(large.getType());
439+ ArrayRef<int64_t> largeShape = largeType.getShape();
450 ArrayRef<int64_t> smallShape =440 ArrayRef<int64_t> smallShape =
451 cast<RankedTensorType>(tensorToInsert.getType()).getShape();441 cast<RankedTensorType>(tensorToInsert.getType()).getShape();
452 Type elementType = largeType.getElementType();442 Type elementType = largeType.getElementType();
443+ Type i32Type = builder.getI32Type();
453 const int64_t S = plan.S;444 const int64_t S = plan.S;
454 const int64_t N = plan.N;445 const int64_t N = plan.N;
455- const int axis = plan.axis;446+ const size_t axis = plan.axis;
447+ const Value offset = plan.dynamicR;
448+ const size_t rank = largeShape.size();
456 449 
457- // Calculate shape after expanding by inserting 1 before dimension S450+
458- SmallVector<int64_t> curShape(smallShape.begin(), smallShape.end());451+ // Make range
459- curShape.insert(curShape.begin() + axis, 1);452+ SmallVector<int64_t> curShape;
453+ curShape.reserve(rank);
454+ curShape.push_back(N);
455+ RankedTensorType rangeType = RankedTensorType::get(curShape, i32Type);
456+ RankedTensorType finalIndicesType = RankedTensorType::get(largeShape, i32Type);
457+ Value indices = builder.create<triton::MakeRangeOp>(loc, rangeType, 0, N);
458+ Value indexToIntCast = builder.create<arith::IndexCastOp>(loc, i32Type, offset);
459+ Value splatOffset = builder.create<triton::SplatOp>(loc, rangeType, indexToIntCast);
460 460 
461- // Expand461+ indices = builder.create<arith::SubIOp>(loc, indices, splatOffset);
462- auto expandedTensor = builder.create<triton::ExpandDimsOp>(
463- loc, RankedTensorType::get(curShape, elementType), tensorToInsert, axis);
464 462 
465- // Calculate shape after broadcasting (the 1 we inserted will turn into N/S)463+ DenseElementsAttr shiftConstAttr = DenseElementsAttr::get(rangeType, builder.getI32IntegerAttr(N));
466- curShape[axis] = N / S;464+ Value shiftConst = builder.create<arith::ConstantOp>(loc, shiftConstAttr);
465+ Value maskIndices = builder.create<arith::AddIOp>(loc, indices, shiftConst);
467 466 
468- // Perform broadcast467+ DenseElementsAttr lowerBoundAttr = DenseElementsAttr::get(rangeType, builder.getI32IntegerAttr(0));
469- auto broadcastedTensor = builder.create<triton::BroadcastOp>(468+ Value lowerBound = builder.create<arith::ConstantOp>(loc, lowerBoundAttr);
470- loc, RankedTensorType::get(curShape, elementType), expandedTensor);469+ Value indexMask = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sge, indices, lowerBound);
471 470 
472- // Reshape to collapse the N/S and S dimensions to a dimension of size N471+ indices = builder.create<arith::SelectOp>(loc, indexMask, indices, maskIndices);
473- auto reshapedTensor =
474- builder.create<triton::ReshapeOp>(loc, largeType, broadcastedTensor);
475 472 
476- // Select reshaped insert tensor when true, original tensor when false473+ // We now expand our tensor cur to be tensor<1x1x...xNx1x...x1xT>
474+ size_t dimensionInsertLoc = 0;
475+ for (size_t i = 0; i < rank; i++) {
476+ if (i == axis) {
477+ dimensionInsertLoc = i + 1;
478+ continue;
479+ }
480+ if (dimensionInsertLoc == 0) {
481+ curShape.insert(curShape.begin(), 1);
482+ } else {
483+ curShape.push_back(1);
484+ }
485+ 
486+ RankedTensorType curType = RankedTensorType::get(curShape, i32Type);
487+ indices = builder.create<triton::ExpandDimsOp>(loc, curType, indices,
488+ dimensionInsertLoc);
489+ }
490+ indices = builder.create<triton::BroadcastOp>(loc, finalIndicesType, indices);
491+ Value expandedInsertTensor = builder.create<triton::ExpandDimsOp>(loc, tensorToInsert, axis);
492+ SmallVector<int64_t> intermediateShape(smallShape);
493+ intermediateShape.insert(intermediateShape.begin() + axis, N / S);
494+ RankedTensorType intermediateType = RankedTensorType::get(intermediateShape, elementType);
495+ expandedInsertTensor = builder.create<triton::BroadcastOp>(loc, intermediateType, expandedInsertTensor);
496+ expandedInsertTensor = builder.create<triton::ReshapeOp>(loc, largeType, expandedInsertTensor);
497+ 
498+ Value name = builder.create<triton::GatherOp>(loc, expandedInsertTensor, indices, axis);
477 Value res = builder.create<arith::SelectOp>(loc, largeType, mask,499 Value res = builder.create<arith::SelectOp>(loc, largeType, mask,
478- reshapedTensor, large);500+ name, large);
479 return res;501 return res;
480}502}
481 503