Fixed permutation propagation during transpose (#297)

Author: enjustli

include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp

@@ -113,18 +113,29 @@ static SmallVector<utils::IteratorType> getNParallelLoopsAttrs(unsigned n) {
   return SmallVector<utils::IteratorType>(n, utils::IteratorType::parallel);
 }
 
+// if order is empty, transpose the last two dimensions
+// otherwise, use the provided order. 
+// The order must be a permutation of the source rank.
 static Value getTransposedValue(Value source, const Location loc,
-                                ConversionPatternRewriter &rewriter) {
-
+                                ConversionPatternRewriter &rewriter,
+                                llvm::ArrayRef<int32_t> order = {}) {
   auto sourceType = cast<RankedTensorType>(source.getType());
   auto sourceRank = sourceType.getRank();
 
   SmallVector<int64_t> perm(sourceRank);
-  std::iota(std::begin(perm), std::end(perm), 0);
-  std::swap(perm[sourceRank - 1], perm[sourceRank - 2]);
-
   SmallVector<int64_t> transposedShape(sourceType.getShape());
-  std::swap(transposedShape[sourceRank - 1], transposedShape[sourceRank - 2]);
+  if (order.empty()) {
+    std::iota(std::begin(perm), std::end(perm), 0);
+    std::swap(perm[sourceRank - 1], perm[sourceRank - 2]);
+    std::swap(transposedShape[sourceRank - 1], transposedShape[sourceRank - 2]);
+  } else {
+    // Use the provided order
+    assert(order.size() == sourceRank && "Order size must match source rank");
+    for (unsigned i = 0; i < sourceRank; ++i) {
+      perm[i] = order[i];
+      transposedShape[i] = sourceType.getShape()[order[i]];
+    }
+  }
 
   Value transposeInit = rewriter.create<tensor::EmptyOp>(
       loc, transposedShape, sourceType.getElementType());
@@ -769,11 +780,8 @@ struct TransposeConverter : public OpConversionPattern<triton::TransOp> {
   LogicalResult
   matchAndRewrite(triton::TransOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto src = adaptor.getSrc();
-    auto srcRank = cast<ShapedType>(src.getType()).getRank();
-    assert(srcRank == 2 && "only expect transposing 2D data");
-
-    auto res = getTransposedValue(src, op.getLoc(), rewriter);
+    auto res = getTransposedValue(adaptor.getSrc(), op.getLoc(), rewriter,
+                                  op.getOrder());
     rewriter.replaceOp(op, res);
     return success();
   }

python/examples/conftest.py

@@ -10,13 +10,27 @@
 def empty_decorator(func):
     return func
 
+
 pytest.mark.interpreter = empty_decorator
 
 
 @pytest.fixture
 def device(request):
     return "cpu"
 
+# this fixture is used for test_trans_4d && test_trans_reshape
+@pytest.fixture
+def with_allocator():
+    import triton
+    from triton.runtime._allocation import NullAllocator
+    from triton._internal_testing import default_alloc_fn
+
+    triton.set_allocator(default_alloc_fn)
+    try:
+        yield
+    finally:
+        triton.set_allocator(NullAllocator())
+
 
 tests_supported = {
     "test_store_eviction_policy",
@@ -56,9 +70,12 @@ def device(request):
     "test_load_cache_modifier",
     "test_dot_without_load",
     "test_cat",
-    "test_addptr"
+    "test_addptr",
+    "test_transpose",
+    "test_trans_4d",
 }
 
+
 def pytest_collection_modifyitems(config, items):
     skip_marker = pytest.mark.skip(reason="CPU backend does not support it yet")
     # There is a dependency issue on build machine which breaks bfloat16