Add support for permute operation to Direct Python Bindings

[rdspring1]

This PR adds support for replacing linear layers with TensorParallel NvFuser layer in deepseek model using Direct Python Bindings.

PR Stack:

• Add Multi-Gpu Support to Direct Python Bindings #4689
• Add support for cast operations to Direct Python Bindings #4697
• Add support for matmul and linear ops to Direct Python Bindings #4698
• Add size, shape, define_vector, and reshape ops to Direct Python Bindings #4704
• Add support for permute operation to Direct Python Bindings #4701 <<< This PR.
• Transition test_dtensor.py and test_deepseek_v3.py to direct bindings  #4809

[github-actions]

Review updated until commit b6004e2

Description

• Added support for permute operation in Direct Python Bindings.

• Implemented permute handling in PythonTranslator.

• Updated opinfo to support direct bindings for permute.

• Adjusted error handling in test_direct_ops.

---

Changes walkthrough 📝

Relevant files

Enhancement

ops.cpp Add permute operation                                                                        python/python_direct/ops.cpp Added permute function to ops module. Included error checking for permute dimensions. +26/-0    python_translate.cpp Implement permute handling in PythonTranslator                      python/python_direct/python_translate.cpp Implemented handlePermute method. Updated handle method to call handlePermute for permutations. Added error check for input and output dtype consistency. +23/-2    opinfos.py Update opinfo for permute                                                                tests/python/opinfo/opinfos.py Added supports_direct_bindings for permute operation. +1/-0      test_direct_ops.py Adjust error handling in tests                                                      tests/python/opinfo/test_direct_ops.py Temporarily skipped regex check in error handling. +2/-3

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Permutation Handling The handling of permutations in handlePermute should be validated to ensure it correctly identifies and translates permutations without introducing errors. void handle(const LoadStoreOp* lsop) final { // TODO short-circuit: lsop is a permutation. if (lsop->out()->isA<TensorView>() && lsop->out()->as<TensorView>()->hasRoot()) { return handlePermute(lsop); } NVF_ERROR( lsop->in()->dtype() == lsop->out()->dtype(), "Expected the dtype for input and output to be the same"); visited_vals_.insert(lsop->out()); static const std::vector<std::string> argument_names = {"dtype"}; printer_.generateKwargsOperation( "fd.ops.cast", std::make_tuple(lsop->in()), argument_names, std::make_tuple(lsop->out()->dtype()), {lsop->out()}); } void handlePermute(const LoadStoreOp* lsop) { TensorView* out_tv = lsop->out()->as<TensorView>(); std::optional<std::vector<int64_t>> new2old = ir_utils::computePermutation( out_tv->getRootDomain(), out_tv->getLogicalDomain()); NVF_ERROR(new2old.has_value(), "Expected permutation"); visited_vals_.insert(lsop->out()); static const std::vector<std::string> argument_names = {"dims"}; printer_.generateKwargsOperation( "fd.ops.permute", std::make_tuple(lsop->in()), argument_names, std::make_tuple(new2old.value()), {lsop->out()}); } Error Handling The skipped regex check in test_errors should be reviewed to ensure that the error messages from direct bindings are appropriately handled and do not mask potential issues. with pytest.raises(exception_type): errors_test_fn(op, sample)

[rdspring1]

!test

[rdspring1]

!build