/*!
 * \file target/codegen.cc
 */

#include "codegen_hip.h"
#include "support/check.h"
#include <tvm/arith/analyzer.h>
#include <tvm/ir/cast.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/index_map.h>
#include <tvm/tirx/op.h>

#include <cmath>
#include <cstdint>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "backend/common/target_utils.h"
#include "op/builtin.h"

namespace tvm {
namespace codegen {

using namespace ffi;

namespace {

bool IsValidCPAsyncTransferBytes(int64_t bytes) {
  return bytes == 4 || bytes == 8 || bytes == 16;
}

std::optional<DataType> GetAccessPtrElementType(const PrimExpr &expr) {
  const auto *ptr_call = expr.as<CallNode>();
  if (ptr_call == nullptr) {
    return std::nullopt;
  }
  if (ptr_call->op.same_as(builtin::address_of())) {
    const auto *buffer_load = ptr_call->args[0].as<BufferLoadNode>();
    ICHECK(buffer_load) << "address_of arg must be BufferLoad";
    return buffer_load->buffer->dtype;
  }
  if (ptr_call->op.same_as(builtin::tvm_access_ptr())) {
    ICHECK(!ptr_call->args.empty());
    return ptr_call->args[0].dtype();
  }
  if (ptr_call->op.same_as(tl::access_ptr())) {
    ICHECK_EQ(ptr_call->args.size(), 3U)
        << "tl.access_ptr expects 3 args: (BufferLoad, extent, rw_mask)";
    const auto *buffer_load = ptr_call->args[0].as<BufferLoadNode>();
    ICHECK(buffer_load) << "tl.access_ptr arg0 must be BufferLoad";
    return buffer_load->buffer->dtype;
  }
  return std::nullopt;
}

int GetTileLangCPAsyncTransferBytes(const CallNode *op) {
  ICHECK(op->args.size() == 3 || op->args.size() == 4)
      << "tl::ptx_cp_async expects 3 or 4 arguments (dst_access_ptr, "
         "src_access_ptr, num_elems, [predicate])";
  const auto *num_elems_imm = op->args[2].as<IntImmNode>();
  ICHECK(num_elems_imm) << "tl::ptx_cp_async num_elems must be IntImm, but got "
                        << op->args[2];
  int64_t num_elems = num_elems_imm->value;
  ICHECK_GT(num_elems, 0);

  auto dst_elem_type = GetAccessPtrElementType(op->args[0]);
  auto src_elem_type = GetAccessPtrElementType(op->args[1]);
  ICHECK(dst_elem_type.has_value() && src_elem_type.has_value())
      << "tl::ptx_cp_async expects address_of, tl.access_ptr, or "
         "tvm_access_ptr operands";

  int64_t dst_total_bits =
      num_elems * dst_elem_type.value().bits() * dst_elem_type.value().lanes();
  int64_t src_total_bits =
      num_elems * src_elem_type.value().bits() * src_elem_type.value().lanes();
  ICHECK_EQ(dst_total_bits, src_total_bits)
      << "tl::ptx_cp_async requires src/dst transfer widths to match, but got "
      << dst_total_bits << " vs " << src_total_bits << " bits";
  ICHECK_EQ(dst_total_bits % 8, 0)
      << "tl::ptx_cp_async requires byte-aligned transfers, but got "
      << dst_total_bits << " bits";

  int64_t total_bytes = dst_total_bits / 8;
  ICHECK(IsValidCPAsyncTransferBytes(total_bytes))
      << "tl::ptx_cp_async requires a final PTX byte width in {4, 8, 16}, but "
         "got "
      << total_bytes;
  return static_cast<int>(total_bytes);
}

} // namespace

static std::string GetFP8Type(DataType type) {
  std::stringstream stream;
  int32_t lanes = type.lanes();
  std::string vec;
  if (type.is_scalar()) {
    vec = "";
  } else if (lanes == 2) {
    vec = "_2";
  } else if (lanes == 4) {
    vec = "_4";
  } else if (lanes == 8) {
    vec = "_8";
  } else if (lanes == 16) {
    vec = "_16";
  } else {
    LOG(FATAL) << "Only support scalar and vector types of width (2, 4, 8, 16) "
                  "for FP8";
  }
  if (type.is_float8_e4m3fn() || type.is_float8_e4m3fnuz() ||
      type.is_float8_e4m3() || type.code() == DataType::kFloat8_e4m3b11fnuz) {
    stream << "fp8_e4" << vec << "_t";
  } else if (type.is_float8_e5m2() || type.is_float8_e5m2fnuz() ||
             type.code() == DataType::kFloat8_e5m2) {
    stream << "fp8_e5" << vec << "_t";
  } else if (type.code() == DataType::kFloat8_e8m0fnu) {
    stream << "fp8_e8" << vec << "_t";
  } else {
    LOG(FATAL) << "Unsupported FP8 type in HIP codegen: " << type;
  }
  return stream.str();
}

// Returns the HIP C type string for a float4_e2m1fn DataType.
// Only used on gfx950; caller sets enable_fp4_ before invoking.
static std::string GetFP4Type(DataType type) {
  std::stringstream stream;
  int32_t lanes = type.lanes();
  if (type.is_scalar()) {
    stream << "fp4_e2_t";
  } else if (lanes == 2) {
    stream << "fp4_e2_2_t";
  } else if (lanes == 4) {
    stream << "fp4_e2_4_t";
  } else if (lanes == 8) {
    stream << "fp4_e2_8_t";
  } else if (lanes == 16) {
    stream << "fp4_e2_16_t";
  } else if (lanes == 32) {
    stream << "fp4_e2_32_t";
  } else {
    LOG(FATAL) << "Only support scalar and vector types of width (2,4,8,16,32) "
                  "for FP4 on HIP";
  }
  return stream.str();
}

/*!
 * \brief Replace patterns with replacement strings.
 * \note should use std::format instead when codebase is ported to C++20.
 */
class Replacer {
public:
  void register_rule(const std::string &pattern,
                     const std::string &replacement) {
    _rules.emplace_back(pattern, replacement);
  }
  std::string rewrite(std::string str) {
    for (auto &&rule : _rules) {
      auto [pattern, replacement] = rule;
      size_t len = pattern.size();
      size_t new_len = replacement.size();
      size_t pos = str.find(pattern);
      while (pos != std::string::npos) {
        str = str.replace(pos, len, replacement);
        pos = str.find(pattern, pos + new_len);
      }
    }
    return str;
  }
  void empty_rules() { _rules.clear(); }

private:
  std::vector<std::pair<std::string, std::string>> _rules;
};

CodeGenTileLangHIP::CodeGenTileLangHIP() { restrict_keyword_ = "__restrict__"; }

void CodeGenTileLangHIP::PrintFuncPrefix(std::ostream &os) {
  os << "extern \"C\" __global__ ";
}

class LaunchConfigExtractor : public tirx::StmtVisitor {
private:
  void VisitStmt_(const AttrStmtNode *op) final {
    if (op->attr_key == tirx::attr::thread_extent) {
      IterVar iv = Downcast<IterVar>(op->node);
      if (iv->var->name_hint == "threadIdx.x" ||
          iv->thread_tag == "threadIdx.x") {
        threadIdx_x_ext = op->value;
      } else if (iv->var->name_hint == "threadIdx.y" ||
                 iv->thread_tag == "threadIdx.y") {
        threadIdx_y_ext = op->value;
      } else if (iv->var->name_hint == "threadIdx.z" ||
                 iv->thread_tag == "threadIdx.z") {
        threadIdx_z_ext = op->value;
      }
    }
    StmtVisitor::VisitStmt_(op);
  }

public:
  PrimExpr threadIdx_x_ext = Integer(1);
  PrimExpr threadIdx_y_ext = Integer(1);
  PrimExpr threadIdx_z_ext = Integer(1);
};

void CodeGenTileLangHIP::PrintExtraAttrs(const PrimFunc &f, std::ostream &os) {
  LaunchConfigExtractor extractor;
  extractor(f->body);
  arith::Analyzer analyzer;
  PrimExpr threadIdx_ext =
      analyzer.Simplify(extractor.threadIdx_x_ext * extractor.threadIdx_y_ext *
                        extractor.threadIdx_z_ext);
  if (const IntImmNode *const threadIdx_ext_int =
          threadIdx_ext.as<IntImmNode>()) {
    if (threadIdx_ext_int->value == 1) {
      // unable to extract the number of threads per block, hence directly
      // return
      return;
    }
    // AMD wavefront size is 64.  Sub-wavefront thread counts (e.g. 32)
    // with __launch_bounds__ can cause the compiler to miscompile 64-bit
    // shuffle operations when combined with #pragma unroll.  Only emit
    // launch_bounds when we have at least a full wavefront.
    if (threadIdx_ext_int->value >= 64) {
      stream << " __launch_bounds__(" << threadIdx_ext_int->value << ")";
    }
  }
}

std::string CodeGenTileLangHIP::Finish() {
  // hip must need a header file.
  decl_stream << "#include <hip/hip_runtime.h>\n";
  if (need_mma_h_) {
    decl_stream << "#include <mma.h>\n";
  }

  if (enable_fp8_) {
    decl_stream << "#include <tl_templates/hip/hip_fp8.h>\n";
  }

  if (enable_fp4_) {
    decl_stream << "#include <tl_templates/hip/hip_fp4.h>\n";
  }

  if (need_cooperative_groups_) {
    decl_stream << "#include <hip/hip_cooperative_groups.h>\n";
  }

  decl_stream << "#include <tl_templates/hip/gemm.h>\n";
  decl_stream << "#include <tl_templates/hip/copy.h>\n";
  decl_stream << "#include <tl_templates/hip/reduce.h>\n";
  decl_stream << "#include <tl_templates/hip/scan.h>\n";
  decl_stream << "#include <tl_templates/hip/ldsm.h>\n";
  decl_stream << "#include <tl_templates/hip/threadblock_swizzle.h>\n";
  decl_stream << "#include <tl_templates/hip/debug.h>\n";
  decl_stream << "\n";
  return CodeGenC::Finish();
}

void CodeGenTileLangHIP::VisitStmt_(const tirx::ForNode *op) {
  if (op->kind == tirx::ForKind::kUnrolled) {
    PrintIndent();
    stream << "#pragma unroll\n";
  }
  std::string extent =
      PrintExpr(arith::Analyzer().Simplify(op->extent + op->min));
  PrintIndent();
  std::string vid = AllocVarID(op->loop_var.get());
  std::string start = PrintExpr(op->min);
  stream << "for (";
  PrintType(op->loop_var.dtype(), stream);
  stream << ' ' << vid << " = " << start << "; " << vid << " < " << extent
         << "; ++" << vid << ") {\n";
  int for_scope = BeginScope();
  // A lexical_alloc_scope spanning the entire loop body is redundant with
  // the loop's own braces; unwrap it to avoid emitting `{ {`.
  Stmt body = op->body;
  while (const auto *attr = body.as<AttrStmtNode>()) {
    if (attr->attr_key != tl::attr::kLexicalAllocScope)
      break;
    body = attr->body;
  }
  PrintStmt(body);
  this->EndScope(for_scope);
  PrintIndent();
  stream << "}\n";
}

void CodeGenTileLangHIP::BindThreadIndex(const IterVar &iv) {
  ICHECK(!var_idmap_.count(iv->var.get()));
  var_idmap_[iv->var.get()] =
      CastFromTo(iv->thread_tag, DataType::UInt(32), iv->var.dtype());
}

void CodeGenTileLangHIP::PrintType(DataType t, std::ostream &os) { // NOLINT(*)
  int lanes = t.lanes();
  if (t.is_handle()) {
    ICHECK(t.is_scalar()) << "do not yet support vector types";
    os << "void*";
    return;
  }

  if (t.is_void()) {
    os << "void";
    return;
  }

  if (t == tl::CuTensorMapType()) {
    os << "CUtensorMap";
    return;
  }

  bool fail = false;
  if (t.is_float()) {
    switch (t.bits()) {
    case 16:
      if (t.is_scalar()) {
        os << "half_t";
      } else if (lanes <= 8) {
        // Emit CUDA code to access fp16 vector elements.
        //
        // half4 is stored as uint2
        //
        // h4.x is emitted as *(half2*)(&(u2.x)).x
        // h4.y is emitted as *(half2*)(&(u2.x)).y
        // h4.z is emitted as *(half2*)(&(u2.y)).x
        // h4.w is emitted as *(half2*)(&(u2.y)).y
        //
        ICHECK_EQ(lanes % 2, 0) << "only support even lane for half type";
        os << "uint" << lanes / 2;
      } else {
        fail = true;
      }
      break;
    case 32:
      if (lanes <= 4) {
        os << "float";
      } else if (lanes <= 8) {
        // Emit CUDA code to access fp32 vector elements for 4 < lanes <= 8.
        //
        // float8 is stored as ulonglong4
        //
        // f8.v1 is emitted as *(float2*)(&(ul4.x)).x
        // f8.v2 is emitted as *(float2*)(&(ul4.x)).y
        //
        ICHECK_EQ(lanes % 2, 0)
            << "only support even lane for float type with lanes > 4";
        os << "ulonglong" << lanes / 2;
      } else if (lanes == 16) {
        // float32x16: GCC vector extension type used by MFMA accumulators.
        os << "float32x16";
        return;
      } else if (lanes == 32) {
        // float32x32: GCC vector extension type used by MFMA accumulators.
        os << "float32x32";
        return;
      } else {
        fail = true;
      }
      break;
    case 64:
      os << "double";
      break;
    default:
      fail = true;
      break;
    }
    if (!fail && (t.is_scalar() || t.bits() == 16))
      return;
    if (!fail && (lanes > 4 && lanes <= 8 && t.bits() == 32))
      return;
    if (!fail && (lanes >= 2 && lanes <= 4)) {
      os << lanes;
      return;
    }
  } else if (t.is_bfloat16()) {
    if (t.is_scalar()) {
      os << "bfloat16_t";
    } else if (lanes <= 8) {
      ICHECK_EQ(lanes % 2, 0) << "only support even lane for half type";
      os << "uint" << lanes / 2;
    } else if (lanes == 16) {
      // bfloat16x16: struct { bfloat16_t data[16]; } used by MFMA accumulators.
      os << "bfloat16x16";
      return;
    } else {
      fail = true;
    }
    if (!fail)
      return;
  } else if (t.is_float8()) {
    enable_fp8_ = true;
    os << GetFP8Type(t);
    return;
  } else if (t.is_float4()) {
    // FP4 E2M1 is only supported on gfx950 (CDNA4). Setting enable_fp4_ will
    // cause Finish() to include hip_fp4.h which is itself guarded by
    // #if defined(__gfx950__), so this is safe to emit on all HIP targets
    // (the compiler will reject FP4 kernel code on non-gfx950 anyway).
    enable_fp4_ = true;
    if (t.lanes() <= 32) {
      os << GetFP4Type(t);
    } else {
      LOG(FATAL) << "Cannot convert FP4 type with " << t.lanes()
                 << " lanes to HIP type (max 32 lanes supported)";
    }
    return;
  } else if (t == DataType::Bool()) {
    os << "bool";
    return;
  } else if (t.is_vector_bool()) {
    // CUDA does not support bool vectors.
    // Use ushort vectors to represent instead.
    int n = t.lanes();
    if (n <= 4) {
      os << "ushort" << n;
      return;
    }
  } else if (t.is_uint() || t.is_int()) {
    if (t.is_uint()) {
      os << "u";
    }
    switch (t.bits()) {
    case 1: {
      if (t.is_scalar()) {
        os << "int";
        return;
      } else if (t.lanes() == 8) {
        os << "int8_t";
        return;
      } else if (t.lanes() == 16) {
        os << "int16_t";
        return;
      } else if (t.lanes() == 32) {
        os << "int";
        return;
      } else {
        LOG(FATAL) << "Cannot convert type " << t << " to CUDA type!";
      }
    }
    case 4: {
      if (t.is_scalar()) {
        os << "int";
        return;
      } else if (t.lanes() == 4) {
        os << "int16_t";
        return;
      } else if (t.lanes() == 8) {
        // directly 8 4-bit int in integer.
        os << "int";
        return;
      } else if (t.lanes() == 16) {
        os << "int2";
        return;
      } else if (t.lanes() == 32) {
        os << "int4";
        return;
      } else if (t.lanes() == 64) {
        os << "int8";
        return;
      } else {
        LOG(FATAL) << "Cannot convert type " << t << " to CUDA type!";
      }
    }
    case 8: {
      if (t.lanes() == 4) {
        // directly 4 8 bit int in integer.

        // We use int for int8x4 instead of char4 because using char4 is
        // likely to produce extra instructions to pack four int8 elements
        // into 32-bit data.
        os << "int";
        return;
      } else if (t.lanes() == 8) {
        os << "int2";
        return;
      } else if (t.lanes() == 16) {
        os << "int4";
        return;
      } else if (!t.is_uint() && t.is_scalar()) {
        os << "signed char";
        break;
      } else {
        os << "char";
        break;
      }
    }
    case 16: {
      if (t.is_scalar()) {
        os << "short";
      } else if (t.lanes() <= 4) {
        os << "short" << lanes;
      } else if (t.lanes() <= 8) {
        // Emit CUDA code to access int16 vector elements.
        //
        // short4 is stored as int2
        //
        // s4.x is emitted as *(short2*)(&(i2.x)).x
        // s4.y is emitted as *(short2*)(&(i2.x)).y
        // s4.z is emitted as *(short2*)(&(i2.y)).x
        // s4.w is emitted as *(short2*)(&(i2.y)).y
        //
        ICHECK_EQ(t.lanes() % 2, 0)
            << "only support even lane for shorT type with lanes > 4";
        os << "int" << t.lanes() / 2;
      } else {
        fail = true;
      }
      if (!fail) {
        return;
      }
      break;
    }
    case 32: {
      if (t.is_scalar()) {
        os << "int";
      } else if (t.lanes() <= 4) {
        os << "int" << t.lanes();
      } else if (t.lanes() <= 8) {
        // Emit CUDA code to access int32 vector elements for 4 < lanes <= 8.
        //
        // int8 is stored as longlong4
        //
        // i8.v1 is emitted as *(int2*)(&(l4.x)).x
        // i8.v2 is emitted as *(int2*)(&(l4.x)).y
        //
        ICHECK_EQ(lanes % 2, 0)
            << "only support even lane for int32 type with lanes > 4";
        os << "longlong" << lanes / 2;
      } else {
        fail = true;
      }
      if (!fail) {
        return;
      }
      break;
    }
    case 64: {
      if (t.is_scalar()) {
        os << "int64_t";
      } else if (t.lanes() == 2) {
        os << "longlong2";
      } else if (t.lanes() == 3) {
        os << "longlong3";
      } else if (t.lanes() == 4) {
        os << "longlong4";
      }
      return;
    }
    default:
      fail = true;
      break;
    }
    if (!fail && lanes == 1) {
      return;
    }
    if (!fail && (lanes >= 2 && lanes <= 4)) {
      os << lanes;
      return;
    }
  }
  LOG(FATAL) << "Cannot convert type " << t << " to CUDA type";
}

void CodeGenTileLangHIP::PrintVecBinaryOp(const std::string &op, DataType t,
                                          PrimExpr lhs, PrimExpr rhs,
                                          std::ostream &os) { // NOLINT(*)

  // Declare the result.
  std::string sret = name_supply_->FreshName("_");
  this->PrintIndent();
  this->PrintType(t, stream);
  stream << ' ' << sret << ";\n";
  int ssa_scope = BeginScope();
  {
    // Unpack into individual ops.
    std::string vlhs = SSAGetID(PrintExpr(lhs), lhs.dtype());
    std::string vrhs = SSAGetID(PrintExpr(rhs), rhs.dtype());

    for (int i = 0, lanes = t.lanes(); i < lanes; ++i) {
      std::ostringstream value_temp;
      if (isalpha(op[0])) {
        value_temp << op << "(";
        PrintVecElemLoad(vlhs, lhs.dtype(), i, value_temp);
        value_temp << ", ";
        PrintVecElemLoad(vrhs, rhs.dtype(), i, value_temp);
        value_temp << ")";
      } else {
        value_temp << "(";
        PrintVecElemLoad(vlhs, lhs.dtype(), i, value_temp);
        value_temp << op;
        PrintVecElemLoad(vrhs, rhs.dtype(), i, value_temp);
        value_temp << ")";
      }
      PrintVecElemStore(sret, t, i, value_temp.str());
    }
  }
  EndScope(ssa_scope);
  os << sret;
}

void CodeGenTileLangHIP::PrintVecElemLoad(const std::string &vec, DataType t,
                                          int i,
                                          std::ostream &os) { // NOLINT(*)
  if (t.is_scalar()) {
    os << vec;
    return;
  }

  static const char access[] = {'x', 'y', 'z', 'w'};
  ICHECK(i >= 0 && i < (t.bits() == 8                        ? 16
                        : (t.lanes() == 16)                  ? 16
                        : (t.lanes() == 32)                  ? 32
                        : (t.bits() == 16 || t.bits() == 32) ? 8
                                                             : 4));
  if (t.bits() == 8 && (t.is_int() || t.is_uint())) {
    std::string type_name = t.is_int() ? "char" : "unsigned char";
    if (t.lanes() == 2 || t.lanes() == 3) {
      os << vec << "." << access[i % t.lanes()];
    } else {
      std::string ac = t.lanes() == 4 ? vec : (vec + "." + access[i / 4]);
      os << "((" << type_name << ")(" << ac << " >> " << i % 4 * 8 << "))";
    }
  } else if ((t.lanes() == 16 || t.lanes() == 32) && t.bits() == 32 &&
             t.is_float()) {
    // float32x16/float32x32: __attribute__((__vector_size__(...))) supports
    // subscript.
    os << vec << "[" << i << "]";
  } else if (t.lanes() == 16 && t.is_bfloat16()) {
    // bfloat16x16: struct { bfloat16_t data[16]; }
    os << vec << ".data[" << i << "]";
  } else if (t.is_float16()) {
    os << "((half2*)(&(" << vec << "." << access[i / 2] << ")))->"
       << access[i % 2];
  } else if (t.is_bfloat16()) {
    os << "((bfloat16x2*)(&(" << vec << "." << access[i / 2] << ")))->"
       << access[i % 2];
  } else if (t.lanes() > 4 && t.lanes() <= 8) {
    std::string type_name;
    if (t.bits() == 16) {
      if (t.is_int()) {
        type_name = "short";
      } else if (t.is_uint()) {
        type_name = "ushort";
      }
    } else if (t.bits() == 32) {
      if (t.is_int()) {
        type_name = "int";
      } else if (t.is_uint()) {
        type_name = "uint";
      } else if (t.is_float()) {
        type_name = "float";
      }
    }
    ICHECK(!type_name.empty());
    os << "((" << type_name << "2*)(&(" << vec << "." << access[i / 2]
       << ")))->" << access[i % 2];
  } else {
    os << vec << "." << access[i];
  }
}

void CodeGenTileLangHIP::PrintVecElemStore(const std::string &vec, DataType t,
                                           int i, const std::string &value) {
  this->PrintIndent();
  static const char access[] = {'x', 'y', 'z', 'w'};

  ICHECK(i >= 0 && i < (t.bits() == 8                        ? 16
                        : (t.lanes() == 16)                  ? 16
                        : (t.lanes() == 32)                  ? 32
                        : (t.bits() == 16 || t.bits() == 32) ? 8
                                                             : 4));
  if (t.bits() == 8 && (t.is_int() || t.is_uint())) {
    if (t.lanes() == 2 || t.lanes() == 3) {
      stream << vec << '.' << access[i % t.lanes()] << "=" << "(" << value
             << ");\n";
    } else {
      std::string ac = t.lanes() == 4 ? vec : (vec + "." + access[i / 4]);
      stream << ac << "=";
      // Do not read the first undef lane.
      if (i != 0) {
        stream << ac << " & ~(0x000000ff << " << i % 4 * 8 << ") |";
      }
      stream << "(" << value << " << " << i % 4 * 8 << ");\n";
    }
  } else if ((t.lanes() == 16 || t.lanes() == 32) && t.bits() == 32 &&
             t.is_float()) {
    // float32x16/float32x32: __attribute__((__vector_size__(...))) supports
    // subscript.
    stream << vec << "[" << i << "] = " << value << ";\n";
  } else if (t.lanes() == 16 && t.is_bfloat16()) {
    // bfloat16x16: struct { bfloat16_t data[16]; }
    stream << vec << ".data[" << i << "] = " << value << ";\n";
  } else if (t.is_float16()) {
    stream << "*((half_t*)(&(((half2*)(&(" << vec << "." << access[i / 2]
           << ")))->" << access[i % 2] << "))) = " << value << ";\n";
  } else if (t.is_bfloat16()) {
    stream << "((bfloat16_t*)(&(" << vec << "." << access[i / 2] << ")))["
           << (i % 2) << "] = " << value << ";\n";
  } else if (t.lanes() > 4 && t.lanes() <= 8) {
    std::string type_name;
    if (t.bits() == 16) {
      if (t.is_int()) {
        type_name = "short";
      } else if (t.is_uint()) {
        type_name = "ushort";
      }
    } else if (t.bits() == 32) {
      if (t.is_int()) {
        type_name = "int";
      } else if (t.is_uint()) {
        type_name = "uint";
      } else if (t.is_float()) {
        type_name = "float";
      }
    }
    ICHECK(!type_name.empty());
    stream << "((" << type_name << "2*)(&(" << vec << "." << access[i / 2]
           << ")))->" << access[i % 2] << " = " << value << ";\n";
  } else {
    stream << vec << "." << access[i] << " = " << value << ";\n";
  }
}

void CodeGenTileLangHIP::PrintStorageSync(const CallNode *op) {
  const std::string &sync = op->args[0].as<StringImmNode>()->value;
  if (sync == "warp") {
    // DO nothing.
  } else if (sync == "shared" || sync == "shared.dyn") {
    this->PrintIndent();
    // TODO: change to __builtin_amdgcn_s_barrier() later
    // __syncthreads() will add a vmcnt(0) to final asm, which will break all
    // buffer_load_async and buffer_load_async...lds
    // in tilelang vmcnt should be managed explicitly by cp_async_wait.
    this->stream << "__syncthreads();\n";
  }
}

void CodeGenTileLangHIP::PrintStorageScope(const std::string &scope,
                                           std::ostream &os) { // NOLINT(*)
  ICHECK_NE(scope, "global")
      << "Cannot allocate global memory when targeting CUDA. You must pass "
         "all global arrays as input instead";
  if (scope == "shared") {
    os << "__shared__ ";
  } else if (scope == "shared.dyn") {
    os << "extern __shared__ __align__(1024) ";
  }
}

std::string CodeGenTileLangHIP::CastFromTo(std::string value, DataType from,
                                           DataType target) {
  if (from == target)
    return value;
  std::ostringstream os;
  os << "((";
  this->PrintType(target, os);
  os << ")";
  if (from.is_float16() && (target.is_int() || target.is_uint()) &&
      target.bits() == 8) {
    os << "(";
    if (target.is_uint()) {
      os << "u";
    }
    os << "int)";
  }
  os << value << ")";
  return os.str();
}

void CodeGenTileLangHIP::VisitExpr_(const ShuffleNode *op,
                                    std::ostream &os) { // NOLINT(*)
  // For bfloat16x2 / float16x2 construction from two scalar lanes, emit the
  // HIP pack intrinsic instead of the invalid `uint1(a, b)` that CodeGenC
  // would generate (HIP's uint1 has no two-argument constructor).
  // `uint1{value}` aggregate initialisation is valid: uint1 is defined by ROCm
  // as HIP_vector_type<unsigned int, 1> which has a single .x member.
  DataType t = op->dtype;
  bool is_bf16x2 = t.is_bfloat16() && t.lanes() == 2;
  bool is_fp16x2 = t.is_float16() && t.lanes() == 2;
  if ((is_bf16x2 || is_fp16x2) && op->vectors.size() == 2 &&
      op->vectors[0].dtype().lanes() == 1 &&
      op->vectors[1].dtype().lanes() == 1) {
    std::string e0 = PrintExpr(op->vectors[0]);
    std::string e1 = PrintExpr(op->vectors[1]);
    if (is_bf16x2) {
      os << "uint1{__pack_bfloat162(" << e0 << ", " << e1 << ")}";
    } else {
      os << "uint1{__pack_half2(" << e0 << ", " << e1 << ")}";
    }
    return;
  }
  // Default path for all other shuffle patterns.
  CodeGenC::VisitExpr_(op, os);
}

void CodeGenTileLangHIP::VisitExpr_(const CastNode *op, std::ostream &os) {
  DataType from_ty = op->value.dtype();
  DataType target_ty = op->dtype;
  ICHECK_EQ(target_ty.lanes(), from_ty.lanes());

  // Emit simple C-style type conversion.
  if (from_ty.is_scalar())
    return CodeGenC::VisitExpr_(op, os);

  // ---------------------------------------------------------------------------
  // Vectorized FP4 <-> float/half/bfloat16 conversions (gfx950 only).
  // These use the pair-wise helpers from hip_fp4.h.  We process two lanes at a
  // time: odd-indexed lane is the "high" nibble, even-indexed is "low".
  // Only enabled when the FP4 header is included (enable_fp4_ is set by
  // PrintType when a float4 type is encountered).
  // ---------------------------------------------------------------------------
  int fp4_lanes = from_ty.lanes();
  // Pairwise cast: process 2 FP4 lanes at a time via packed uint8_t byte.
  // Supported lane widths: 2, 4, 8, 16, 32 (all even widths up to FP4x32).
  bool fp4_pair_cast = (fp4_lanes == 2 || fp4_lanes == 4 || fp4_lanes == 8 ||
                        fp4_lanes == 16 || fp4_lanes == 32);

  // FP4 -> float16 : use __tl_cvt_fp4x2_to_half2 per 2-element pair
  if (from_ty.is_float4_e2m1fn() && target_ty.is_float16() && fp4_pair_cast) {
    std::string sret = name_supply_->FreshName("_");
    this->PrintIndent();
    this->PrintType(target_ty, stream);
    stream << ' ' << sret << ";\n";
    std::string src = SSAGetID(PrintExpr(op->value), from_ty);
    // Iterate over pairs: src is stored as fp4_e2_{lanes}_t; we access the
    // packed byte for each pair via reinterpret as uint8_t array.
    // Materialize the uint1 result into a local variable to avoid taking the
    // address of a prvalue (C++ UB).
    for (int i = 0; i < fp4_lanes; i += 2) {
      std::string tmp = name_supply_->FreshName("_fp4h2_");
      this->PrintIndent();
      stream << "uint1 " << tmp << " = __tl_cvt_fp4x2_to_half2(((uint8_t*)&("
             << src << "))[" << i / 2 << "]);\n";
      PrintVecElemStore(sret, target_ty, i, "((half_t*)(&(" + tmp + ")))[0]");
      PrintVecElemStore(sret, target_ty, i + 1,
                        "((half_t*)(&(" + tmp + ")))[1]");
    }
    os << sret;
    return;
  }

  // FP4 -> float32 : use __tl_cvt_fp4x2_to_float2 per 2-element pair
  if (from_ty.is_float4_e2m1fn() && target_ty.is_float() &&
      target_ty.bits() == 32 && fp4_pair_cast) {
    std::string sret = name_supply_->FreshName("_");
    this->PrintIndent();
    this->PrintType(target_ty, stream);
    stream << ' ' << sret << ";\n";
    std::string src = SSAGetID(PrintExpr(op->value), from_ty);
    for (int i = 0; i < fp4_lanes; i += 2) {
      std::ostringstream val;
      val << "__tl_cvt_fp4x2_to_float2(((uint8_t*)&(" << src << "))[" << i / 2
          << "])";
      std::string tmp = name_supply_->FreshName("_fp4f2_");
      this->PrintIndent();
      stream << "float2 " << tmp << " = " << val.str() << ";\n";
      PrintVecElemStore(sret, target_ty, i, tmp + ".x");
      PrintVecElemStore(sret, target_ty, i + 1, tmp + ".y");
    }
    os << sret;
    return;
  }

  // FP4 -> double : use __tl_cvt_fp4x2_to_double2 per 2-element pair
  if (from_ty.is_float4_e2m1fn() && target_ty.is_float() &&
      target_ty.bits() == 64 && fp4_pair_cast) {
    std::string sret = name_supply_->FreshName("_");
    this->PrintIndent();
    this->PrintType(target_ty, stream);
    stream << ' ' << sret << ";\n";
    std::string src = SSAGetID(PrintExpr(op->value), from_ty);
    for (int i = 0; i < fp4_lanes; i += 2) {
      std::ostringstream val;
      val << "__tl_cvt_fp4x2_to_double2(((uint8_t*)&(" << src << "))[" << i / 2
          << "])";
      std::string tmp = name_supply_->FreshName("_fp4d2_");
      this->PrintIndent();
      stream << "double2 " << tmp << " = " << val.str() << ";\n";
      PrintVecElemStore(sret, target_ty, i, tmp + ".x");
      PrintVecElemStore(sret, target_ty, i + 1, tmp + ".y");
    }
    os << sret;
    return;
  }

  // FP4 -> bfloat16 : use __tl_cvt_fp4x2_to_bfloat162 per 2-element pair
  if (from_ty.is_float4_e2m1fn() && target_ty.is_bfloat16() && fp4_pair_cast) {
    std::string sret = name_supply_->FreshName("_");
    this->PrintIndent();
    this->PrintType(target_ty, stream);
    stream << ' ' << sret << ";\n";
    std::string src = SSAGetID(PrintExpr(op->value), from_ty);
    for (int i = 0; i < fp4_lanes; i += 2) {
      std::ostringstream val;
      val << "__tl_cvt_fp4x2_to_bfloat162(((uint8_t*)&(" << src << "))["
          << i / 2 << "])";
      std::string tmp = name_supply_->FreshName("_fp4bf2_");
      this->PrintIndent();
      stream << "uint1 " << tmp << " = " << val.str() << ";\n";
      PrintVecElemStore(sret, target_ty, i,
                        "((bfloat16_t*)(&(" + tmp + ")))[0]");
      PrintVecElemStore(sret, target_ty, i + 1,
                        "((bfloat16_t*)(&(" + tmp + ")))[1]");
    }
    os << sret;
    return;
  }

  // float16 -> FP4
  if (from_ty.is_float16() && target_ty.is_float4_e2m1fn() && fp4_pair_cast) {
    std::string sret = name_supply_->FreshName("_");
    this->PrintIndent();
    this->PrintType(target_ty, stream);
    stream << ' ' << sret << ";\n";
    std::string src = SSAGetID(PrintExpr(op->value), from_ty);
    for (int i = 0; i < fp4_lanes; i += 2) {
      std::ostringstream h0, h1;
      PrintVecElemLoad(src, from_ty, i, h0);
      PrintVecElemLoad(src, from_ty, i + 1, h1);
      std::string tmp = name_supply_->FreshName("_h2fp4_");
      this->PrintIndent();
      stream << "uint1 " << tmp << " = uint1{__pack_half2(" << h0.str() << ", "
             << h1.str() << ")};\n";
      this->PrintIndent();
      stream << "((uint8_t*)&(" << sret << "))[" << i / 2
             << "] = __tl_cvt_half2_to_fp4x2(" << tmp << ");\n";
    }
    os << sret;
    return;
  }

  // float32 -> FP4
  if (from_ty.is_float() && from_ty.bits() == 32 &&
      target_ty.is_float4_e2m1fn() && fp4_pair_cast) {
    std::string sret = name_supply_->FreshName("_");
    this->PrintIndent();
    this->PrintType(target_ty, stream);
    stream << ' ' << sret << ";\n";
    std::string src = SSAGetID(PrintExpr(op->value), from_ty);
    for (int i = 0; i < fp4_lanes; i += 2) {
      std::ostringstream f0, f1;
      PrintVecElemLoad(src, from_ty, i, f0);
      PrintVecElemLoad(src, from_ty, i + 1, f1);
      this->PrintIndent();
      stream << "((uint8_t*)&(" << sret << "))[" << i / 2
             << "] = __tl_cvt_float2_to_fp4x2("
             << "float2{" << f0.str() << ", " << f1.str() << "});\n";
    }
    os << sret;
    return;
  }

  // double -> FP4
  if (from_ty.is_float() && from_ty.bits() == 64 &&
      target_ty.is_float4_e2m1fn() && fp4_pair_cast) {
    std::string sret = name_supply_->FreshName("_");
    this->PrintIndent();
    this->PrintType(target_ty, stream);
    stream << ' ' << sret << ";\n";
    std::string src = SSAGetID(PrintExpr(op->value), from_ty);
    for (int i = 0; i < fp4_lanes; i += 2) {
      std::ostringstream d0, d1;
      PrintVecElemLoad(src, from_ty, i, d0);
      PrintVecElemLoad(src, from_ty, i + 1, d1);
      this->PrintIndent();
      stream << "((uint8_t*)&(" << sret << "))[" << i / 2
             << "] = __tl_cvt_double2_to_fp4x2("
             << "double2{" << d0.str() << ", " << d1.str() << "});\n";
    }
    os << sret;
    return;
  }

  // bfloat16 -> FP4
  if (from_ty.is_bfloat16() && target_ty.is_float4_e2m1fn() && fp4_pair_cast) {
    std::string sret = name_supply_->FreshName("_");
    this->PrintIndent();
    this->PrintType(target_ty, stream);
    stream << ' ' << sret << ";\n";
    std::string src = SSAGetID(PrintExpr(op->value), from_ty);
    for (int i = 0; i < fp4_lanes; i += 2) {
      std::ostringstream b0, b1;
      PrintVecElemLoad(src, from_ty, i, b0);
      PrintVecElemLoad(src, from_ty, i + 1, b1);
      std::string tmp = name_supply_->FreshName("_bf2fp4_");
      this->PrintIndent();
      stream << "uint1 " << tmp << " = uint1{__pack_bfloat162(" << b0.str()
             << ", " << b1.str() << ")};\n";
      this->PrintIndent();
      stream << "((uint8_t*)&(" << sret << "))[" << i / 2
             << "] = __tl_cvt_bfloat162_to_fp4x2(" << tmp << ");\n";
    }
    os << sret;
    return;
  }

  // We could emit make_float4 like calls, but the emitted code looks
  // too compact to read. Emit this as vectorized unary ops.
  std::string sret = name_supply_->FreshName("_");
  this->PrintIndent();
  this->PrintType(target_ty, stream);
  stream << ' ' << sret << ";\n";
  {
    std::string src = SSAGetID(PrintExpr(op->value), from_ty);
    for (int i = 0, lanes = from_ty.lanes(); i < lanes; ++i) {
      std::ostringstream val;
      val << "(";
      PrintType(target_ty.element_of(), val);
      val << ")(";
      PrintVecElemLoad(src, from_ty, i, val);
      val << ")";
      PrintVecElemStore(sret, target_ty, i, val.str());
    }
  }
  os << sret;
}

void CodeGenTileLangHIP::PrintCallExtern(Type ret_type, String global_symbol,
                                         const Array<PrimExpr> &args,
                                         bool skip_first_arg,
                                         std::ostream &os) { // NOLINT(*)
  DataType ret_dtype = GetRuntimeDataType(ret_type);
  if (ret_dtype.is_vector()) {
    //
    // Emit an unsupported vector call
    //
    // v = intrin_f((float4*)A[0], (float4*)B[0])
    //
    // as
    //
    // float4 __ret;
    // {
    //   float4 __arg0 = ((float4*)A)[0];
    //   float4 __arg1 = ((float4*)B)[0];
    //   __ret.x = intrin_f(__arg0.x, __arg1.x);
    //   __ret.y = intrin_f(__arg0.y, __arg1.y);
    //   __ret.z = intrin_f(__arg0.z, __arg1.z);
    //   __ret.w = intrin_f(__arg0.w, __arg1.w);
    // }
    // v = __ret;
    //
    // Declare the result vector.
    std::string sret = name_supply_->FreshName("_");
    this->PrintIndent();
    this->PrintType(ret_dtype, stream);
    stream << ' ' << sret << ";\n";
    {
      // Load arguments.
      std::vector<std::string> sargs;
      size_t arg_begin = static_cast<size_t>(skip_first_arg);
      for (size_t i = arg_begin; i < args.size(); ++i) {
        std::string val = SSAGetID(PrintExpr(args[i]), args[i].dtype());
        sargs.push_back(std::move(val));
      }

      // Emit a scalar call for each lane.
      for (int i = 0; i < ret_dtype.lanes(); ++i) {
        std::ostringstream scall;
        scall << global_symbol << "(";
        for (size_t j = 0; j < sargs.size(); ++j) {
          if (j > 0)
            scall << ", ";
          PrintVecElemLoad(sargs[j], args[arg_begin + j].dtype(), i, scall);
        }
        scall << ")";
        PrintVecElemStore(sret, ret_dtype, i, scall.str());
      }
    }
    os << sret;
  } else {
    CodeGenC::PrintCallExtern(ret_type, global_symbol, args, skip_first_arg,
                              os);
  }
}

// Print a reference expression to a buffer.
std::string CodeGenTileLangHIP::GetBufferRef(DataType t,
                                             const BufferNode *buffer,
                                             PrimExpr index) {
  Var buffer_var = buffer->data;
  const VarNode *buffer_var_node = buffer_var.get();
  std::ostringstream os;
  std::string vid = GetVarID(buffer_var_node);
  std::string scope;
  if (alloc_storage_scope_.count(buffer_var_node)) {
    scope = alloc_storage_scope_.at(buffer_var_node);
  }

  // FP4 scalar access on gfx950: redirect to tl_fp4_packed_load helper.
  // Non-scalar FP4 accesses fall through to the normal path (the vector
  // types fp4_e2_4_t etc. are directly addressable as structs).
  if (t.is_float4() && t.is_scalar()) {
    std::string idx_str = PrintExpr(index);
    auto packed_it = fp4_packed_buffers_.find(buffer_var);
    if (packed_it != fp4_packed_buffers_.end()) {
      // Packed local buffer: use the pre-allocated fp4_e2_2_t array.
      os << "tl_fp4_packed_load(" << packed_it->second << ", " << idx_str
         << ")";
    } else {
      // Non-packed (e.g. shared) buffer: reinterpret as fp4_e2_2_t*.
      os << "tl_fp4_packed_load((fp4_e2_2_t*)" << vid << ", " << idx_str << ")";
    }
    return os.str();
  }

  // bool is_vol = IsVolatile(buffer_var);
  // always false for tl cutlass backend.
  bool is_vol = false;

  auto ptr_cast = [this, is_vol, scope](DataType pointed_to) {
    std::ostringstream ptr_os;
    ptr_os << "(";
    if (is_vol) {
      ptr_os << "volatile ";
    }
    if (!scope.empty() && IsScopePartOfType()) {
      PrintStorageScope(scope, ptr_os);
    }
    PrintType(pointed_to, ptr_os);
    ptr_os << "*)";
    return ptr_os.str();
  };

  DataType buffer_element_dtype = buffer->dtype;

  std::string buffer_str = vid;
  if (!HandleTypeMatch(buffer_var_node, buffer_element_dtype) || is_vol) {
    std::stringstream temp;
    temp << "(" << ptr_cast(buffer_element_dtype) << vid << ")";
    buffer_str = temp.str();
  }

  if (scope.empty()) {
    scope = GetPtrStorageScope(buffer->data);
  }
  // local.var is a scalar — no indexing needed.
  if (scope == "local.var") {
    os << vid;
    return os.str();
  }

  std::string index_str = PrintExpr(index);
  if (t.bits() == 4 || (t.bits() == 1 && t.is_int())) {
    // This is a special case, because CodegenCUDA::PrintType()
    // returns "int" for bool and for 4-bit integers. In most cases,
    // we divide by the number of lanes to determine the index.
    // However, the backing type for scalar int4 and scalar bool is
    // int32.  Therefore, we need to divide by the ratio of their
    // sizes in that case.
    int div_factor = (t.lanes() == 1) ? (32 / t.bits()) : t.lanes();

    os << "*(" << "(" << ptr_cast(t) << vid << ")" << " + " << index_str
       << " / " << div_factor << ")";
  } else if (t == buffer_element_dtype) {
    os << buffer_str << "[" << index_str << "]";
  } else {
    os << "*" << ptr_cast(t) << "(" << buffer_str << " + " << index_str << ")";
  }

  return os.str();
}

void CodeGenTileLangHIP::VisitExpr_(const CallNode *op, std::ostream &os) {
  auto print_extern_call_stmt = [&](std::string name, size_t start = 0,
                                    size_t end = 0) {
    this->PrintIndent();
    this->stream << name << "(";
    for (size_t i = start; i < op->args.size() - end; i++) {
      if (i > start)
        this->stream << ", ";
      this->stream << this->PrintExpr(op->args[i]);
    }
    this->stream << ");\n";
  };
  if (op->op.same_as(builtin::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 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 {
      // Predicated version
      std::string condition = this->PrintExpr(op->args[3]);
      this->stream << "tl::cp_async_gs_conditional<" << size << ">(" << dst
                   << ", " << src << ", " << condition << ");\n";
    }
  } else if (op->op.same_as(tl::ptx_cp_async())) {
    int total_bytes = GetTileLangCPAsyncTransferBytes(op);
    std::string dst = this->PrintExpr(op->args[0]);
    std::string src = this->PrintExpr(op->args[1]);
    std::string size = std::to_string(total_bytes);
    this->PrintIndent();
    if (op->args.size() == 3) {
      this->stream << "tl::cp_async_gs<" << size << ">(" << dst << ", " << src
                   << ");\n";
    } else {
      std::string condition = this->PrintExpr(op->args[3]);
      this->stream << "tl::cp_async_gs_conditional<" << size << ">(" << dst
                   << ", " << src << ", " << condition << ");\n";
    }
  } else if (op->op.same_as(builtin::ptx_commit_group())) {
    print_extern_call_stmt("tl::cp_async_commit");
  } else if (op->op.same_as(builtin::ptx_wait_group())) {
    int n = Downcast<IntImm>(op->args[0])->value;
    std::string func_name = "tl::cp_async_wait<" + std::to_string(n) + ">";
    print_extern_call_stmt(func_name, 1);
  } else if (op->op.same_as(builtin::create_barriers())) {
    this->PrintIndent();
    int barrier_count = Downcast<IntImm>(op->args[0])->value;
    std::string barrier_name = "_mbarrier";
    this->stream << "__shared__ uint64_t " << barrier_name << "["
                 << barrier_count << "];\n";
  } else if (op->op.same_as(builtin::ptx_arrive_barrier())) {
    print_extern_call_stmt("tl::mbarrier_arrive");
  } else if (op->op.same_as(builtin::ptx_init_barrier_thread_count())) {
    print_extern_call_stmt("tl::mbarrier_init");
  } else if (op->op.same_as(builtin::ptx_arrive_barrier_expect_tx())) {
    print_extern_call_stmt("tl::mbarrier_arrive_expect_tx");
  } else if (op->op.same_as(builtin::ptx_cp_async_barrier())) {
    print_extern_call_stmt("tl::mbarrier_cp_async_arrive");
  } else if (op->op.same_as(tl::mbarrier_expect_tx())) {
    print_extern_call_stmt("tl::mbarrier_expect_tx");
  } else if (op->op.same_as(tl::mbarrier_wait_parity())) {
    print_extern_call_stmt("tl::mbarrier_wait");
  } else if (op->op.same_as(tl::ptx_stmatrix())) {
    int trans = Downcast<IntImm>(op->args[0])->value;
    int num = Downcast<IntImm>(op->args[1])->value;
    bool has_shape =
        op->args.size() >= 4 && op->args.back().as<StringImmNode>();
    if (has_shape) {
      ICHECK_EQ(Downcast<StringImm>(op->args.back())->value, "m8n8")
          << "HIP stmatrix codegen only supports m8n8";
    }
    std::string func_name = "tl::ptx_stmatrix_x" + std::to_string(num);
    if (trans == 1)
      func_name += "_trans";
    print_extern_call_stmt(func_name, 2, has_shape ? 1 : 0);
  } else if (op->op.same_as(tl::wait_wgmma())) {
    this->PrintIndent();
    int num_mma = Downcast<IntImm>(op->args[0])->value;
    this->stream << "tl::wait_wgmma<" << std::to_string(num_mma) << ">();\n";
  } else if (op->op.same_as(tl::pack_b16())) {
    os << "__pack_half2(" << this->PrintExpr(op->args[0]) << ", "
       << this->PrintExpr(op->args[1]) << ")";
  } else if (op->op.same_as(tl::sync_grid())) {
    this->need_cooperative_groups_ = true;
    this->PrintIndent();
    this->stream << "cooperative_groups::this_grid().sync();\n";
  } else if (op->op.same_as(tl::sync_warp())) {
    // AMD wavefronts execute in lockstep, so intra-wavefront convergence is
    // guaranteed by the hardware. __syncwarp() has no HIP equivalent and is a
    // no-op here. The mask argument (if present) is intentionally ignored.
  } else if (op->op.same_as(tl::any_sync())) {
    ICHECK_EQ(op->args.size(), 2U) << "tl.any_sync expects <mask, predicate>.";
    // HIP __any takes only the predicate; the mask is ignored because
    // wavefront execution is always convergent across the full wave.
    os << "__any(" << PrintExpr(op->args[1]) << ")";
  } else if (op->op.same_as(tl::all_sync())) {
    ICHECK_EQ(op->args.size(), 2U) << "tl.all_sync expects <mask, predicate>.";
    os << "__all(" << PrintExpr(op->args[1]) << ")";
  } else if (op->op.same_as(tl::ballot_sync())) {
    ICHECK_EQ(op->args.size(), 2U)
        << "tl.ballot_sync expects <mask, predicate>.";
    // HIP __ballot returns uint64 natively, covering every lane of the
    // wavefront; the CUDA-style mask argument is ignored.
    os << "((unsigned long long)__ballot(" << PrintExpr(op->args[1]) << "))";
  } else if (op->op.same_as(tl::ballot())) {
    ICHECK_EQ(op->args.size(), 1U) << "tl.ballot expects <predicate>.";
    os << "((unsigned long long)__ballot(" << PrintExpr(op->args[0]) << "))";
  } else if (op->op.same_as(tl::activemask())) {
    ICHECK(op->args.empty()) << "tl.activemask takes no arguments.";
    os << "((unsigned long long)__ballot(1))";
  } else if (op->op.same_as(tl::syncthreads_count())) {
    ICHECK_EQ(op->args.size(), 1U)
        << "tl.syncthreads_count expects <predicate>.";
    os << "__syncthreads_count(" << PrintExpr(op->args[0]) << ")";
  } else if (op->op.same_as(tl::syncthreads_and())) {
    ICHECK_EQ(op->args.size(), 1U) << "tl.syncthreads_and expects <predicate>.";
    os << "__syncthreads_and(" << PrintExpr(op->args[0]) << ")";
  } else if (op->op.same_as(tl::syncthreads_or())) {
    ICHECK_EQ(op->args.size(), 1U) << "tl.syncthreads_or expects <predicate>.";
    os << "__syncthreads_or(" << PrintExpr(op->args[0]) << ")";
  } else if (op->op.same_as(tl::shfl_sync())) {
    ICHECK_EQ(op->args.size(), 4U)
        << "tl.shfl_sync expects <mask, value, src_lane, width>.";
    // HIP __shfl takes only (value, src_lane, width); the mask is ignored.
    os << "__shfl(" << PrintExpr(op->args[1]) << ", " << PrintExpr(op->args[2])
       << ", " << PrintExpr(op->args[3]) << ")";
  } else if (op->op.same_as(tl::shfl_xor_sync())) {
    ICHECK_EQ(op->args.size(), 4U)
        << "tl.shfl_xor_sync expects <mask, value, lane_mask, width>.";
    os << "__shfl_xor(" << PrintExpr(op->args[1]) << ", "
       << PrintExpr(op->args[2]) << ", " << PrintExpr(op->args[3]) << ")";
  } else if (op->op.same_as(tl::shfl_down_sync())) {
    ICHECK_EQ(op->args.size(), 4U)
        << "tl.shfl_down_sync expects <mask, value, delta, width>.";
    os << "__shfl_down(" << PrintExpr(op->args[1]) << ", "
       << PrintExpr(op->args[2]) << ", " << PrintExpr(op->args[3]) << ")";
  } else if (op->op.same_as(tl::shfl_up_sync())) {
    ICHECK_EQ(op->args.size(), 4U)
        << "tl.shfl_up_sync expects <mask, value, delta, width>.";
    os << "__shfl_up(" << PrintExpr(op->args[1]) << ", "
       << PrintExpr(op->args[2]) << ", " << PrintExpr(op->args[3]) << ")";
  } else if (op->op.same_as(tl::match_any_sync()) ||
             op->op.same_as(tl::match_all_sync())) {
    LOG(FATAL) << "tl." << op->op << " is not supported on HIP: the "
               << "__match_{any,all}_sync primitives have no HIP equivalent.";
  } else if (op->op.same_as(tl::add2()) || op->op.same_as(tl::sub2()) ||
             op->op.same_as(tl::mul2()) || op->op.same_as(tl::fma2()) ||
             op->op.same_as(tl::max2()) || op->op.same_as(tl::min2()) ||
             op->op.same_as(tl::abs2())) {
    // Packed x2 element-wise math intrinsics.
    // HIP currently only supports float32x2 (float2).
    std::string op_name;
    if (op->op.same_as(tl::add2()))
      op_name = "add2";
    else if (op->op.same_as(tl::sub2()))
      op_name = "sub2";
    else if (op->op.same_as(tl::mul2()))
      op_name = "mul2";
    else if (op->op.same_as(tl::fma2()))
      op_name = "fma2";
    else if (op->op.same_as(tl::max2()))
      op_name = "max2";
    else if (op->op.same_as(tl::min2()))
      op_name = "min2";
    else
      op_name = "abs2";

    os << "tl::" << op_name << "(";
    os << PrintExpr(op->args[0]);
    for (size_t i = 1; i < op->args.size(); ++i) {
      os << ", " << PrintExpr(op->args[i]);
    }
    os << ")";
  } else if (op->op.same_as(tl::ds_read_tr16_b64())) {
    ICHECK_EQ(op->args.size(), 1U)
        << "tl.ds_read_tr16_b64 expects one argument (smem_access_ptr).";
    os << "tl::ds_read_tr16_b64(" << PrintExpr(op->args[0]) << ")";
  } else if (op->op.same_as(tl::ds_read_tr8_b64())) {
    ICHECK_EQ(op->args.size(), 1U)
        << "tl.ds_read_tr8_b64 expects one argument (smem_access_ptr).";
    os << "tl::ds_read_tr8_b64(" << PrintExpr(op->args[0]) << ")";
  } else if (op->op.same_as(tl::__ldg())) {
    // HIP fallback: regular load
    const BufferLoadNode *bl = op->args[0].as<BufferLoadNode>();
    ICHECK(bl) << "T.__ldg expects a BufferLoad as the first argument.";
    ICHECK_EQ(bl->indices.size(), 1)
        << "T.__ldg currently supports flattened 1D buffer accesses.";
    const BufferNode *buffer = bl->buffer.get();
    PrimExpr base = bl->indices[0];
    auto buffer_ref = this->GetBufferRef(op->dtype, buffer, base);
    os << buffer_ref;
  } else if (op->op.same_as(builtin::tvm_fill_fragment())) {
    need_mma_h_ = true;
    ICHECK_EQ(op->args.size(), 6U);
    os << "nvcuda::wmma::fill_fragment(";
    this->PrintExpr(op->args[0], os);
    os << "[";
    this->PrintExpr(op->args[4], os);
    os << "], ";
    this->PrintExpr(op->args[5], os);
    os << ")";
  } else if (op->op.same_as(builtin::tvm_load_matrix_sync())) {
    need_mma_h_ = true;
    ICHECK_EQ(op->args.size(), 8U);
    os << "nvcuda::wmma::load_matrix_sync(";
    this->PrintExpr(op->args[0], os);
    os << "[";
    this->PrintExpr(op->args[4], os);
    os << "], ";
    this->PrintExpr(op->args[5], os);
    os << ", ";
    this->PrintExpr(op->args[6], os);
    os << ")";
  } else if (op->op.same_as(builtin::tvm_store_matrix_sync())) {
    need_mma_h_ = true;
    ICHECK_EQ(op->args.size(), 8U);
    os << "nvcuda::wmma::store_matrix_sync(";
    this->PrintExpr(op->args[5], os);
    os << ", ";
    this->PrintExpr(op->args[0], os);
    os << "[";
    this->PrintExpr(op->args[4], os);
    os << "], ";
    this->PrintExpr(op->args[6], os);
    if (const StringImmNode *str = op->args[7].as<StringImmNode>()) {
      os << ", nvcuda::wmma::mem_" << str->value;
    } else {
      LOG(FATAL) << "Invalid parameters";
    }
    os << ")";
  } else if (op->op.same_as(builtin::tvm_mma_sync())) {
    need_mma_h_ = true;
    ICHECK_EQ(op->args.size(), 8U);
    os << "nvcuda::wmma::mma_sync(";
    for (int i = 0; i < 4; ++i) {
      this->PrintExpr(op->args[i * 2], os);
      os << "[";
      this->PrintExpr(op->args[i * 2 + 1], os);
      os << "]" << ((i < 3) ? ", " : ")");
    }
  } else if (op->op.same_as(builtin::tvm_bmma_sync())) {
    need_mma_h_ = true;
    ICHECK_EQ(op->args.size(), 8U);
    os << "nvcuda::wmma::bmma_sync(";
    for (int i = 0; i < 4; ++i) {
      this->PrintExpr(op->args[i * 2], os);
      os << "[";
      this->PrintExpr(op->args[i * 2 + 1], os);
      os << "]" << ((i < 3) ? ", " : ")");
    }
  } else if (op->op.same_as(tl::tvm_mfma())) {
    // arg 0: prefix: {otype}_{intrM}x{intrN}x{intrK}_{itype}
    // arg 1: A layout: row/col
    // arg 2: B layout: row/col
    // arg 3: A precision: float16, float32, ...
    // arg 4: B precision: float16, float32, ...
    // arg 5: C precision: float32, float64, ...
    // arg 6: A multiplicand
    // arg 7: A multiplicand index
    // arg 8: B multiplicand
    // arg 9: B multiplicand index
    // arg 10: C accumulator
    // arg 11: C accumulator index

    ICHECK(op->args.size() == 12U)
        << "Invalid number of arguments for tvm_mfma";
    std::string prefix = Downcast<StringImm>(op->args[0])->value;
    std::string A_layout = Downcast<StringImm>(op->args[1])->value;
    std::string B_layout = Downcast<StringImm>(op->args[2])->value;
    std::string A_dtype = Downcast<StringImm>(op->args[3])->value;
    std::string B_dtype = Downcast<StringImm>(op->args[4])->value;
    std::string C_dtype = Downcast<StringImm>(op->args[5])->value;
    std::string a_ref = this->PrintExpr(op->args[6]);
    std::string a_bias = this->PrintExpr(op->args[7]);
    std::string b_ref = this->PrintExpr(op->args[8]);
    std::string b_bias = this->PrintExpr(op->args[9]);
    std::string c_ref = this->PrintExpr(op->args[10]);
    std::string c_bias = this->PrintExpr(op->args[11]);
    ICHECK(A_layout == "row" || B_layout == "row")
        << "Matrix core only support row major";
    // map for dtype -> float32x4 -> float4
    std::unordered_map<std::string, std::string> dtype_map = {
        {"int8", "char"},
        {"int32", "int"},
        {"int8x4", "int32_t"},
        {"int8x8", "int64_t"},
        {"int8x16", "int32x4"},
        {"int32x4", "int32x4"},
        {"int32x16", "int32x16"},
        {"float16", "half"},
        {"float32", "float"},
        {"float64", "double"},
        {"float16x4", "float16x4"},
        {"float16x8", "float16x8"},
        {"bfloat16x4", "bfloat16x4_vec"},
        {"bfloat16x8", "bfloat16x8_vec"},
        {"float32x4", "float32x4"},
        {"float8_e4m3fnuzx4", "fp8_e4_4_t"},
        {"float8_e4m3fnx4", "fp8_e4_4_t"},
        {"float8_e4m3fnuzx8", "long"},
        {"float8_e4m3fnx8", "long"},
        {"float8_e5m2fnuzx4", "fp8_e5_4_t"},
        {"float8_e5m2fnuzx8", "long"},
        {"float8_e5m2x4", "fp8_e5_4_t"},
        {"float8_e5m2x8", "long"},
        {"float32x16", "float32x16"},
        {"float32x32", "float32x32"}};
    std::string call_mfma_code = R"({
      *((({C_dtype}*){c_ref}) + {c_bias}) = {mfma_buildin}(*((({A_dtype}*){a_ref}) + {a_bias}),
                    *((({B_dtype}*){b_ref}) + {b_bias}),
                    *((({C_dtype}*){c_ref}) + {c_bias}), 0, 0, 0);
    })";
    std::string mfma_buildin = "__builtin_amdgcn_mfma_" + prefix;
    Replacer replacer;

    replacer.register_rule("{mfma_buildin}", mfma_buildin);
    replacer.register_rule("{A_dtype}", dtype_map[A_dtype]);
    replacer.register_rule("{B_dtype}", dtype_map[B_dtype]);
    replacer.register_rule("{C_dtype}", dtype_map[C_dtype]);
    replacer.register_rule("{a_ref}", a_ref);
    replacer.register_rule("{a_bias}", a_bias);
    replacer.register_rule("{b_ref}", b_ref);
    replacer.register_rule("{b_bias}", b_bias);
    replacer.register_rule("{c_ref}", c_ref);
    replacer.register_rule("{c_bias}", c_bias);
    os << replacer.rewrite(call_mfma_code);
  } else if (op->op.same_as(tl::tvm_rdna_wmma())) {
    // arg 0: shape string, e.g. "f32_16x16x16_f16_w32"
    // arg 1: A layout: "row"/"col"
    // arg 2: B layout: "row"/"col"
    // arg 3: A dtype, e.g. "float16x8"
    // arg 4: B dtype
    // arg 5: C dtype, e.g. "float32x8"
    // arg 6: A data pointer
    // arg 7: A element index (in units of the vectorized type)
    // arg 8: B data pointer
    // arg 9: B element index
    // arg 10: C data pointer
    // arg 11: C element index
    ICHECK(op->args.size() == 12U) << "tvm_rdna_wmma expects 12 arguments";
    std::string shape = Downcast<StringImm>(op->args[0])->value;
    std::string A_layout = Downcast<StringImm>(op->args[1])->value;
    std::string B_layout = Downcast<StringImm>(op->args[2])->value;
    std::string A_dtype = Downcast<StringImm>(op->args[3])->value;
    std::string B_dtype = Downcast<StringImm>(op->args[4])->value;
    std::string C_dtype = Downcast<StringImm>(op->args[5])->value;
    std::string a_ref = this->PrintExpr(op->args[6]);
    std::string a_bias = this->PrintExpr(op->args[7]);
    std::string b_ref = this->PrintExpr(op->args[8]);
    std::string b_bias = this->PrintExpr(op->args[9]);
    std::string c_ref = this->PrintExpr(op->args[10]);
    std::string c_bias = this->PrintExpr(op->args[11]);

    // Get RDNA Generation
    ICHECK(target_.defined()) << "CodeGenTileLangHIP target is not set";
    int rdna_gen = tvm::tl::TargetGetRDNAGeneration(target_);
    ICHECK(rdna_gen == 11 || rdna_gen == 12)
        << "Unsupported RDNA target for WMMA: gfx" << target_->str();

    // Determine wmma builtin name from shape
    // shape = "f32_16x16x16_f16_w32" ->
    // "__builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12" For gfx12 targets use
    // the _gfx12 suffix variant, which gfx11 targets don't have.
    std::string wmma_builtin = "__builtin_amdgcn_wmma_" + shape;

    int ab_half_elems = 16;
    std::string ab_vec_typedef = "tl_v16f16";
    if (rdna_gen == 12) {
      wmma_builtin += "_gfx12";
      ab_half_elems = 8;
      ab_vec_typedef = "tl_v8f16";
    }

    // Emit the WMMA call.
    // For gfx12:
    // Signature: v8f32 = wmma_builtin(v8f16 a, v8f16 b, v8f32 c)
    // where v8f16 = __fp16 x 8, v8f32 = float x 8.
    // For gfx11:
    // Signature: v8f32 = wmma_builtin(v16f16 a, v16f16 b, v8f32 c)
    // where v16f16 = __fp16 x 16, v8f32 = float x 8.
    //
    // A/B buffers hold half_t (fp16), C/D buffers hold float.
    // Each element index accesses a packed vector of 8/16 elements.
    //
    // Using typedef'd vector types for the cast:
    //   typedef __attribute__((__vector_size__(8/16 * sizeof(__fp16)))) __fp16
    //   tl_v8/16f16; typedef __attribute__((__vector_size__(8 *
    //   sizeof(float)))) float tl_v8f32;
    std::string call_wmma_code = R"({
      typedef __attribute__((__vector_size__({ab_half_elems} * sizeof(__fp16)))) __fp16 {ab_vec_typedef};
      typedef __attribute__((__vector_size__(8 * sizeof(float)))) float tl_v8f32;
      *((tl_v8f32*){c_ref} + {c_bias}) = {wmma_builtin}(
          *(({ab_vec_typedef}*){a_ref} + {a_bias}),
          *(({ab_vec_typedef}*){b_ref} + {b_bias}),
          *((tl_v8f32*){c_ref} + {c_bias}));
    })";
    Replacer wmma_replacer;
    wmma_replacer.register_rule("{wmma_builtin}", wmma_builtin);
    wmma_replacer.register_rule("{ab_half_elems}",
                                std::to_string(ab_half_elems));
    wmma_replacer.register_rule("{ab_vec_typedef}", ab_vec_typedef);
    wmma_replacer.register_rule("{a_ref}", a_ref);
    wmma_replacer.register_rule("{a_bias}", a_bias);
    wmma_replacer.register_rule("{b_ref}", b_ref);
    wmma_replacer.register_rule("{b_bias}", b_bias);
    wmma_replacer.register_rule("{c_ref}", c_ref);
    wmma_replacer.register_rule("{c_bias}", c_bias);
    os << wmma_replacer.rewrite(call_wmma_code);
  } else if (op->op.same_as(builtin::thread_return())) {
    os << "return";
  } else if (op->op.same_as(tl::loop_break())) {
    this->PrintIndent();
    this->stream << "break;\n";
  } else if (op->op.same_as(tl::no_set_max_nreg())) {
    // HIP doesn't need explicit register management like CUDA
    // This is a no-op for HIP
    return;
  } else if (op->op.same_as(tl::annotate_producer_reg_dealloc()) ||
             op->op.same_as(tl::annotate_consumer_reg_alloc())) {
    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 if (op->op.same_as(tl::atomic_add_elem_op())) {
    // atomic_add_elem_op(dst_ptr, src_value[, memory_order])
    std::string dst_ptr = PrintExpr(op->args[0]);
    std::string src_value = PrintExpr(op->args[1]);
    this->PrintIndent();
    this->stream << "AtomicAdd(" << dst_ptr << ", " << src_value;
    if (op->args.size() > 2) {
      this->stream << ", " << PrintExpr(op->args[2]);
    }
    this->stream << ");\n";
  } else if (op->op.same_as(tl::atomic_add_ret_elem_op())) {
    // atomic_add_ret_elem_op(dst_ptr, src_value[, memory_order]) -> returns
    // prev value
    os << "AtomicAddRet(" << PrintExpr(op->args[0]) << ", "
       << PrintExpr(op->args[1]);
    if (op->args.size() > 2) {
      os << ", " << PrintExpr(op->args[2]);
    }
    os << ")";
  } else if (op->op.same_as(tl::atomic_addx2_elem_op())) {
    // atomic_addx2_elem_op(dst_ptr, src_ptr[, memory_order])
    std::string dst_ptr = PrintExpr(op->args[0]);
    std::string src_ptr = PrintExpr(op->args[1]);
    this->PrintIndent();
    this->stream << "AtomicAddx2(" << dst_ptr << ", " << src_ptr;
    if (op->args.size() > 2) {
      this->stream << ", " << PrintExpr(op->args[2]);
    }
    this->stream << ");\n";
  } else if (op->op.same_as(tl::atomic_addx4_elem_op())) {
    // atomic_addx4_elem_op(dst_ptr, src_ptr[, memory_order])
    std::string dst_ptr = PrintExpr(op->args[0]);
    std::string src_ptr = PrintExpr(op->args[1]);
    this->PrintIndent();
    this->stream << "AtomicAddx4(" << dst_ptr << ", " << src_ptr;
    if (op->args.size() > 2) {
      this->stream << ", " << PrintExpr(op->args[2]);
    }
    this->stream << ");\n";
  } else if (op->op.same_as(tl::atomic_load_elem_op())) {
    // atomic_load_elem_op(src_ptr, memory_order) -> returns loaded value
    os << "AtomicLoad(" << PrintExpr(op->args[0]) << ", "
       << PrintExpr(op->args[1]) << ")";
  } else if (op->op.same_as(tl::atomic_store_elem_op())) {
    // atomic_store_elem_op(dst_ptr, value, memory_order)
    std::string dst_ptr = PrintExpr(op->args[0]);
    std::string value = PrintExpr(op->args[1]);
    std::string memory_order = PrintExpr(op->args[2]);
    this->PrintIndent();
    this->stream << "AtomicStore(" << dst_ptr << ", " << value << ", "
                 << memory_order << ");\n";
  } else if (op->op.same_as(tl::atomic_or_elem_op())) {
    // atomic_or_elem_op(dst_ptr, src_value[, memory_order])
    std::string dst_ptr = PrintExpr(op->args[0]);
    std::string src_value = PrintExpr(op->args[1]);
    this->PrintIndent();
    this->stream << "AtomicOr(" << dst_ptr << ", " << src_value;
    if (op->args.size() > 2) {
      this->stream << ", " << PrintExpr(op->args[2]);
    }
    this->stream << ");\n";
  } else if (op->op.same_as(tl::atomic_max_elem_op())) {
    // atomic_max_elem_op(dst_ptr, src_value[, memory_order])
    std::string dst_ptr = PrintExpr(op->args[0]);
    std::string src_value = PrintExpr(op->args[1]);
    this->PrintIndent();
    this->stream << "AtomicMax(" << dst_ptr << ", " << src_value;
    if (op->args.size() > 2) {
      this->stream << ", " << PrintExpr(op->args[2]);
    }
    this->stream << ");\n";
  } else if (op->op.same_as(tl::atomic_max_ret_elem_op())) {
    // atomic_max_ret_elem_op(dst_ptr, src_value[, memory_order]) -> returns
    // prev value
    os << "AtomicMaxRet(" << PrintExpr(op->args[0]) << ", "
       << PrintExpr(op->args[1]);
    if (op->args.size() > 2) {
      os << ", " << PrintExpr(op->args[2]);
    }
    os << ")";
  } else if (op->op.same_as(tl::atomic_min_elem_op())) {
    // atomic_min_elem_op(dst_ptr, src_value[, memory_order])
    std::string dst_ptr = PrintExpr(op->args[0]);
    std::string src_value = PrintExpr(op->args[1]);
    this->PrintIndent();
    this->stream << "AtomicMin(" << dst_ptr << ", " << src_value;
    if (op->args.size() > 2) {
      this->stream << ", " << PrintExpr(op->args[2]);
    }
    this->stream << ");\n";
  } else if (op->op.same_as(tl::atomic_min_ret_elem_op())) {
    // atomic_min_ret_elem_op(dst_ptr, src_value[, memory_order]) -> returns
    // prev value
    os << "AtomicMinRet(" << PrintExpr(op->args[0]) << ", "
       << PrintExpr(op->args[1]);
    if (op->args.size() > 2) {
      os << ", " << PrintExpr(op->args[2]);
    }
    os << ")";
  } else {
    CodeGenC::VisitExpr_(op, os);
  }
}

void CodeGenTileLangHIP::VisitStmt_(const AttrStmtNode *op) {
  if (op->attr_key == tl::attr::kLexicalAllocScope) {
    PrintIndent();
    stream << "{\n";
    int scope = BeginScope();
    PrintStmt(op->body);
    EndScope(scope);
    PrintIndent();
    stream << "}\n";
    return;
  } else if (op->attr_key == "threadblock_swizzle_pattern") {
    this->PrintIndent();
    std::string func_name;
    int panel_size = 0;
    if (const auto *call = op->value.as<CallNode>()) {
      if (call->op.same_as(tirx::builtin::tvm_tuple()) &&
          call->args.size() >= 2) {
        const auto *name_node = call->args[0].as<StringImmNode>();
        const auto *size_node = call->args[1].as<IntImmNode>();
        ICHECK(name_node && size_node) << "threadblock_swizzle_pattern expects "
                                          "tvm_tuple(device_func, panel_size)";
        func_name = name_node->value;
        panel_size = static_cast<int>(size_node->value);
      }
    }
    ICHECK(!func_name.empty() && panel_size > 0)
        << "threadblock_swizzle_pattern: failed to extract func_name and "
           "panel_size";
    this->stream << "const dim3 blockIdx = tl::" << func_name << "<"
                 << panel_size << ">();\n";
    this->VisitStmt(op->body);
    return;
  }
  CodeGenC::VisitStmt_(op);
}

void CodeGenTileLangHIP::VisitStmt_(const AllocBufferNode *op) {
  std::string vid = AllocVarID(op->buffer->data.get());

  this->PrintIndent();
  std::string scope = GetPtrStorageScope(op->buffer->data);
  DataType dtype = op->buffer->dtype;

  // FP4 scalar local buffers use packed storage (2 elements per byte).
  // Skip the normal type+scope header; emit the packed type directly below.
  bool is_fp4_scalar_local =
      dtype.is_float4() && dtype.is_scalar() && scope == "local";

  if (!is_fp4_scalar_local) {
    PrintStorageScope(scope, stream);
    PrintType(dtype, stream);
  }

  if (scope == "shared.dyn") {
    stream << ' ' << vid << "[];\n";
  } else {
    auto alloc_ref = GetRef<AllocBuffer>(op);
    auto opt_size = alloc_ref.ConstantAllocationSize();
    ICHECK(opt_size.has_value() && opt_size.value() > 0)
        << "Can only handle constant size stack allocation for now";
    size_t constant_size = static_cast<size_t>(opt_size.value());

    if (scope == "shared") {
      if (dtype.is_float4_e2m1fn() && dtype.is_scalar()) {
        constant_size = (constant_size + 1) / 2;
      } else if (dtype == DataType::Int(4) || dtype == DataType::UInt(4) ||
                 dtype == DataType::Int(1)) {
        size_t elements_per_storage_word = 32 / dtype.bits();
        constant_size = (constant_size + elements_per_storage_word - 1) /
                        elements_per_storage_word;
      }
    }

    if (is_fp4_scalar_local) {
      // Use fp4_e2_2_t (2 elements per byte) as packed storage for FP4
      // local register arrays.  Record the mapping so BufferLoad/Store
      // can emit tl_fp4_packed_load / tl_fp4_packed_store calls.
      enable_fp4_ = true;
      auto vid_packed = vid + "_packed";
      stream << "fp4_e2_2_t " << vid_packed << '[' << (constant_size + 1) / 2
             << "];\n";
      fp4_packed_buffers_[op->buffer->data] = vid_packed;
    } else if (scope == "local.var") {
      // Single-element variable: emit an initializer so the value is defined.
      // Default to 0; respect the user-provided tl.local_var_init annotation.
      PrimExpr init = tirx::make_const(op->buffer->dtype, 0);
      auto init_it = op->annotations.find(tl::attr::kLocalVarInit);
      if (init_it != op->annotations.end()) {
        PrimExpr user_init = Downcast<PrimExpr>((*init_it).second);
        if (!user_init.dtype().is_void() &&
            user_init.dtype() != op->buffer->dtype) {
          user_init = tirx::Cast(op->buffer->dtype, user_init);
        }
        init = user_init;
      }
      stream << ' ' << vid << " = " << PrintExpr(init) << ";\n";
    } else {
      stream << ' ' << vid << '[' << constant_size << "];\n";
    }
  }

  RegisterHandleType(op->buffer->data.get(), op->buffer->dtype);
}

void CodeGenTileLangHIP::VisitStmt_(const BufferStoreNode *op) {
  ICHECK_EQ(op->indices.size(), 1) << "Store to non-flat memory not supported.";
  ICHECK(!op->predicate.defined())
      << "Predicated buffer store is not supported.";

  DataType value_dtype = op->value.dtype();
  DataType element_dtype = op->buffer->dtype;
  Var buffer_var = op->buffer->data;

  // FP4 scalar store: use tl_fp4_packed_store to correctly handle nibble-level
  // writes without corrupting the neighbouring nibble.
  if (element_dtype.is_float4() && element_dtype.is_scalar() &&
      value_dtype.is_scalar()) {
    std::string idx_str = PrintExpr(op->indices[0]);
    std::string value = this->PrintExpr(op->value);
    this->PrintIndent();
    auto packed_it = fp4_packed_buffers_.find(buffer_var);
    if (packed_it != fp4_packed_buffers_.end()) {
      stream << "tl_fp4_packed_store(" << packed_it->second << ", " << idx_str
             << ", " << value << ");\n";
    } else {
      stream << "tl_fp4_packed_store((fp4_e2_2_t*)"
             << GetVarID(buffer_var.get()) << ", " << idx_str << ", " << value
             << ");\n";
    }
    return;
  }

  // Default path for all other types.
  CodeGenC::VisitStmt_(op);
}

void CodeGenTileLangHIP::VisitExpr_(const RampNode *op, std::ostream &os) {
  int lanes = static_cast<int>(Downcast<IntImm>(op->lanes)->value);
  ICHECK_LE(lanes, 4)
      << "ValueError: Ramp of more than 4 lanes is not allowed.";
  os << "(make_";
  PrintType(op->dtype, os);
  os << "(";
  for (int i = 0; i < lanes; i++) {
    os << "(" << PrintExpr(op->base) << ")" << "+(" << PrintExpr(op->stride)
       << "*" << i << ")";
    if (i != lanes - 1)
      os << ", ";
  }
  os << "))";
}

void CodeGenTileLangHIP::VisitExpr_(const BroadcastNode *op,
                                    std::ostream &os) { // NOLINT(*)
  int lanes = static_cast<int>(Downcast<IntImm>(op->lanes)->value);
  if ((op->dtype.is_int() || op->dtype.is_uint()) && op->dtype.bits() == 8 &&
      lanes == 4) {
    // make_int8x4
    const int64_t *p = as_const_int(op->value);
    ICHECK(p);
    int64_t v = *p & 0xFF;
    v = (v << 24) | (v << 16) | (v << 8) | v;
    if (op->dtype.is_uint()) {
      os << "(uint)" << v;
    } else {
      os << "(int)" << v;
    }
    return;
  }

  if (op->dtype.is_float16()) {
    std::string v = PrintExpr(op->value);
    os << "make_";
    PrintType(op->dtype, os);
    os << '(';
    for (int i = 0; i < lanes / 2; ++i) {
      if (i != 0)
        os << ", ";
      os << "__pack_half2(" << v << ", " << v << ")";
    }
    os << ')';
    return;
  }

  if (op->dtype.is_bfloat16()) {
    std::string v = PrintExpr(op->value);
    os << "make_";
    PrintType(op->dtype, os);
    os << '(';
    for (int i = 0; i < lanes / 2; ++i) {
      if (i != 0)
        os << ", ";
      os << "__pack_bfloat162(" << v << ", " << v << ")";
    }
    os << ')';
    return;
  }

  if (op->dtype.is_float() && op->dtype.bits() == 32 &&
      op->dtype.lanes() == 8) {
    std::string v = PrintExpr(op->value);
    // HIP does not allow taking the address of a temporary, so use a union
    // to reinterpret float2 as unsigned long long without UB or temp-address.
    os << "make_ulonglong4(";
    for (int i = 0; i < 4; ++i) {
      if (i != 0)
        os << ", ";
      os << "([&]{ union { float2 f; unsigned long long u; } _tmp;"
         << " _tmp.f = make_float2(" << v << ", " << v
         << "); return _tmp.u; }())";
    }
    os << ')';
    return;
  }

  if (op->dtype.is_float() && op->dtype.bits() == 32 &&
      (op->dtype.lanes() == 16 || op->dtype.lanes() == 32)) {
    // float32x16/float32x32: GCC vector extension — initialize with compound
    // literal.
    std::string v = PrintExpr(op->value);
    os << "(float32x" << op->dtype.lanes() << "){";
    for (int i = 0; i < op->dtype.lanes(); ++i) {
      if (i != 0)
        os << ", ";
      os << v;
    }
    os << "}";
    return;
  }

  if (op->dtype.is_bfloat16() && op->dtype.lanes() == 16) {
    // bfloat16x16: struct aggregate initializer.
    std::string v = PrintExpr(op->value);
    os << "bfloat16x16{";
    for (int i = 0; i < 16; ++i) {
      if (i != 0)
        os << ", ";
      os << v;
    }
    os << "}";
    return;
  }

  if ((op->dtype.is_int() || op->dtype.is_uint()) && op->dtype.bits() == 4) {
    bool fail = false;
    const int64_t *p = as_const_int(op->value);
    ICHECK(p);
    int64_t v = *p & 0xF;

    if (lanes == 4) {
      v = (v << 12) | (v << 8) | (v << 4) | v;
      if (op->dtype.is_uint()) {
        os << "(uint16_t)" << v;
      } else {
        os << "(int16_t)" << v;
      }
    } else {
      v = (v << 28) | (v << 24) | (v << 20) | (v << 16) | (v << 12) | (v << 8) |
          (v << 4) | v;
      if (lanes == 8) {
        if (op->dtype.is_uint()) {
          os << "(uint)" << v;
        } else {
          os << "(int)" << v;
        }
      } else if (lanes == 16 || lanes == 32) {
        os << "make_";
        PrintType(op->dtype, os);
        os << '(';
        for (int i = 0; i < lanes / 8; ++i) {
          if (i != 0)
            os << ", ";
          if (op->dtype.is_uint()) {
            os << "(uint)" << v;
          } else {
            os << "(int)" << v;
          }
        }
        os << ')';
      } else {
        fail = true;
      }
    }

    if (!fail) {
      return;
    }
  }

  std::string v = PrintExpr(op->value);
  os << "make_";
  PrintType(op->dtype, os);
  os << '(';
  for (int i = 0; i < lanes; ++i) {
    if (i != 0)
      os << ", ";
    os << v;
  }
  os << ')';
}

inline void PrintConst(const FloatImmNode *op, std::ostream &os,
                       CodeGenTileLangHIP *p) { // NOLINT(*)
  // Type code is kBFloat
  if (op->dtype.is_bfloat16()) {
    os << "bfloat16_t(";
    FloatImm const_f32 = FloatImm(DataType::Float(32), op->value);
    PrintConst(const_f32.get(), os, p);
    os << ')';
    return;
  } else if (op->dtype.is_float8_e4m3fnuz() || op->dtype.is_float8_e4m3() ||
             op->dtype.is_float8_e4m3fn()) {
    os << "fp8_e4_t(";
    FloatImm const_f32 = FloatImm(DataType::Float(32), op->value);
    PrintConst(const_f32.get(), os, p);
    os << ')';
    return;
  }
  // Type code is kFloat
  switch (op->dtype.bits()) {
  case 64:
  case 32: {
    std::ostringstream temp;
    if (std::isinf(op->value)) {
      if (op->value < 0) {
        temp << "-";
      }
      temp << ((op->dtype.bits() == 32) ? "HUGE_VALF" : "HUGE_VAL");
    } else if (std::isnan(op->value)) {
      temp << ((op->dtype.bits() == 32) ? "NAN" : "NAN");
    } else {
      temp << std::scientific << op->value;
      if (op->dtype.bits() == 32)
        temp << 'f';
    }
    p->MarkConst(temp.str());
    os << temp.str();
    break;
  }
  case 16: {
    os << "half_t" << '(';
    FloatImm const_f32 = FloatImm(DataType::Float(32), op->value);
    PrintConst(const_f32.get(), os, p);
    os << ')';
    break;
  }
  default:
    LOG(FATAL) << "Bad bit-width for float: " << op->dtype << "\n";
  }
}

void CodeGenTileLangHIP::VisitExpr_(const FloatImmNode *op,
                                    std::ostream &os) { // NOLINT(*)
  PrintConst(op, os, this);
}

void CodeGenTileLangHIP::HandleVolatileLoads(const std::string &value,
                                             const BufferLoadNode *op,
                                             std::ostream &os) {
  // Cast away volatile qualifier for fp16 types. That is, only loads and
  // stores are volatile. The loaded objects are not marked as volatile.
  //
  if ((op->dtype.is_float16() || op->dtype.is_bfloat16()) &&
      IsVolatile(op->buffer->data.get())) {
    os << "(";
    PrintType(op->dtype, os);
    os << ")(" << value << ")";
  } else {
    os << value;
  }
}

void CodeGenTileLangHIP::PrintVecElemLoadExpr(DataType t, int i,
                                              const std::string &value,
                                              std::ostream &os) {
  ICHECK_GT(t.lanes(), 1);
  if (t.bits() == 8 && (t.is_int() || t.is_uint())) {
    if (!(t.lanes() == 2 || t.lanes() == 3)) {
      if (i != 0) {
        os << "|";
      }
      os << "((0x000000ff << " << i * 8 << ") & (" << value << " << " << i * 8
         << "))";
      return;
    }
  }

  if (t.is_float16()) {
    if (i == 0) {
      os << "make_";
      PrintType(t, os);
      os << '(';
    }
    if (i % 2 == 0) {
      os << "__pack_half2(" << value;
    } else {
      os << "," << value << ")";
      if (i != t.lanes() - 1) {
        os << ",";
      } else {
        os << ")";
      }
    }
    return;
  }

  if (t.is_bfloat16() && t.lanes() != 16) {
    if (i == 0) {
      os << "make_";
      PrintType(t, os);
      os << '(';
    }
    if (i % 2 == 0) {
      os << "__pack_bfloat162(" << value;
    } else {
      os << "," << value << ")";
      if (i != t.lanes() - 1) {
        os << ",";
      } else {
        os << ")";
      }
    }
    return;
  }

  if ((t.lanes() == 16 || t.lanes() == 32) && t.bits() == 32 && t.is_float()) {
    // float32x16/float32x32: compound literal.
    if (i == 0)
      os << "(float32x" << t.lanes() << "){";
    os << value;
    if (i != t.lanes() - 1)
      os << ",";
    else
      os << "}";
    return;
  }

  if (t.lanes() == 16 && t.is_bfloat16()) {
    // bfloat16x16: struct aggregate initializer.
    if (i == 0)
      os << "bfloat16x16{";
    os << value;
    if (i != t.lanes() - 1)
      os << ",";
    else
      os << "}";
    return;
  }

  if (i == 0) {
    os << "make_";
    PrintType(t, os);
    os << "(";
  }
  os << value;
  if (i != t.lanes() - 1) {
    os << ",";
  } else {
    os << ")";
  }
  return;
}

void CodeGenTileLangHIP::AddFunction(const PrimFunc &f) {
  // clear previous generated state.
  this->InitFuncState(f);
  // reserve keywords
  ReserveKeywordsAsUnique();

  auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
  ICHECK(global_symbol.has_value())
      << "CodeGenC: Expect PrimFunc to have the global_symbol attribute";
  bool no_alias = f->HasNonzeroAttr(tirx::attr::kNoAlias);
  std::unordered_set<const VarNode *> non_restrict;
  if (auto opt = f->GetAttr<Array<tirx::Var>>(tl::attr::kNonRestrictParams)) {
    for (const tirx::Var &v : opt.value())
      non_restrict.insert(v.get());
  }

  this->PrintFuncPrefix(stream);
  CodeGenC::PrintType(f->ret_type, stream);
  this->PrintExtraAttrs(f, stream);
  this->stream << " " << static_cast<std::string>(global_symbol.value()) << "(";

  for (size_t i = 0; i < f->params.size(); ++i) {
    tirx::Var v = f->params[i];
    std::string vid = AllocVarID(v.get());
    if (i != 0)
      stream << ", ";
    if (v.dtype().is_handle()) {
      // work around for grid constant parameters.
      if (auto *ptr = v->type_annotation.as<PointerTypeNode>()) {
        if (ptr->storage_scope == "grid_constant") {
          stream << "__grid_constant__ const ";
          CodeGenC::PrintType(ptr->element_type, stream);
          stream << ' ' << vid;
          continue;
        }
      }

      auto it = alloc_storage_scope_.find(v.get());
      if (it != alloc_storage_scope_.end()) {
        PrintStorageScope(it->second, stream);
      }

      CodeGenC::PrintType(GetType(v), stream);
      if (auto *ptr = v->type_annotation.as<PointerTypeNode>()) {
        if (auto *prim = ptr->element_type.as<PrimTypeNode>()) {
          RegisterHandleType(v.get(), prim->dtype);
        }
      }

      if (no_alias && !non_restrict.count(v.get())) {
        PrintRestrict(v, stream);
      }
    } else {
      CodeGenC::PrintType(GetType(v), stream);
    }
    stream << ' ' << vid;
  }
  stream << ") {\n";
  this->PreFunctionBody(f);
  int func_scope = this->BeginScope();
  this->PrintStmt(f->body);
  this->EndScope(func_scope);
  this->PrintIndent();
  this->stream << "}\n\n";
}

} // namespace codegen
} // namespace tvm
