[OPT] Add fly-int-swizzle-simplify pass

[sjfeng1999]

Add an algebraic simplification pass that recognizes the canonical three-instruction swizzle sequence (andi/shrui/xori) emitted by applySwizzle and peels period-aligned addends out of the swizzle:

swizzle(base + d)  →  swizzle(base) + d   when d % period == 0

In example04 (preshuffle_gemm) with SwizzleType S<3,3,3> (period=512), the LDS address for each pipeline stage is swizzle(thread_offset + stage * lds_stride). Since lds_stride (BLOCK_M × BLOCK_K × sizeof(f16) = 16384) is a multiple of 512, the pass peels the stage offset out:

Before: each stage recomputes the full swizzle (andi+shrui+xori) on
(thread_offset + stage_offset), duplicating 3 ALU ops per stage.
After: swizzle is computed once on thread_offset; stage offsets become
plain arith.addi with an immediate, which then fold/CSE across
stages.

Concrete effects on the ds_read/ds_write address computation in the GEMM hot loop:

• Instruction count: eliminates redundant andi+shrui+xori triplets per pipeline stage (−6 SALU/VALU ops for 2-stage, −9 for 3-stage).
• Immediate offsets: the peeled constant addends enable ds_read_b128 with immediate byte offsets instead of a full v_add + swizzle, which the ISA can encode directly in the ds instruction offset field.
• Register pressure: fewer live intermediates (%masked, %shifted) per stage reduces VGPR demand in the address-computation section.

The pass also includes a divisibility lattice that propagates alignment info both forward (through arith add/mul/shl) and backward (from IntAttr.divisibility annotations on Make*Op results), enabling the optimization even for dynamic strides with known alignment.

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

Adds an MLIR optimization pass (fly-int-swizzle-simplify) to algebraically simplify the canonical swizzle sequence emitted by applySwizzle, peeling period-aligned addends out of the swizzle input to reduce redundant ALU ops and improve downstream folding/CSE, and wires it into the ROCm backend pipeline.

Changes:

• Introduce fly-int-swizzle-simplify pass with a divisibility lattice to detect period-aligned addends and rewrite swizzle-shaped andi/shrui/xori patterns.
• Register/build the new transform (CMake + Passes.td) and enable it in the ROCm compilation pipeline.
• Add MLIR tests covering peel/no-peel behavior for i32 and i64.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
lib/Dialect/Fly/Transforms/IntSwizzleSimplify.cpp	Implements the new swizzle simplification pass and divisibility propagation.
include/flydsl/Dialect/Fly/Transforms/Passes.td	Registers the new pass and documents its intent/pattern.
python/flydsl/compiler/backends/rocm.py	Inserts the pass into the ROCm lowering pipeline.
tests/mlir/Transforms/int_swizzle_simplify.mlir	Adds FileCheck coverage for the optimization and non-optimization cases.
lib/Dialect/Fly/CMakeLists.txt	Builds the new transform source file into the Fly dialect library.
include/flydsl/Dialect/Fly/Utils/IntUtils.h / lib/Dialect/Fly/Utils/IntUtils.cpp	Adds an isDivisibleBy helper for IntAttr.
include/flydsl/Dialect/Fly/IR/FlyAttrDefs.td	Adds SwizzleAttr::period() convenience method.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.