[Bugfix] Open Rocm ci test and fix some bugs.

.github/workflows/ci.yml

@@ -373,6 +373,7 @@ jobs:
             ./python
 
       # AMD ROCm tests
+      # runtime and transform tests needs to repair, then rm it from ignore list
       - name: Run ROCm tests with Python ${{ matrix.python-version }} (${{ matrix.runner.toolkit }})
         id: rocm-tests
         if: contains(matrix.runner.toolkit, 'ROCm')
@@ -383,7 +384,8 @@ jobs:
             pytest --verbose --color=yes --durations=0 --showlocals --cache-clear
           )
           "${PYTEST[@]}" --maxfail=3 --numprocesses=4 \
-            ./python/amd
+            --ignore=./python/runtime --ignore=./python/transform \
+            ./python
 
       # Apple Metal tests
       - name: Run Metal tests with Python ${{ matrix.python-version }} (${{ matrix.runner.toolkit }})

src/op/gemm.cc

@@ -153,6 +153,8 @@ std::pair<int, int> GemmWarpPolicyNode::computeWarpPartition(
   int kNPerWarp = 8;            // Columns processed by a single warp
   if (TargetIsVolta(target)) {
     kNPerWarp = 16;
+  } else if (TargetIsCDNA(target)) {
+    kNPerWarp = 16;
   }
   ICHECK(M % kMPerWarp == 0)
       << "M must be divisible by " << kMPerWarp << ", but got " << M;

src/op/logical.cc

@@ -42,14 +42,16 @@ TVM_REGISTER_OP("tl.any_of")
     .set_attr<TCallEffectKind>("TCallEffectKind",
                                Integer(CallEffectKind::kPure))
     .set_attr<TScriptPrinterName>("TScriptPrinterName", "any_of")
-    .set_attr<FLowerIntrinsic>("cuda.FLowerIntrinsic", any_of_op);
+    .set_attr<FLowerIntrinsic>("cuda.FLowerIntrinsic", any_of_op)
+    .set_attr<FLowerIntrinsic>("hip.FLowerIntrinsic", any_of_op);
 
 TVM_REGISTER_OP("tl.all_of")
     .set_num_inputs(1)
     .set_attr<TCallEffectKind>("TCallEffectKind",
                                Integer(CallEffectKind::kPure))
     .set_attr<TScriptPrinterName>("TScriptPrinterName", "all_of")
-    .set_attr<FLowerIntrinsic>("cuda.FLowerIntrinsic", all_of_op);
+    .set_attr<FLowerIntrinsic>("cuda.FLowerIntrinsic", all_of_op)
+    .set_attr<FLowerIntrinsic>("hip.FLowerIntrinsic", all_of_op);
 
 } // namespace tl
 } // namespace tvm

src/target/codegen_hip.cc

@@ -767,24 +767,26 @@ void CodeGenTileLangHIP::VisitExpr_(const CallNode *op, std::ostream &os) {
     }
     this->stream << ");\n";
   };
-  if (op->op.same_as(builtin::ptx_cp_async())) {
+  if (op->op.same_as(builtin::ptx_cp_async()) ||
+      op->op.same_as(tl::ptx_cp_async())) {
+    // args[0] = dst_access_ptr, args[1] = src_access_ptr, args[2] = bytes,
+    // args[3] = predicate (optional)
+    ICHECK(op->args.size() == 3 || op->args.size() == 4)
+        << "ptx_cp_async expects 3 or 4 arguments (dst_access_ptr, "
+           "src_access_ptr, bytes, [predicate])";
     std::string dst = this->PrintExpr(op->args[0]);
-    std::string dst_offset = this->PrintExpr(op->args[1]);
-    std::string src = this->PrintExpr(op->args[2]);
-    std::string src_offset = this->PrintExpr(op->args[3]);
-    std::string size = this->PrintExpr(op->args[4]);
-    // use size of argument list to indicate whether or not to use predicated
-    // cp.async
-    if (op->args.size() == 5) {
-      this->PrintIndent();
-      this->stream << "tl::cp_async_gs<" << size << ">(" << dst << "+"
-                   << dst_offset << ", " << src << "+" << src_offset << ");\n";
+    std::string src = this->PrintExpr(op->args[1]);
+    std::string size = this->PrintExpr(op->args[2]);
+    this->PrintIndent();
+    if (op->args.size() == 3) {
+      // Non-predicated version
+      this->stream << "tl::cp_async_gs<" << size << ">(" << dst << ", " << src
+                   << ");\n";
     } else {
-      std::string condition = this->PrintExpr(op->args[5]);
-      this->PrintIndent();
+      // Predicated version
+      std::string condition = this->PrintExpr(op->args[3]);
       this->stream << "tl::cp_async_gs_conditional<" << size << ">(" << dst
-                   << "+" << dst_offset << ", " << src << "+" << src_offset
-                   << ", " << condition << ");\n";
+                   << ", " << src << ", " << condition << ");\n";
     }
   } else if (op->op.same_as(builtin::ptx_commit_group())) {
     print_extern_call_stmt("tl::cp_async_commit");
@@ -942,6 +944,8 @@ void CodeGenTileLangHIP::VisitExpr_(const CallNode *op, std::ostream &os) {
         {"float32x4", "float32x4"},
         {"float8_e4m3fnuzx4", "fp8_e4_4_t"},
         {"float8_e4m3fnuzx8", "long"},
+        {"float8_e5m2fnuzx4", "fp8_e5_4_t"},
+        {"float8_e5m2fnuzx8", "long"},
         {"float32x16", "float32x16"}};
     std::string call_mfma_code = R"({
       *((({C_dtype}*){c_ref}) + {c_bias}) = {mfma_buildin}(*((({A_dtype}*){a_ref}) + {a_bias}),
@@ -980,6 +984,16 @@ void CodeGenTileLangHIP::VisitExpr_(const CallNode *op, std::ostream &os) {
     // HIP doesn't need explicit register management like CUDA
     // This is a no-op for HIP
     return;
+  } else if (op->op.same_as(tl::warp_reduce_sum())) {
+    os << "tl::warp_reduce_sum(" << PrintExpr(op->args[0]) << ")";
+  } else if (op->op.same_as(tl::warp_reduce_max())) {
+    os << "tl::warp_reduce_max(" << PrintExpr(op->args[0]) << ")";
+  } else if (op->op.same_as(tl::warp_reduce_min())) {
+    os << "tl::warp_reduce_min(" << PrintExpr(op->args[0]) << ")";
+  } else if (op->op.same_as(tl::warp_reduce_bitand())) {
+    os << "tl::warp_reduce_bitand(" << PrintExpr(op->args[0]) << ")";
+  } else if (op->op.same_as(tl::warp_reduce_bitor())) {
+    os << "tl::warp_reduce_bitor(" << PrintExpr(op->args[0]) << ")";
   } else {
     CodeGenC::VisitExpr_(op, os);
   }

src/tl_templates/cuda/reduce.h

@@ -175,15 +175,15 @@ template <int threads, int Axis = 0, bool reverse = false> struct CumSum2D {
   static_assert(threads == 1024 or threads == 512 or threads == 256 or
                 threads == 128 or threads == 64 or threads == 32);
   template <typename T, int SEG = 32>
-  static TL_DEVICE T run(const T *__restrict__ src, T *__restrict__ dst, int H,
-                         int W) {
+  static TL_DEVICE void run(const T *__restrict__ src, T *__restrict__ dst,
+                            int H, int W) {
 
     constexpr int TILE_H = threads / SEG;
     constexpr unsigned MASK = 0xffffffff;
     const int num_blocks = (H + TILE_H - 1) / TILE_H;
     const int tid = threadIdx.x;
-    const int lane = tid % 32;
-    const int row = tid / 32;
+    const int lane = tid % SEG;
+    const int row = tid / SEG;
 
     for (int b = 0; b < num_blocks; ++b) {
       const int gRow = b * TILE_H + row;

src/tl_templates/hip/atomic.h

@@ -0,0 +1,104 @@
+#pragma once
+
+#include <hip/hip_runtime.h>
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+template <typename T1, typename T2>
+__forceinline__ __device__ void AtomicAdd(T1 *address, T2 val,
+                                          int memory_order = 0) {
+  atomicAdd(reinterpret_cast<T1 *>(address), static_cast<T1>(val));
+}
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+// Overload for when the first argument is a value instead of a pointer
+template <typename T1, typename T2>
+__forceinline__ __device__ void AtomicAdd(T1 &address, T2 val,
+                                          int memory_order = 0) {
+  atomicAdd(reinterpret_cast<T1 *>(&address), static_cast<T1>(val));
+}
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+template <typename T1, typename T2>
+__forceinline__ __device__ T1 AtomicAddRet(T1 *ref, T2 val,
+                                           int memory_order = 0) {
+  return atomicAdd(ref, static_cast<T1>(val));
+}
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+template <typename T1, typename T2>
+__forceinline__ __device__ void AtomicMax(T1 *address, T2 val,
+                                          int memory_order = 0) {
+  atomicMax(reinterpret_cast<T1 *>(address), static_cast<T1>(val));
+}
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+// Overload for when the first argument is a value instead of a pointer
+template <typename T1, typename T2>
+__forceinline__ __device__ void AtomicMax(T1 &address, T2 val,
+                                          int memory_order = 0) {
+  atomicMax(reinterpret_cast<T1 *>(&address), static_cast<T1>(val));
+}
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+template <typename T1, typename T2>
+__forceinline__ __device__ void AtomicMin(T1 *address, T2 val,
+                                          int memory_order = 0) {
+  atomicMin(reinterpret_cast<T1 *>(address), static_cast<T1>(val));
+}
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+// Overload for when the first argument is a value instead of a pointer
+template <typename T1, typename T2>
+__forceinline__ __device__ void AtomicMin(T1 &address, T2 val,
+                                          int memory_order = 0) {
+  atomicMin(reinterpret_cast<T1 *>(&address), static_cast<T1>(val));
+}
+
+__forceinline__ __device__ void AtomicAddx2(float *ref, float *val,
+                                            int memory_order = 0) {
+  float2 add_val = *reinterpret_cast<float2 *>(val);
+  atomicAdd(ref + 0, add_val.x);
+  atomicAdd(ref + 1, add_val.y);
+}
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+__forceinline__ __device__ float2 AtomicAddx2Ret(float *ref, float *val,
+                                                 int memory_order = 0) {
+  float2 add_val = *reinterpret_cast<float2 *>(val);
+  float2 ret;
+  ret.x = atomicAdd(ref + 0, add_val.x);
+  ret.y = atomicAdd(ref + 1, add_val.y);
+  return ret;
+}
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+__forceinline__ __device__ void AtomicAddx4(float *ref, float *val,
+                                            int memory_order = 0) {
+  float4 add_val = *reinterpret_cast<float4 *>(val);
+  atomicAdd(ref + 0, add_val.x);
+  atomicAdd(ref + 1, add_val.y);
+  atomicAdd(ref + 2, add_val.z);
+  atomicAdd(ref + 3, add_val.w);
+}
+
+// Add an extra unused input to accommodate the additional 'memory_order'
+// argument during lowering.
+__forceinline__ __device__ float4 AtomicAddx4Ret(float *ref, float *val,
+                                                 int memory_order = 0) {
+  float4 add_val = *reinterpret_cast<float4 *>(val);
+  float4 ret;
+  ret.x = atomicAdd(ref + 0, add_val.x);
+  ret.y = atomicAdd(ref + 1, add_val.y);
+  ret.z = atomicAdd(ref + 2, add_val.z);
+  ret.w = atomicAdd(ref + 3, add_val.w);
+  return ret;
+}

src/tl_templates/hip/common.h

@@ -1,6 +1,8 @@
 #pragma once
 
+#include "atomic.h"
 #include <ck_tile/core.hpp>
+#include <hip/amd_detail/amd_warp_functions.h>
 #include <hip/hip_bf16.h>
 #include <hip/hip_fp16.h>
 #include <hip/hip_runtime.h>
@@ -105,18 +107,94 @@ TL_DEVICE unsigned __pack_bfloat162(const bfloat16_t x, const bfloat16_t y) {
   return (v1 << 16) | v0;
 }
 
-template <typename T1, typename T2>
-TL_DEVICE void AtomicAdd(T1 *address, T2 val) {
-  atomicAdd(reinterpret_cast<T1 *>(address), static_cast<T1>(val));
+namespace tl {
+
+// Any
+template <typename T> TL_DEVICE bool Any(T *a, int size) {
+  for (int i = 0; i < size; i++) {
+    if (a[i]) {
+      return true;
+    }
+  }
+  return false;
+}
+
+// All
+template <typename T> TL_DEVICE bool All(T *a, int size) {
+  for (int i = 0; i < size; i++) {
+    if (!a[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// TODO(gong): support shfl_sync(rocm 7.1.1 provide shfl_sync)
+// shfl_sync func
+template <typename T> TL_DEVICE T shfl_xor(T val, int delta) {
+  return __shfl_xor(val, delta);
+}
+
+template <typename T> TL_DEVICE T shfl_down(T val, int delta) {
+  return __shfl_down(val, delta);
+}
+
+template <typename T> TL_DEVICE T shfl_up(T val, int delta) {
+  return __shfl_up(val, delta);
+}
+
+template <typename T> TL_DEVICE T shfl(T val, int srcLane) {
+  return __shfl(val, srcLane);
+}
+
+// specialize half_t
+template <> TL_DEVICE half_t shfl_xor(half_t val, int delta) {
+  float f = static_cast<float>(val);
+  float r = __shfl_xor(f, delta);
+  return half_t(r);
+}
+
+template <> TL_DEVICE half_t shfl_down(half_t val, int delta) {
+  float f = static_cast<float>(val);
+  float r = __shfl_down(f, delta);
+  return half_t(r);
+}
+
+template <> TL_DEVICE half_t shfl_up(half_t val, int delta) {
+  float f = static_cast<float>(val);
+  float r = __shfl_up(f, delta);
+  return half_t(r);
+}
+
+template <> TL_DEVICE half_t shfl(half_t val, int srcLane) {
+  float f = static_cast<float>(val);
+  float r = __shfl(f, srcLane);
+  return half_t(r);
+}
+
+// specialize bfloat16_t
+template <> TL_DEVICE bfloat16_t shfl_xor(bfloat16_t val, int laneMask) {
+  float f = static_cast<float>(val);
+  float r = __shfl_xor(f, laneMask);
+  return bfloat16_t(r);
 }
 
-// Overload for when the first argument is a value instead of a pointer
-template <typename T1, typename T2>
-TL_DEVICE void AtomicAdd(T1 address, T2 val) {
-  atomicAdd(reinterpret_cast<T1 *>(&address), static_cast<T1>(val));
+template <> TL_DEVICE bfloat16_t shfl_down(bfloat16_t val, int delta) {
+  float f = static_cast<float>(val);
+  float r = __shfl_down(f, delta);
+  return bfloat16_t(r);
 }
 
-template <typename T1, typename T2>
-TL_DEVICE T1 AtomicAddRet(T1 *address, T2 val) {
-  return atomicAdd(reinterpret_cast<T1 *>(address), static_cast<T1>(val));
+template <> TL_DEVICE bfloat16_t shfl_up(bfloat16_t val, int delta) {
+  float f = static_cast<float>(val);
+  float r = __shfl_up(f, delta);
+  return bfloat16_t(r);
 }
+
+template <> TL_DEVICE bfloat16_t shfl(bfloat16_t val, int srcLane) {
+  float f = static_cast<float>(val);
+  float r = __shfl(f, srcLane);
+  return bfloat16_t(r);
+}
+
+} // namespace tl