[BACKEND] Support of ConvertLayoutOp from blocked to blocked and SliceLayout with blocked parent (#658)

python/src/triton.cc

@@ -19,6 +19,7 @@
 #include "triton/Dialect/TritonGPU/Transforms/Passes.h"
 #include "triton/Target/LLVMIR/LLVMIRTranslation.h"
 #include "triton/Target/PTX/PTXTranslation.h"
+#include "triton/tools/sys/getenv.hpp"
 
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
@@ -100,14 +101,6 @@ long pow2_divisor(long N) {
   return 1;
 }
 
-bool getBoolEnv(const std::string &env) {
-  const char *s = std::getenv(env.c_str());
-  std::string str(s ? s : "");
-  std::transform(str.begin(), str.end(), str.begin(),
-                 [](unsigned char c) { return std::tolower(c); });
-  return (str == "on" || str == "true" || str == "1");
-}
-
 // Returns something like "int16", whether dtype is a torch.dtype or
 // triton.language.dtype.
 std::string dtype_cache_key_part(const py::object &dtype) {
@@ -1635,7 +1628,7 @@ void init_triton_ir(py::module &&m) {
                  /*shouldPrintBeforePass=*/nullptr,
                  /*shouldPrintAfterPass=*/
                  [](mlir::Pass *pass, mlir::Operation *) {
-                   return getBoolEnv("MLIR_ENABLE_DUMP");
+                   return ::triton::tools::getBoolEnv("MLIR_ENABLE_DUMP");
                  },
                  /*printModuleScope=*/false,
                  /*printAfterOnlyOnChange=*/true,

python/tests/test_transpose.py

@@ -0,0 +1,68 @@
+import pytest
+import torch
+from torch.testing import assert_allclose
+
+import triton
+import triton.language as tl
+import triton.runtime as runtime
+
+
+@triton.jit
+def kernel(x_ptr, stride_xm,
+           z_ptr, stride_zn,
+           SIZE_M: tl.constexpr, SIZE_N: tl.constexpr):
+    off_m = tl.arange(0, SIZE_M)
+    off_n = tl.arange(0, SIZE_N)
+    Xs = x_ptr + off_m[:, None] * stride_xm + off_n[None, :] * 1
+    Zs = z_ptr + off_m[:, None] * 1 + off_n[None, :] * stride_zn
+    tl.store(Zs, tl.load(Xs))
+
+# These sizes cover the case of:
+# - blocked layout and sliced layout with block parent
+#  -- blocked layout in which sizePerThread/threadsPerWarp/warpsPerCTA
+#     need/need not to be wrapped
+#  -- sliced layout incase sizePerThread need to be wrapped
+#  -- different orders
+# - LayoutConversion from blocked -> blocked
+# - tt.Broadcast which requires for broadcast in either/both of
+#   CTA/perThread level
+
+# What is not covered and requires for TODO:
+# - vectorization load/store of shared memory
+# - multiple replication of layout conversion
+
+
+@pytest.mark.parametrize('NUM_WARPS,SIZE_M,SIZE_N', [
+    [1, 16, 16],
+    [1, 32, 32],
+    [1, 32, 64],
+    [2, 64, 128],
+    [2, 128, 64]
+])
+def test_convert_layout_impl(NUM_WARPS, SIZE_M, SIZE_N):
+    # TODO: this is to initialize the cuda context since it is not properly
+    #       dealed with in the existing runtime, remove this when the runtime
+    #       is updated
+    torch.zeros([10], device=torch.device('cuda'))
+    device = torch.cuda.current_device()
+    binary = runtime.build_kernel(kernel,
+                                  "*fp32,i32,*fp32,i32",
+                                  constants={"SIZE_M": SIZE_M,
+                                             "SIZE_N": SIZE_N},
+                                  num_warps=NUM_WARPS,
+                                  num_stages=3)
+    grid = lambda META: (1, )
+
+    x = torch.randn((SIZE_M, SIZE_N), device='cuda', dtype=torch.float32)
+    z = torch.empty((SIZE_N, SIZE_M), device=x.device, dtype=x.dtype)
+    runtime.launch_kernel(kernel=binary,
+                          device=device,
+                          grid=grid,
+                          x_ptr=x,
+                          stride_xm=x.stride(0),
+                          z_ptr=z,
+                          stride_zn=z.stride(0),
+                          SIZE_M=tl.constexpr(SIZE_M),
+                          SIZE_N=tl.constexpr(SIZE_N))
+    golden_z = torch.t(x)
+    assert_allclose(z, golden_z, rtol=1e-7, atol=1e-7)

python/tests/test_vecadd_no_scf.py

@@ -48,6 +48,7 @@ def vecadd_no_scf_tester(num_warps, block_size):
 
 
 def test_vecadd_no_scf():
+    vecadd_no_scf_tester(num_warps=4, block_size=256)
     vecadd_no_scf_tester(num_warps=2, block_size=256)
     vecadd_no_scf_tester(num_warps=1, block_size=256)