[llvm][nvptx] Add sm_90a

[guraypp]

Summary

This works adds sm_90a as nvptx target. sm_90a is required to generate wgmma and setmaxnreg instructions.

Here is information about "a" prefix in PTX document:
Target architectures with suffix “a”, such as sm_90a, include architecture-accelerated features that are supported on the specified architecture only, hence such targets do not follow the onion layer model. Therefore, PTX code generated for such targets cannot be run on later generation devices. Architecture-accelerated features can only be used with targets that support these features.

[guraypp]

I am using NVPTX backend from MLIR. Adding sm_90a like this works fine to generate wgmma instructions in my case.

Please let me if I am missing something here.

[tra]

It's unfortunate that NVIDIA decided to break existing naming convention. :-/

PTX code generated for such targets cannot be run on later generation devices. Architecture-accelerated features can only be used with targets that support these features.

We do have a number of places where we compare SM values, so adding support for sm90a will require changing the way we encode SM version and the way we determine availability of various instructions/intrinsics/builtins. Properly implementing it will be a bit of a pain.

This patch effectively makes sm_90a a synonym for sm_90, only with additional constraint of requiring PTX 8.0.

It does not hurt, but it also does not do anything you would not be able to do by specifying sm_90 and ptx80 manually. Do you really need this change? Or do you plan making further changes to fully implement support for sm_90a?

llvm/lib/Target/NVPTX/NVPTX.td

46

Nit: I'd move it up to where we define other SMxx records.

[guraypp]

Thanks for review! Having a synonym sm_90a is actually a big help for generating .target sm_90a in PTX code. ptxas throws an error for wgmma instructions with .target sm_90.

By the way, wgmma are new tensor core instructions for hopper, and they're only supported for sm_90a. MLIR generates them as inline assembly, and one can generate them using asm in CUDA (like CUTLASS).

It would be really great if we can land this workaround until we find a proper solution.

I saw that sm_90a requires PTX 8.0, see below:
https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#release-notes-ptx-release-history

[tra]

SGTM.

[guraypp]

rebase

[guraypp]

move sm90a definition to up

[guraypp]

rebase

[guraypp]

rebase

llvm/lib/Target/NVPTX/NVPTX.td

 // - Need at least one feature to avoid generating zero sized array by
 //   TableGen in NVPTXGenSubtarget.inc.
 //===----------------------------------------------------------------------===//
 
 class FeatureSM<int version>:
    SubtargetFeature<"sm_"# version, "SmVersion",
                     "" # version,
                     "Target SM " # version>;
+def SM90a: FeatureSM<90>;
+
 class FeaturePTX<int version>:
    SubtargetFeature<"ptx"# version, "PTXVersion",
                     "" # version,
                     "Use PTX version " # version>;
 
 foreach version = [20, 21, 30, 32, 35, 37, 50, 52, 53,
                    60, 61, 62, 70, 72, 75, 80, 86, 87, 89, 90] in
   def SM#version: FeatureSM<version>;
 
 foreach version = [32, 40, 41, 42, 43, 50, 60, 61, 63, 64, 65,
                    70, 71, 72, 73, 74, 75, 76, 77, 78, 80, 81] in
   def PTX#version: FeaturePTX<version>;
 
 //===----------------------------------------------------------------------===//
 // NVPTX supported processors.
 //===----------------------------------------------------------------------===//
 
 class Proc<string Name, list<SubtargetFeature> Features>
  : Processor<Name, NoItineraries, Features>;
 
 def : Proc<"sm_20", [SM20, PTX32]>;
 def : Proc<"sm_21", [SM21, PTX32]>;

 def : Proc<"sm_70", [SM70, PTX60]>;
 def : Proc<"sm_72", [SM72, PTX61]>;
 def : Proc<"sm_75", [SM75, PTX63]>;
 def : Proc<"sm_80", [SM80, PTX70]>;
 def : Proc<"sm_86", [SM86, PTX71]>;
 def : Proc<"sm_87", [SM87, PTX74]>;
 def : Proc<"sm_89", [SM89, PTX78]>;
 def : Proc<"sm_90", [SM90, PTX78]>;
+def : Proc<"sm_90a", [SM90a, PTX80]>;
 
 def NVPTXInstrInfo : InstrInfo {
 }
 
 def NVPTX : Target {
   let InstructionSet = NVPTXInstrInfo;
 }