[AMDGPU] Add MIMG NSA threshold configuration attribute

[critson]

Summary

Make MIMG NSA minimum addresses threshold an attribute that can
be set on a function or configured via command line.
This enables frontend tuning which allows increased NSA usage
where beneficial.

[foad]

Seems reasonable.

I would also be in favour of changing the default to 2. That would tend to introduce fewer "mov"s to shuffle data around, at the expense of using a larger encoding on average for mimg instructions. But I think that is a good trade-off to make for the sake of performance.

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

959

Use max()

963

Use max()

[critson]

In D134780#3820209, @foad wrote:

Seems reasonable.

I would also be in favour of changing the default to 2. That would tend to introduce fewer "mov"s to shuffle data around, at the expense of using a larger encoding on average for mimg instructions. But I think that is a good trade-off to make for the sake of performance.

Thank you for the quick review!
I did test the idea of using default 2, but I am not convinced yet that the benefit is universal enough.
Out of ~20000 pipelines, ~2000 had higher VGPR usage with threshold 2 and ~1000 had lower VGPR usage.
When VGPR usage decreases, it does so more beneficial ways though.
Spilling decreased in 2 pipelines, and only increased in 1.

[foad]

Out of ~20000 pipelines, ~2000 had higher VGPR usage with threshold 2 and ~1000 had lower VGPR usage.

That's weird. I can't see why enabling NSA would consistently cause higher vgpr usage.

[critson]

• Address reviewer feedback

[foad]

LGTM, thanks.

[foad]

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

59

On the "don't repeat yourself" principle you could either remove the cl::init here...

959

I don't think you need the .getValue().

965

... or return NSAThreshold here.

[foad]

Out of ~20000 pipelines, ~2000 had higher VGPR usage with threshold 2 and ~1000 had lower VGPR usage.

That's weird. I can't see why enabling NSA would consistently cause higher vgpr usage.

FYI I looked at one case where this happens and it was caused by GCNNSAReassign making strange (well, different) decisions. So now I need to try to understand what that pass does.

[critson]

In D134780#3821260, @foad wrote:

Out of ~20000 pipelines, ~2000 had higher VGPR usage with threshold 2 and ~1000 had lower VGPR usage.

That's weird. I can't see why enabling NSA would consistently cause higher vgpr usage.

FYI I looked at one case where this happens and it was caused by GCNNSAReassign making strange (well, different) decisions. So now I need to try to understand what that pass does.

Yes, I think we probably need to revisit the logic in GCNNSAReassign.
I seem to remember it lack some support for subregisters which might be part of the issue.

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

59

True, it would be nice to have this number only once.

959

Pretty sure it is required, because std:max does not know how to handle cl::opt.

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

 /// want a selected instruction entering RegBankSelect. In order to avoid
 /// defining a multitude of intermediate image instructions, directly hack on
 /// the intrinsic's arguments. In cases like a16 addresses, this requires
 /// padding now unnecessary arguments with $noreg.
 bool AMDGPULegalizerInfo::legalizeImageIntrinsic(
     MachineInstr &MI, MachineIRBuilder &B, GISelChangeObserver &Observer,
     const AMDGPU::ImageDimIntrinsicInfo *Intr) const {
 
+  const MachineFunction &MF = *MI.getMF();
   const unsigned NumDefs = MI.getNumExplicitDefs();
   const unsigned ArgOffset = NumDefs + 1;
   bool IsTFE = NumDefs == 2;
   // We are only processing the operands of d16 image operations on subtargets
   // that use the unpacked register layout, or need to repack the TFE result.
 
   // TODO: Do we need to guard against already legalized intrinsics?
   const AMDGPU::MIMGBaseOpcodeInfo *BaseOpcode =
bool AMDGPULegalizerInfo::legalizeImageIntrinsic(
   if (IsA16 || IsG16) {
     if (Intr->NumVAddrs > 1) {
       SmallVector<Register, 4> PackedRegs;
 
       packImage16bitOpsToDwords(B, MI, PackedRegs, ArgOffset, Intr, IsA16,
                                 IsG16);
 
       // See also below in the non-a16 branch
-      const bool UseNSA = ST.hasNSAEncoding() && PackedRegs.size() >= 3 &&
+      const bool UseNSA = ST.hasNSAEncoding() &&
+                          PackedRegs.size() >= ST.getNSAThreshold(MF) &&
                           PackedRegs.size() <= ST.getNSAMaxSize();
 
       if (!UseNSA && PackedRegs.size() > 1) {
         LLT PackedAddrTy = LLT::fixed_vector(2 * PackedRegs.size(), 16);
         auto Concat = B.buildConcatVectors(PackedAddrTy, PackedRegs);
         PackedRegs[0] = Concat.getReg(0);
         PackedRegs.resize(1);
       }
if (IsA16 || IsG16) {
     // do so, so force non-NSA for the common 2-address case as a heuristic.
     //
     // SIShrinkInstructions will convert NSA encodings to non-NSA after register
     // allocation when possible.
     //
     // TODO: we can actually allow partial NSA where the final register is a
     // contiguous set of the remaining addresses.
     // This could help where there are more addresses than supported.
-    const bool UseNSA = ST.hasNSAEncoding() && CorrectedNumVAddrs >= 3 &&
+    const bool UseNSA = ST.hasNSAEncoding() &&
+                        CorrectedNumVAddrs >= ST.getNSAThreshold(MF) &&
                         CorrectedNumVAddrs <= ST.getNSAMaxSize();
 
     if (!UseNSA && Intr->NumVAddrs > 1)
       convertImageAddrToPacked(B, MI, ArgOffset + Intr->VAddrStart,
                                Intr->NumVAddrs);
   }
 
   int Flags = 0;

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

static cl::opt<bool> EnableVGPRIndexMode(
   "amdgpu-vgpr-index-mode",
   cl::desc("Use GPR indexing mode instead of movrel for vector indexing"),
   cl::init(false));
 
 static cl::opt<bool> UseAA("amdgpu-use-aa-in-codegen",
                            cl::desc("Enable the use of AA during codegen."),
                            cl::init(true));
 
+static cl::opt<unsigned> NSAThreshold("amdgpu-nsa-threshold",
+                                      cl::desc("Number of addresses from which to enable MIMG NSA."),
+                                      cl::init(3), cl::Hidden);
+
 GCNSubtarget::~GCNSubtarget() = default;
 
 GCNSubtarget &
 GCNSubtarget::initializeSubtargetDependencies(const Triple &TT,
                                               StringRef GPU, StringRef FS) {
   // Determine default and user-specified characteristics
   //
   // We want to be able to turn these off, but making this a subtarget feature

 }
 
 std::unique_ptr<ScheduleDAGMutation>
 GCNSubtarget::createFillMFMAShadowMutation(const TargetInstrInfo *TII) const {
   return EnablePowerSched ? std::make_unique<FillMFMAShadowMutation>(&InstrInfo)
                           : nullptr;
 }
 
+unsigned GCNSubtarget::getNSAThreshold(const MachineFunction &MF) const {
+  if (NSAThreshold.getNumOccurrences() > 0)
+    return std::max(NSAThreshold.getValue(), 2u);
+
+  int Value = AMDGPU::getIntegerAttribute(MF.getFunction(), "amdgpu-nsa-threshold", -1);
+  if (Value > 0)
+    return std::max(Value, 2);
+
+  return 3;
+}
+
 const AMDGPUSubtarget &AMDGPUSubtarget::get(const MachineFunction &MF) {
   if (MF.getTarget().getTargetTriple().getArch() == Triple::amdgcn)
     return static_cast<const AMDGPUSubtarget&>(MF.getSubtarget<GCNSubtarget>());
   else
     return static_cast<const AMDGPUSubtarget&>(MF.getSubtarget<R600Subtarget>());
 }
 
 const AMDGPUSubtarget &AMDGPUSubtarget::get(const TargetMachine &TM, const Function &F) {
   if (TM.getTargetTriple().getArch() == Triple::amdgcn)
     return static_cast<const AMDGPUSubtarget&>(TM.getSubtarget<GCNSubtarget>(F));
   else
     return static_cast<const AMDGPUSubtarget&>(TM.getSubtarget<R600Subtarget>(F));
 }

llvm/lib/Target/AMDGPU/GCNSubtarget.h

   }
 
   void adjustSchedDependency(SUnit *Def, int DefOpIdx, SUnit *Use, int UseOpIdx,
                              SDep &Dep) const override;
 
   // \returns true if it's beneficial on this subtarget for the scheduler to
   // cluster stores as well as loads.
   bool shouldClusterStores() const { return getGeneration() >= GFX11; }
+
+  // \returns the number of address arguments from which to enable MIMG NSA
+  // on supported architectures.
+  unsigned getNSAThreshold(const MachineFunction &MF) const;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AMDGPU_GCNSUBTARGET_H

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::lowerImage(SDValue Op,
   //
   // SIShrinkInstructions will convert NSA encodings to non-NSA after register
   // allocation when possible.
   //
   // TODO: we can actually allow partial NSA where the final register is a
   // contiguous set of the remaining addresses.
   // This could help where there are more addresses than supported.
   bool UseNSA = ST->hasFeature(AMDGPU::FeatureNSAEncoding) &&
-                VAddrs.size() >= 3 &&
+                VAddrs.size() >= (unsigned)ST->getNSAThreshold(MF) &&
                 VAddrs.size() <= (unsigned)ST->getNSAMaxSize();
   SDValue VAddr;
   if (!UseNSA)
     VAddr = getBuildDwordsVector(DAG, DL, VAddrs);
 
   SDValue True = DAG.getTargetConstant(1, DL, MVT::i1);
   SDValue False = DAG.getTargetConstant(0, DL, MVT::i1);
   SDValue Unorm;

llvm/test/CodeGen/AMDGPU/amdgpu-nsa-threshold.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -march=amdgcn -mcpu=gfx1100 -verify-machineinstrs < %s | FileCheck -check-prefix=ATTRIB %s
+; RUN: llc -march=amdgcn -mcpu=gfx1100 -amdgpu-nsa-threshold=2 -verify-machineinstrs < %s | FileCheck -check-prefix=FORCE-2 %s
+; RUN: llc -march=amdgcn -mcpu=gfx1100 -amdgpu-nsa-threshold=3 -verify-machineinstrs < %s | FileCheck -check-prefix=FORCE-3 %s
+; RUN: llc -march=amdgcn -mcpu=gfx1100 -amdgpu-nsa-threshold=4 -verify-machineinstrs < %s | FileCheck -check-prefix=FORCE-4 %s
+
+; Note: command line argument should override function attribute.
+
+define amdgpu_ps <4 x float> @sample_2d_nsa2(<8 x i32> inreg %rsrc, <4 x i32> inreg %samp, float %t, float %s) #2 {
+; ATTRIB-LABEL: sample_2d_nsa2:
+; ATTRIB:       ; %bb.0: ; %main_body
+; ATTRIB-NEXT:    s_mov_b32 s12, exec_lo
+; ATTRIB-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; ATTRIB-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; ATTRIB-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; ATTRIB-NEXT:    image_sample v[0:3], [v1, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; ATTRIB-NEXT:    s_waitcnt vmcnt(0)
+; ATTRIB-NEXT:    ; return to shader part epilog
+;
+; FORCE-2-LABEL: sample_2d_nsa2:
+; FORCE-2:       ; %bb.0: ; %main_body
+; FORCE-2-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-2-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-2-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; FORCE-2-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-2-NEXT:    image_sample v[0:3], [v1, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; FORCE-2-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-2-NEXT:    ; return to shader part epilog
+;
+; FORCE-3-LABEL: sample_2d_nsa2:
+; FORCE-3:       ; %bb.0: ; %main_body
+; FORCE-3-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-3-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-3-NEXT:    v_mov_b32_e32 v2, v0
+; FORCE-3-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-3-NEXT:    image_sample v[0:3], v[1:2], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; FORCE-3-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-3-NEXT:    ; return to shader part epilog
+;
+; FORCE-4-LABEL: sample_2d_nsa2:
+; FORCE-4:       ; %bb.0: ; %main_body
+; FORCE-4-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-4-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-4-NEXT:    v_mov_b32_e32 v2, v0
+; FORCE-4-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-4-NEXT:    image_sample v[0:3], v[1:2], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; FORCE-4-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-4-NEXT:    ; return to shader part epilog
+main_body:
+  %v = call <4 x float> @llvm.amdgcn.image.sample.2d.v4f32.f32(i32 15, float %s, float %t, <8 x i32> %rsrc, <4 x i32> %samp, i1 0, i32 0, i32 0)
+  ret <4 x float> %v
+}
+
+define amdgpu_ps <4 x float> @sample_3d_nsa2(<8 x i32> inreg %rsrc, <4 x i32> inreg %samp, float %r, float %s, float %t) #2 {
+; ATTRIB-LABEL: sample_3d_nsa2:
+; ATTRIB:       ; %bb.0: ; %main_body
+; ATTRIB-NEXT:    s_mov_b32 s12, exec_lo
+; ATTRIB-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; ATTRIB-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; ATTRIB-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; ATTRIB-NEXT:    image_sample v[0:3], [v1, v2, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; ATTRIB-NEXT:    s_waitcnt vmcnt(0)
+; ATTRIB-NEXT:    ; return to shader part epilog
+;
+; FORCE-2-LABEL: sample_3d_nsa2:
+; FORCE-2:       ; %bb.0: ; %main_body
+; FORCE-2-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-2-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-2-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; FORCE-2-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-2-NEXT:    image_sample v[0:3], [v1, v2, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; FORCE-2-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-2-NEXT:    ; return to shader part epilog
+;
+; FORCE-3-LABEL: sample_3d_nsa2:
+; FORCE-3:       ; %bb.0: ; %main_body
+; FORCE-3-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-3-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-3-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; FORCE-3-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-3-NEXT:    image_sample v[0:3], [v1, v2, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; FORCE-3-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-3-NEXT:    ; return to shader part epilog
+;
+; FORCE-4-LABEL: sample_3d_nsa2:
+; FORCE-4:       ; %bb.0: ; %main_body
+; FORCE-4-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-4-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-4-NEXT:    v_mov_b32_e32 v3, v0
+; FORCE-4-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-4-NEXT:    image_sample v[0:3], v[1:3], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; FORCE-4-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-4-NEXT:    ; return to shader part epilog
+main_body:
+  %v = call <4 x float> @llvm.amdgcn.image.sample.3d.v4f32.f32(i32 15, float %s, float %t, float %r, <8 x i32> %rsrc, <4 x i32> %samp, i1 0, i32 0, i32 0)
+  ret <4 x float> %v
+}
+
+define amdgpu_ps <4 x float> @sample_2d_nsa3(<8 x i32> inreg %rsrc, <4 x i32> inreg %samp, float %t, float %s) #3 {
+; ATTRIB-LABEL: sample_2d_nsa3:
+; ATTRIB:       ; %bb.0: ; %main_body
+; ATTRIB-NEXT:    s_mov_b32 s12, exec_lo
+; ATTRIB-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; ATTRIB-NEXT:    v_mov_b32_e32 v2, v0
+; ATTRIB-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; ATTRIB-NEXT:    image_sample v[0:3], v[1:2], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; ATTRIB-NEXT:    s_waitcnt vmcnt(0)
+; ATTRIB-NEXT:    ; return to shader part epilog
+;
+; FORCE-2-LABEL: sample_2d_nsa3:
+; FORCE-2:       ; %bb.0: ; %main_body
+; FORCE-2-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-2-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-2-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; FORCE-2-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-2-NEXT:    image_sample v[0:3], [v1, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; FORCE-2-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-2-NEXT:    ; return to shader part epilog
+;
+; FORCE-3-LABEL: sample_2d_nsa3:
+; FORCE-3:       ; %bb.0: ; %main_body
+; FORCE-3-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-3-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-3-NEXT:    v_mov_b32_e32 v2, v0
+; FORCE-3-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-3-NEXT:    image_sample v[0:3], v[1:2], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; FORCE-3-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-3-NEXT:    ; return to shader part epilog
+;
+; FORCE-4-LABEL: sample_2d_nsa3:
+; FORCE-4:       ; %bb.0: ; %main_body
+; FORCE-4-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-4-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-4-NEXT:    v_mov_b32_e32 v2, v0
+; FORCE-4-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-4-NEXT:    image_sample v[0:3], v[1:2], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; FORCE-4-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-4-NEXT:    ; return to shader part epilog
+main_body:
+  %v = call <4 x float> @llvm.amdgcn.image.sample.2d.v4f32.f32(i32 15, float %s, float %t, <8 x i32> %rsrc, <4 x i32> %samp, i1 0, i32 0, i32 0)
+  ret <4 x float> %v
+}
+
+define amdgpu_ps <4 x float> @sample_3d_nsa3(<8 x i32> inreg %rsrc, <4 x i32> inreg %samp, float %r, float %s, float %t) #3 {
+; ATTRIB-LABEL: sample_3d_nsa3:
+; ATTRIB:       ; %bb.0: ; %main_body
+; ATTRIB-NEXT:    s_mov_b32 s12, exec_lo
+; ATTRIB-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; ATTRIB-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; ATTRIB-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; ATTRIB-NEXT:    image_sample v[0:3], [v1, v2, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; ATTRIB-NEXT:    s_waitcnt vmcnt(0)
+; ATTRIB-NEXT:    ; return to shader part epilog
+;
+; FORCE-2-LABEL: sample_3d_nsa3:
+; FORCE-2:       ; %bb.0: ; %main_body
+; FORCE-2-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-2-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-2-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; FORCE-2-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-2-NEXT:    image_sample v[0:3], [v1, v2, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; FORCE-2-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-2-NEXT:    ; return to shader part epilog
+;
+; FORCE-3-LABEL: sample_3d_nsa3:
+; FORCE-3:       ; %bb.0: ; %main_body
+; FORCE-3-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-3-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-3-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; FORCE-3-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-3-NEXT:    image_sample v[0:3], [v1, v2, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; FORCE-3-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-3-NEXT:    ; return to shader part epilog
+;
+; FORCE-4-LABEL: sample_3d_nsa3:
+; FORCE-4:       ; %bb.0: ; %main_body
+; FORCE-4-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-4-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-4-NEXT:    v_mov_b32_e32 v3, v0
+; FORCE-4-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-4-NEXT:    image_sample v[0:3], v[1:3], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; FORCE-4-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-4-NEXT:    ; return to shader part epilog
+main_body:
+  %v = call <4 x float> @llvm.amdgcn.image.sample.3d.v4f32.f32(i32 15, float %s, float %t, float %r, <8 x i32> %rsrc, <4 x i32> %samp, i1 0, i32 0, i32 0)
+  ret <4 x float> %v
+}
+
+define amdgpu_ps <4 x float> @sample_2d_nsa4(<8 x i32> inreg %rsrc, <4 x i32> inreg %samp, float %t, float %s) #4 {
+; ATTRIB-LABEL: sample_2d_nsa4:
+; ATTRIB:       ; %bb.0: ; %main_body
+; ATTRIB-NEXT:    s_mov_b32 s12, exec_lo
+; ATTRIB-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; ATTRIB-NEXT:    v_mov_b32_e32 v2, v0
+; ATTRIB-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; ATTRIB-NEXT:    image_sample v[0:3], v[1:2], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; ATTRIB-NEXT:    s_waitcnt vmcnt(0)
+; ATTRIB-NEXT:    ; return to shader part epilog
+;
+; FORCE-2-LABEL: sample_2d_nsa4:
+; FORCE-2:       ; %bb.0: ; %main_body
+; FORCE-2-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-2-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-2-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; FORCE-2-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-2-NEXT:    image_sample v[0:3], [v1, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; FORCE-2-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-2-NEXT:    ; return to shader part epilog
+;
+; FORCE-3-LABEL: sample_2d_nsa4:
+; FORCE-3:       ; %bb.0: ; %main_body
+; FORCE-3-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-3-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-3-NEXT:    v_mov_b32_e32 v2, v0
+; FORCE-3-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-3-NEXT:    image_sample v[0:3], v[1:2], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; FORCE-3-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-3-NEXT:    ; return to shader part epilog
+;
+; FORCE-4-LABEL: sample_2d_nsa4:
+; FORCE-4:       ; %bb.0: ; %main_body
+; FORCE-4-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-4-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-4-NEXT:    v_mov_b32_e32 v2, v0
+; FORCE-4-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-4-NEXT:    image_sample v[0:3], v[1:2], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+; FORCE-4-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-4-NEXT:    ; return to shader part epilog
+main_body:
+  %v = call <4 x float> @llvm.amdgcn.image.sample.2d.v4f32.f32(i32 15, float %s, float %t, <8 x i32> %rsrc, <4 x i32> %samp, i1 0, i32 0, i32 0)
+  ret <4 x float> %v
+}
+
+define amdgpu_ps <4 x float> @sample_3d_nsa4(<8 x i32> inreg %rsrc, <4 x i32> inreg %samp, float %r, float %s, float %t) #4 {
+; ATTRIB-LABEL: sample_3d_nsa4:
+; ATTRIB:       ; %bb.0: ; %main_body
+; ATTRIB-NEXT:    s_mov_b32 s12, exec_lo
+; ATTRIB-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; ATTRIB-NEXT:    v_mov_b32_e32 v3, v0
+; ATTRIB-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; ATTRIB-NEXT:    image_sample v[0:3], v[1:3], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; ATTRIB-NEXT:    s_waitcnt vmcnt(0)
+; ATTRIB-NEXT:    ; return to shader part epilog
+;
+; FORCE-2-LABEL: sample_3d_nsa4:
+; FORCE-2:       ; %bb.0: ; %main_body
+; FORCE-2-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-2-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-2-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; FORCE-2-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-2-NEXT:    image_sample v[0:3], [v1, v2, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; FORCE-2-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-2-NEXT:    ; return to shader part epilog
+;
+; FORCE-3-LABEL: sample_3d_nsa4:
+; FORCE-3:       ; %bb.0: ; %main_body
+; FORCE-3-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-3-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-3-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
+; FORCE-3-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-3-NEXT:    image_sample v[0:3], [v1, v2, v0], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; FORCE-3-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-3-NEXT:    ; return to shader part epilog
+;
+; FORCE-4-LABEL: sample_3d_nsa4:
+; FORCE-4:       ; %bb.0: ; %main_body
+; FORCE-4-NEXT:    s_mov_b32 s12, exec_lo
+; FORCE-4-NEXT:    s_wqm_b32 exec_lo, exec_lo
+; FORCE-4-NEXT:    v_mov_b32_e32 v3, v0
+; FORCE-4-NEXT:    s_and_b32 exec_lo, exec_lo, s12
+; FORCE-4-NEXT:    image_sample v[0:3], v[1:3], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_3D
+; FORCE-4-NEXT:    s_waitcnt vmcnt(0)
+; FORCE-4-NEXT:    ; return to shader part epilog
+main_body:
+  %v = call <4 x float> @llvm.amdgcn.image.sample.3d.v4f32.f32(i32 15, float %s, float %t, float %r, <8 x i32> %rsrc, <4 x i32> %samp, i1 0, i32 0, i32 0)
+  ret <4 x float> %v
+}
+
+declare <4 x float> @llvm.amdgcn.image.sample.2d.v4f32.f32(i32, float, float, <8 x i32>, <4 x i32>, i1, i32, i32) #1
+declare <4 x float> @llvm.amdgcn.image.sample.3d.v4f32.f32(i32, float, float, float, <8 x i32>, <4 x i32>, i1, i32, i32) #1
+
+attributes #1 = { nounwind readonly }
+attributes #2 = { nounwind readonly "amdgpu-nsa-threshold"="2" }
+attributes #3 = { nounwind readonly "amdgpu-nsa-threshold"="3" }
+attributes #4 = { nounwind readonly "amdgpu-nsa-threshold"="4" }

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.image.nsa.ll

 ; RUN: llc -march=amdgcn -mcpu=gfx1010 -mattr=-nsa-encoding -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NONSA,GFX10-NONSA %s
-; RUN: llc -march=amdgcn -mcpu=gfx1010 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NSA,GFX1010-NSA %s
-; RUN: llc -march=amdgcn -mcpu=gfx1030 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NSA,GFX1030-NSA %s
+; RUN: llc -march=amdgcn -mcpu=gfx1010 -amdgpu-nsa-threshold=32 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NONSA,GFX10-NONSA %s
+; RUN: llc -march=amdgcn -mcpu=gfx1010 -amdgpu-nsa-threshold=2 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NSA,NSA-T2 %s
+; RUN: llc -march=amdgcn -mcpu=gfx1010 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NSA,NSA-T3,GFX1010-NSA %s
+; RUN: llc -march=amdgcn -mcpu=gfx1030 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NSA,NSA-T3,GFX1030-NSA %s
 ; RUN: llc -march=amdgcn -mcpu=gfx1100 -mattr=-nsa-encoding -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NONSA,GFX11-NONSA %s
-; RUN: llc -march=amdgcn -mcpu=gfx1100 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NSA,GFX11-NSA %s
+; RUN: llc -march=amdgcn -mcpu=gfx1100 -amdgpu-nsa-threshold=32 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NONSA,GFX11-NONSA %s
+; RUN: llc -march=amdgcn -mcpu=gfx1100 -amdgpu-nsa-threshold=2 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NSA,NSA-T2 %s
+; RUN: llc -march=amdgcn -mcpu=gfx1100 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,NSA,NSA-T3,GFX11-NSA %s
 
+; Default NSA threshold is 3 addresses
 ; GCN-LABEL: {{^}}sample_2d:
-;
-; TODO: use NSA here
-; GCN: v_mov_b32_e32 v2, v0
-;
-; GCN: image_sample v[0:3], v[1:2],
+; NONSA: v_mov_b32_e32 v2, v0
+; NONSA: image_sample v[0:3], v[1:2],
+; NSA-T2: image_sample v[0:3], [v1, v0],
+; NSA-T3: v_mov_b32_e32 v2, v0
+; NSA-T3: image_sample v[0:3], v[1:2],
 define amdgpu_ps <4 x float> @sample_2d(<8 x i32> inreg %rsrc, <4 x i32> inreg %samp, float %t, float %s) {
 main_body:
   %v = call <4 x float> @llvm.amdgcn.image.sample.2d.v4f32.f32(i32 15, float %s, float %t, <8 x i32> %rsrc, <4 x i32> %samp, i1 0, i32 0, i32 0)
   ret <4 x float> %v
 }
 
 ; GCN-LABEL: {{^}}sample_3d:
 ; NONSA: v_mov_b32_e32 v3, v0