[SYCLomatic][PTX] Support migration of PTX instruction ld.cg.global.v2.s16 and ld.cg.global.v4.s16 (#2769)

Signed-off-by: chenwei.sun <chenwei.sun@intel.com>

Author: tomflinda

clang/lib/DPCT/RulesAsm/AsmMigration.cpp

@@ -609,7 +609,7 @@ bool SYCLGenBase::emitAddressExpr(const InlineAsmAddressExpr *Dst) {
     return false;
   };
 
-  if (CurrInst->is(asmtok::op_ld, asmtok::op_red))
+  if (CurrInst->is(asmtok::op_red))
     OS() << "*";
   switch (Dst->getMemoryOpKind()) {
   case InlineAsmAddressExpr::Imm:
@@ -633,7 +633,7 @@ bool SYCLGenBase::emitAddressExpr(const InlineAsmAddressExpr *Dst) {
     if (tryEmitStmt(Reg, Dst->getSymbol()))
       return SYCLGenSuccess();
 
-    if (CurrInst->is(asmtok::op_st))
+    if (CurrInst->is(asmtok::op_st, asmtok::op_ld))
       OS() << llvm::formatv("(uintptr_t){0}", Reg);
     else
       OS() << llvm::formatv("(({0} *)((uintptr_t){1} + {2}))", Type, Reg,
@@ -2694,16 +2694,10 @@ class SYCLGen : public SYCLGenBase {
     return SYCLGenSuccess();
   }
 
-  bool HandleStVec(const InlineAsmInstruction *Inst, int VecNum) {
-    std::string Ops;
-    if (tryEmitStmt(Ops, Inst->getInputOperand(0)))
-      return SYCLGenError();
-
-    // To extract the values from the string like "{x, y, z, w}" and store them
-    // int Values vector
+  std::vector<std::string> extractValues(const std::string &Ops) {
     std::vector<std::string> Values;
-    size_t start = 1;                  // Skip the '{' character
-    size_t end = Ops.find(',', start); // Find the first comma
+    size_t start = 1; // Skip the '{' character
+    size_t end = Ops.find(',', start);
 
     while (end != std::string::npos) {
       std::string Token = Ops.substr(start, end - start);
@@ -2720,11 +2714,24 @@ class SYCLGen : public SYCLGenBase {
     std::string token = Ops.substr(start, Ops.size() - start - 1);
     size_t first = token.find_first_not_of(' ');
     size_t last = token.find_last_not_of(' ');
-
     if (first != std::string::npos && last != std::string::npos) {
       Values.push_back(token.substr(first, last - first + 1));
     }
 
+    return Values;
+  }
+
+  bool HandleStVec(const InlineAsmInstruction *Inst, int VecNum) {
+    std::string Ops;
+    if (tryEmitStmt(Ops, Inst->getInputOperand(0)))
+      return SYCLGenError();
+
+    // To extract the values from the string like "{x, y, z, w}" and store them
+    // int Values vector
+    std::vector<std::string> Values = extractValues(Ops);
+    if (Values.size() != (size_t)VecNum)
+      return SYCLGenError();
+
     std::string Output;
     if (tryEmitStmt(Output, Inst->getOutputOperand()))
       return SYCLGenError();
@@ -2790,6 +2797,43 @@ class SYCLGen : public SYCLGenBase {
     return SYCLGenSuccess();
   }
 
+  bool HandleLdVec(const InlineAsmInstruction *Inst, int VecNum) {
+    std::string Ops;
+    if (tryEmitStmt(Ops, Inst->getOutputOperand()))
+      return SYCLGenError();
+
+    // To extract the values from the string like "{x, y, z, w}" and store them
+    // int Values vector
+    std::vector<std::string> Values = extractValues(Ops);
+    if (Values.size() != (size_t)VecNum)
+      return SYCLGenError();
+
+    std::string Input;
+    if (tryEmitStmt(Input, Inst->getInputOperand(0)))
+      return SYCLGenError();
+
+    std::string Type;
+    if (tryEmitType(Type, Inst->getType(0)))
+      return SYCLGenError();
+
+    const auto *Src =
+        dyn_cast_or_null<InlineAsmAddressExpr>(Inst->getInputOperand(0));
+    if (!Src)
+      return SYCLGenError();
+
+    for (int Index = 0; Index < VecNum; Index++) {
+      OS() << llvm::formatv("{0}  = *(({1} *){2} + {3}){4}", Values[Index],
+                            Type, Input, Index,
+                            Index == VecNum - 1 ? "" : ";\n");
+      if (Index < VecNum - 1) {
+        indent();
+      }
+    }
+
+    endstmt();
+    return SYCLGenSuccess();
+  }
+
   bool handle_ld(const InlineAsmInstruction *Inst) override {
     if (Inst->getNumInputOperands() != 1)
       return SYCLGenError();
@@ -2801,14 +2845,32 @@ class SYCLGen : public SYCLGenBase {
 
     if (!Src)
       return SYCLGenError();
+
+    if (Inst->hasAttr(InstAttr::cg)) {
+      if (Inst->hasAttr(InstAttr::v4))
+        return HandleLdVec(Inst, 4);
+      if (Inst->hasAttr(InstAttr::v2))
+        return HandleLdVec(Inst, 2);
+    }
+
     std::string Type;
     if (tryEmitType(Type, Inst->getType(0)))
       return SYCLGenError();
     if (emitStmt(Dst))
       return SYCLGenError();
     OS() << " = ";
-    if (emitStmt(Src))
+
+    std::string InOp;
+    if (tryEmitStmt(InOp, Inst->getInputOperand(0)))
       return SYCLGenError();
+
+    if (Src->getMemoryOpKind() == InlineAsmAddressExpr::RegImm) {
+      OS() << llvm::formatv("*(({0} *)({1} + {2}))", Type, InOp,
+                            Src->getImmAddr()->getValue().getZExtValue());
+    } else {
+      OS() << "*" << InOp;
+    }
+
     endstmt();
     return SYCLGenSuccess();
   }

clang/test/dpct/asm/ld.cu

@@ -36,4 +36,38 @@ __device__ void shared_address_load32(uint32_t addr, uint32_t &val) {
   asm volatile("{ld.shared.b32 %0, [%1];}" : : "r"(val), "r"(addr) : "memory"); 
 }
 
+// CHECK: inline void load_global_short2(sycl::short2 &a, const sycl::short2 *addr) {
+// CHECK-NEXT:  short x, y, z, w;
+// CHECK-NEXT:  x  = *((int16_t *)(uintptr_t)addr + 0);
+// CHECK-NEXT:  y  = *((int16_t *)(uintptr_t)addr + 1);
+// CHECK-NEXT:  a.x() = x;
+// CHECK-NEXT:  a.y() = y;
+// CHECK-NEXT:}
+__device__ inline void load_global_short2(short2 &a, const short2 *addr) {
+  short x, y, z, w;
+  asm("ld.cg.global.v2.s16 {%0, %1}, [%2+0];" : "=h"(x), "=h"(y) : "l"(addr));
+  a.x = x;
+  a.y = y;
+}
+
+// CHECK: inline void load_global_short4(sycl::short4 &a, const sycl::short4 *addr) {
+// CHECK-NEXT:  short x, y, z, w;
+// CHECK-NEXT:  x  = *((int16_t *)(uintptr_t)addr + 0);
+// CHECK-NEXT:  y  = *((int16_t *)(uintptr_t)addr + 1);
+// CHECK-NEXT:  z  = *((int16_t *)(uintptr_t)addr + 2);
+// CHECK-NEXT:  w  = *((int16_t *)(uintptr_t)addr + 3);
+// CHECK-NEXT:  a.x() = x;
+// CHECK-NEXT:  a.y() = y;
+// CHECK-NEXT:  a.z() = z;
+// CHECK-NEXT:  a.w() = w;
+// CHECK-NEXT:}
+__device__ inline void load_global_short4(short4 &a, const short4 *addr) {
+  short x, y, z, w;
+  asm("ld.cg.global.v4.s16 {%0, %1, %2, %3}, [%4+0];" : "=h"(x), "=h"(y), "=h"(z), "=h"(w) : "l"(addr));
+  a.x = x;
+  a.y = y;
+  a.z = z;
+  a.w = w;
+}
+
 // clang-format on