Cherry pick salu promotion

waveasm/include/waveasm/Transforms/TranslateFromMLIR.h

@@ -385,6 +385,26 @@ class TranslationContext {
     return pendingSRDs.size() + pendingScalarArgs.size();
   }
 
+  /// Get the number of args that fit in hardware preload SGPRs.
+  /// Uses all-or-nothing strategy: preloads ALL args when they fit,
+  /// otherwise preloads NOTHING. Partial preloading (SRDs only) is
+  /// avoided because mixing HW preloading with explicit s_load from
+  /// the same kernarg buffer produces incorrect values on GFX950.
+  static constexpr int64_t kMaxPreloadSGPRs = 16;
+  size_t getNumPreloadedArgs() const {
+    size_t total = pendingSRDs.size() + pendingScalarArgs.size();
+    if (total * 2 <= static_cast<size_t>(kMaxPreloadSGPRs))
+      return total;
+    return 0;
+  }
+
+  /// Whether scalar kernel args are hardware-preloaded (true when all args fit
+  /// in the preload limit) or must be loaded via explicit s_load instructions.
+  bool areScalarsPreloaded() const {
+    size_t total = pendingSRDs.size() + pendingScalarArgs.size();
+    return total * 2 <= static_cast<size_t>(kMaxPreloadSGPRs);
+  }
+
   //===--------------------------------------------------------------------===//
   // Split Vector Result Tracking
   //===--------------------------------------------------------------------===//
@@ -610,12 +630,11 @@ class TranslationContext {
 
   /// Get the number of user SGPRs (for computing system SGPR offsets)
   /// User SGPRs include: kernarg ptr (2) + preloaded args (on gfx950+)
+  /// Hardware limits user SGPRs to 16 on gfx950. Args that don't fit are
+  /// loaded via explicit s_load into overflow SGPR slots after the SRDs.
   int64_t getUserSgprCount() const {
-    // Base: 2 SGPRs for kernarg_segment_ptr
-    int64_t count = 2;
-    // On gfx950+ with kernarg preloading, add preloaded args
+    int64_t count = 2; // kernarg_segment_ptr
     if (llvm::isa<GFX950TargetAttr>(target)) {
-      // Each kernel arg uses 2 SGPRs, capped at 14 (hardware max 16 total).
       count += std::min(size_t(14), getNumKernelArgs() * 2);
     }
     return count;

waveasm/include/waveasm/Transforms/Utils.h

@@ -15,14 +15,19 @@
 namespace waveasm {
 
 /// Extract an integer constant from a Value.
-/// Handles ConstantOp directly, or V_MOV_B32(ConstantOp), or ImmType.
+/// Handles ConstantOp directly, V_MOV_B32(ConstantOp), S_MOV_B32(ConstantOp),
+/// or ImmType.
 inline std::optional<int64_t> getConstantValue(mlir::Value v) {
   if (auto constOp = v.getDefiningOp<ConstantOp>())
     return constOp.getValue();
   if (auto movOp = v.getDefiningOp<V_MOV_B32>()) {
     if (auto constOp = movOp.getSrc().getDefiningOp<ConstantOp>())
       return constOp.getValue();
   }
+  if (auto movOp = v.getDefiningOp<S_MOV_B32>()) {
+    if (auto constOp = movOp.getSrc().getDefiningOp<ConstantOp>())
+      return constOp.getValue();
+  }
   if (auto immType = llvm::dyn_cast<ImmType>(v.getType()))
     return immType.getValue();
   return std::nullopt;

waveasm/lib/Transforms/AssemblyEmitter.cpp

@@ -169,9 +169,16 @@ std::string KernelGenerator::emitBufferStore(Operation *op,
   // Use VMEMStoreOpInterface to access operands by name
   if (auto storeOp = dyn_cast<VMEMStoreOpInterface>(op)) {
     std::string vdata = resolveValue(storeOp.getData());
-    std::string voffset = resolveValue(storeOp.getVoffset());
+    Value voffsetVal = storeOp.getVoffset();
     std::string srd = resolveValue(storeOp.getSaddr());
-    result += " " + vdata + ", " + voffset + ", " + srd + ", 0 offen";
+
+    if (isa<ImmType>(voffsetVal.getType())) {
+      result += " " + vdata + ", off, " + srd + ", 0";
+    } else {
+      std::string voffset = resolveValue(voffsetVal);
+      result += " " + vdata + ", " + voffset + ", " + srd + ", 0 offen";
+    }
+
     if (auto instOffsetAttr = op->getAttrOfType<IntegerAttr>("instOffset")) {
       int64_t offset = instOffsetAttr.getInt();
       if (offset > 0) {
@@ -290,11 +297,13 @@ std::string KernelGenerator::emitDefaultFormat(Operation *op,
   }
 
   for (Value result : op->getResults()) {
+    // Skip SCC-typed results (hardware-implicit, not emitted).
     if (isa<SCCType>(result.getType()))
       continue;
     operands.push_back(resolveValue(result));
   }
   for (Value operand : op->getOperands()) {
+    // Skip SCC-typed operands (hardware-implicit, not emitted).
     if (isa<SCCType>(operand.getType()))
       continue;
     if (isScalarOp) {
@@ -348,8 +357,12 @@ KernelGenerator::emitScaledMFMA(Operation *scaledOp, llvm::StringRef mnemonic) {
 
 std::optional<std::string> KernelGenerator::generateOp(Operation *op) {
   return llvm::TypeSwitch<Operation *, std::optional<std::string>>(op)
-      .Case<ProgramOp, LabelOp, CommentOp, RawOp>(
+      .Case<ProgramOp, LabelOp, CommentOp, PrecoloredVRegOp, PrecoloredSRegOp,
+            PrecoloredARegOp, ConstantOp, PackOp, ExtractOp, DCEProtectOp>(
           [](auto) { return std::nullopt; })
+      .Case<RawOp>([&](RawOp rawOp) -> std::optional<std::string> {
+        return generateRaw(rawOp);
+      })
 
       .Case<S_WAITCNT>([&](S_WAITCNT waitcntOp) {
         std::optional<int64_t> vmcnt, lgkmcnt, expcnt;
@@ -539,16 +552,18 @@ std::optional<std::string> KernelGenerator::generateOp(Operation *op) {
             std::string src = resolveValue(srcVal);
             std::string lines;
             if (isAGPR) {
-              // v_accvgpr_write_b32 requires a VGPR source in this backend.
-              // Materialize immediate sources into the reserved scratch VGPR.
+              auto [isLit, litVal] = getLiteralValue(srcVal);
               std::string writeSrc = src;
-              if (srcIsImm) {
+              if (srcIsImm && !(isLit && isInlineConstant(litVal))) {
+                // Non-inline literal: materialize into scratch VGPR first.
                 lines += "  v_mov_b32 " + formatVGPRRange(kScratchVGPR, 1) +
                          ", " + src;
                 writeSrc = formatVGPRRange(kScratchVGPR, 1);
                 peakVGPRs = std::max(peakVGPRs, kScratchVGPR + 1);
                 invalidateScratchCache();
               }
+              // Inline constants (e.g. 0) go directly into
+              // v_accvgpr_write_b32 without a scratch VGPR.
               for (int64_t i = 0; i < size; ++i) {
                 if (!lines.empty())
                   lines += "\n";
@@ -569,6 +584,12 @@ std::optional<std::string> KernelGenerator::generateOp(Operation *op) {
         if (isAGPR) {
           if (srcIsImm) {
             std::string src = resolveValue(srcVal);
+            auto [isLit, litVal] = getLiteralValue(srcVal);
+            if (isLit && isInlineConstant(litVal)) {
+              // Inline constant: emit directly without scratch VGPR.
+              return "  v_accvgpr_write_b32 " + resolveValue(result) + ", " +
+                     src;
+            }
             std::string scratch = formatVGPRRange(kScratchVGPR, 1);
             peakVGPRs = std::max(peakVGPRs, kScratchVGPR + 1);
             invalidateScratchCache();
@@ -701,10 +722,7 @@ std::optional<std::string> KernelGenerator::generateOp(Operation *op) {
 
               SmallVector<bool> handled(pendingCopies.size(), false);
 
-              // Allocate swap temps once and reuse across all swaps.
-              // Swaps are emitted sequentially so the temp is dead after
-              // each 3-instruction sequence and can be reused.
-              int64_t vSwapTemp = -1;
+              // SGPR swaps still need a temp; VGPR swaps use v_swap_b32.
               int64_t sSwapTemp = -1;
 
               for (size_t i = 0; i < pendingCopies.size(); ++i) {
@@ -734,13 +752,7 @@ std::optional<std::string> KernelGenerator::generateOp(Operation *op) {
                     }
                     int64_t regA = pendingCopies[i].dst;
                     int64_t regB = pendingCopies[j].dst;
-                    if (vSwapTemp < 0) {
-                      vSwapTemp = peakVGPRs;
-                      peakVGPRs = std::max(peakVGPRs, vSwapTemp + 1);
-                    }
-                    os << "  v_mov_b32 v" << vSwapTemp << ", v" << regA << "\n";
-                    os << "  v_mov_b32 v" << regA << ", v" << regB << "\n";
-                    os << "  v_mov_b32 v" << regB << ", v" << vSwapTemp << "\n";
+                    os << "  v_swap_b32 v" << regA << ", v" << regB << "\n";
                     handled[i] = true;
                     handled[j] = true;
                     break;
@@ -856,6 +868,7 @@ std::optional<std::string> KernelGenerator::generateOp(Operation *op) {
             if (opName.starts_with("waveasm.")) {
               mnemonic = opName.drop_front(8);
             }
+            // s_cmp_ne_* → s_cmp_lg_* (ISA mnemonic)
             std::string mnemStr;
             if (mnemonic.contains("_ne_")) {
               mnemStr = mnemonic.str();
@@ -871,6 +884,7 @@ std::optional<std::string> KernelGenerator::generateOp(Operation *op) {
           })
 
       // SALU arithmetic ops that set SCC: emit dst and operands, skip scc.
+      // S_ADDC_U32 has an explicit SCC-in operand that must also be skipped.
       .Case<S_ADD_U32, S_ADDC_U32, S_ADD_I32, S_SUB_U32, S_SUB_I32>(
           [&](auto addOp) -> std::optional<std::string> {
             llvm::StringRef opName = addOp->getName().getStringRef();
@@ -882,6 +896,7 @@ std::optional<std::string> KernelGenerator::generateOp(Operation *op) {
             // Only emit the first result (dst), not the second (scc)
             operands.push_back(resolveValue(addOp.getDst()));
             for (Value operand : addOp->getOperands()) {
+              // Skip SCC-typed operands (carry-in for s_addc_u32).
               if (isa<SCCType>(operand.getType()))
                 continue;
               operands.push_back(resolveValue(operand));
@@ -994,6 +1009,7 @@ std::optional<std::string> KernelGenerator::generateOp(Operation *op) {
                                              operands);
           })
 
+      // S_CSELECT_B32: emit dst, src0, src1 — skip the SCC-in operand.
       .Case<S_CSELECT_B32>(
           [&](S_CSELECT_B32 selOp) -> std::optional<std::string> {
             llvm::SmallVector<std::string> operands;
@@ -1082,6 +1098,8 @@ llvm::SmallVector<std::string> KernelGenerator::generate() {
   });
   peakVGPRs = std::max(peakVGPRs, int64_t(1));
   peakSGPRs = std::max(peakSGPRs, int64_t(2));
+  if (auto minSgprs = program->getAttrOfType<IntegerAttr>("min_sgprs"))
+    peakSGPRs = std::max(peakSGPRs, minSgprs.getInt());
 
   for (Operation &op : program.getBodyBlock()) {
     if (auto labelOp = dyn_cast<LabelOp>(op)) {

waveasm/lib/Transforms/LiteralMaterialization.cpp

@@ -33,11 +33,12 @@ static bool isInlineConstant(int64_t val) {
 }
 
 static const llvm::StringSet<> &getVOP2Instructions() {
+  // v_add_u32, v_sub_u32, v_subrev_u32 are VOP3-only on GFX9+ (the VOP2
+  // carry-producing variants are v_add_co_u32 / v_sub_co_u32 /
+  // v_subrev_co_u32). VOP3 does not support literal operands.
   static const llvm::StringSet<> kVOP2 = {
-      "v_add_u32",     "v_sub_u32", "v_subrev_u32",  "v_and_b32",
-      "v_or_b32",      "v_xor_b32", "v_lshlrev_b32", "v_lshrrev_b32",
-      "v_ashrrev_i32", "v_max_u32", "v_min_u32",     "v_add_i32",
-      "v_sub_i32",
+      "v_and_b32",     "v_or_b32",      "v_xor_b32", "v_lshlrev_b32",
+      "v_lshrrev_b32", "v_ashrrev_i32", "v_max_u32", "v_min_u32",
   };
   return kVOP2;
 }