[Refactor] move buffer_ops to kernels.common; remove expr/vector.py

[sjfeng1999]

Motivation

close #641

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

This PR refactors FlyDSL’s low-level helpers by removing the legacy flydsl.expr.vector module and shifting buffer load/store utilities away from flydsl.expr.buffer_ops toward kernels.common.buffer_ops, while updating kernels/tests/docs to use the new import locations and as_ir_value conversions for raw MLIR builders.

Changes:

• Removed python/flydsl/expr/vector.py and updated call sites to use flydsl._mlir.dialects.vector plus flydsl.expr.typing re-exports (e.g., Vector, ReductionOp, full).
• Retargeted buffer ops usage from flydsl.expr.buffer_ops to kernels.common.buffer_ops and simplified many tests to use Out[idx] = value.
• Updated documentation to reflect the new recommended imports and module layout.

Reviewed changes

Copilot reviewed 69 out of 69 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/unit/test_tdm_mcast_add_gfx1250.py	Switches vector ops usage to MLIR vector dialect and as_ir_value.
tests/unit/test_for_auto_iter_args.py	Replaces explicit buffer stores with tensor indexing writes.
tests/unit/test_expr_optional_rocdl.py	Updates optional-import expectations after removing buffer_ops from flydsl.expr lazy modules.
tests/system/test_while_e2e.py	Replaces buffer_ops stores with Out[...] assignments in while-loop tests.
tests/system/test_inline_compare_scf_if.py	Updates example/test kernel to use Out[gid] = ... instead of buffer_ops.buffer_store.
tests/system/test_ifexp_e2e.py	Replaces buffer_ops stores with Out[0] = ... across if-expression tests.
tests/system/test_if_liveout_minimal.py	Replaces buffer_ops usage with tensor indexing store.
tests/system/test_for_vector_carry_shape_e2e.py	Updates reduction enum usage and replaces buffer store with Out[...] assignment.
tests/system/test_for_auto_iter_args_e2e.py	Removes buffer_ops usage in favor of Out[...] assignments across e2e kernels.
tests/system/test_closure_freevars_mismatch.py	Removes buffer_ops import/usage and uses Out[...] assignment.
tests/perf/bench_tdm_bandwidth_gfx1250.py	Switches vector loads to MLIR vector.load plus as_ir_value.
tests/kernels/test_quant.py	Moves vector-related imports from flydsl.expr.vector to flydsl.expr.typing.
tests/kernels/test_fused_rope_cache.py	Updates comment wording to reflect new raw-store approach.
python/flydsl/expr/vector.py	Removes legacy vector dialect helper module.
python/flydsl/expr/rocdl/tdm_ops.py	Switches to MLIR vector dialect and ensures operands are raw via as_ir_value.
python/flydsl/expr/init.py	Removes vector import and drops lazy buffer_ops module export.
kernels/norm/softmax_kernel.py	Updates ReductionOp/full imports to come from flydsl.expr.typing.
kernels/norm/rmsnorm_kernel.py	Updates ReductionOp/full imports to come from flydsl.expr.typing.
kernels/norm/rmsnorm_common.py	Updates full import to come from flydsl.expr.typing.
kernels/norm/layernorm_kernel.py	Updates ReductionOp/full imports to come from flydsl.expr.typing.
kernels/moe/topk_gating_softmax_kernel.py	Switches to MLIR vector dialect and uses as_ir_value for extract operands.
kernels/moe/silu_and_mul_fq.py	Switches to MLIR vector dialect, uses kernels.common.buffer_ops, and adds as_ir_value conversions.
kernels/moe/moe_sorting_kernel.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/moe/moe_gemm_2stage/reduction.py	Retargets buffer ops import to kernels.common.buffer_ops and switches to MLIR vector.
kernels/moe/moe_gemm_2stage/gemm2.py	Retargets buffer ops import to kernels.common.buffer_ops and adds as_ir_value conversions for vector ops.
kernels/moe/moe_gemm_2stage/gemm1.py	Retargets buffer ops import to kernels.common.buffer_ops and adds as_ir_value conversions for vector ops.
kernels/moe/moe_gemm_2stage_wmma_gfx1250.py	Switches to MLIR vector dialect and kernels.common.buffer_ops, adding as_ir_value conversions.
kernels/moe/moe_gemm_2stage_mxscale_gfx1250.py	Switches to MLIR vector dialect and kernels.common.buffer_ops, adding as_ir_value conversions.
kernels/moe/moe_common.py	Switches to MLIR vector dialect and adds as_ir_value conversions for bitcast/from_elements helpers.
kernels/moe/moe_blockscale_2stage/reduction.py	Retargets buffer ops import to kernels.common.buffer_ops and switches to MLIR vector.
kernels/moe/moe_blockscale_2stage/gemm2.py	Retargets buffer ops import to kernels.common.buffer_ops and adds as_ir_value conversions for vector ops.
kernels/moe/moe_blockscale_2stage/gemm1.py	Retargets buffer ops import to kernels.common.buffer_ops and adds as_ir_value conversions for vector ops.
kernels/moe/mixed_moe_gemm_2stage/gemm2.py	Retargets buffer ops import to kernels.common.buffer_ops and adds as_ir_value conversions for vector ops.
kernels/moe/mixed_moe_gemm_2stage/gemm1.py	Retargets buffer ops import to kernels.common.buffer_ops and adds as_ir_value conversions for vector ops.
kernels/gemm/wmma_gemm_gfx1250.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/gemm/splitk_hgemm.py	Switches to MLIR vector dialect and retargets buffer ops import to kernels.common.buffer_ops.
kernels/gemm/small_m_hgemm.py	Switches to MLIR vector dialect and adds as_ir_value conversions for vector operations.
kernels/gemm/rdna3_f16_gemm.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/gemm/rdna_fp8_preshuffle_gemm.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/gemm/rdna_f16_gemm.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/gemm/preshuffle_gemm.py	Switches to MLIR vector dialect, retargets buffer ops usage to kernels.common.buffer_ops, and adds as_ir_value conversions.
kernels/gemm/hgemm_splitk.py	Switches to MLIR vector dialect, retargets buffer ops import to kernels.common.buffer_ops, and adds as_ir_value conversions.
kernels/gemm/gemm_fp8fp4_gfx1250.py	Retargets buffer ops import to kernels.common.buffer_ops and replaces arith.unwrap uses with as_ir_value in yields.
kernels/gemm/gemm_common_gfx1250.py	Switches to MLIR vector dialect and replaces arith.unwrap with as_ir_value in LDS helpers.
kernels/gemm/fp8_gemm_4wave.py	Switches vector import to MLIR dialect and ensures bitcast inputs are raw via as_ir_value.
kernels/gemm/fp4_gemm_4wave.py	Retargets buffer ops alias import to kernels.common.buffer_ops.
kernels/gemm/blockscale_preshuffle_gemm.py	Switches to MLIR vector dialect, retargets buffer ops import to kernels.common.buffer_ops, and uses as_ir_value for stores/loads.
kernels/conv/conv3d_implicit_8wave.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/conv/conv3d_implicit_8wave_fp8.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/common/tensor_shim.py	Switches to MLIR vector dialect and retargets buffer ops import to kernels.common.buffer_ops.
kernels/common/mem_ops.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/common/buffer_ops.py	Hosts the buffer ops implementation under kernels.common and updates imports accordingly.
kernels/comm/flydsl_dispatch_combine_intranode_kernel.py	Retargets buffer ops imports to kernels.common.buffer_ops and replaces address unwrapping with as_ir_value.
kernels/comm/custom_all_reduce_kernel.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/attention/qk_norm_rope_quant.py	Switches to MLIR vector dialect, retargets buffer ops import to kernels.common.buffer_ops, and updates ReductionOp import.
kernels/attention/pa_metadata.py	Switches to MLIR vector dialect, retargets buffer ops import to kernels.common.buffer_ops, and adds as_ir_value conversions.
kernels/attention/pa_decode_swa.py	Switches to MLIR vector dialect, retargets buffer ops import to kernels.common.buffer_ops, and adds as_ir_value conversions.
kernels/attention/pa_decode_fp8.py	Switches to MLIR vector dialect, retargets buffer ops import to kernels.common.buffer_ops, and adds as_ir_value conversions.
kernels/attention/pa_common.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/attention/fused_rope_cache_kernel.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/attention/flash_attn_gfx950.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/attention/flash_attn_generic.py	Retargets buffer ops import to kernels.common.buffer_ops.
kernels/attention/flash_attn_fp8_gfx950.py	Retargets buffer ops import to kernels.common.buffer_ops.
docs/prebuilt_kernels_guide.md	Updates example imports/commentary for the refactor (needs clarification).
docs/kernel_authoring_guide.md	Removes the fx.buffer_ops section and updates module list references after the refactor.
docs/architecture_guide.md	Updates repository structure docs to remove expr/vector.py and expr/buffer_ops.py entries.
docs/api/dsl.rst	Removes the flydsl.expr.buffer_ops API section from the DSL docs.
docs/api/compiler.rst	Removes the flydsl.expr.buffer_ops section from compiler docs.

---

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