use super::nvptx_instr_info::{NvptxInstrInfo, NvptxOpcode};
use super::nvptx_mc_encoder::{ComputeCapability, ComputeFeature, NvptxMCEncoder, PtxType};
use super::nvptx_register_info::NvptxRegisterInfo;
use crate::codegen::*;
use crate::module::Module;
use std::collections::HashMap;
pub mod address_space {
pub const GLOBAL: u32 = 1;
pub const CONSTANT: u32 = 4;
pub const SHARED: u32 = 3;
pub const LOCAL: u32 = 5;
pub const PARAM: u32 = 101;
pub const GENERIC: u32 = 0;
pub const TEXTURE: u32 = 2;
pub const SURFACE: u32 = 102;
}
pub fn addrspace_to_ptx_space(addrspace: u32) -> &'static str {
match addrspace {
address_space::GLOBAL => "global",
address_space::CONSTANT => "const",
address_space::SHARED => "shared",
address_space::LOCAL => "local",
address_space::PARAM => "param",
address_space::GENERIC => "generic",
address_space::TEXTURE => "tex",
address_space::SURFACE => "surf",
_ => "global",
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AddressSpaceProperties {
pub name: &'static str,
pub is_readable: bool,
pub is_writable: bool,
pub supports_atomics: bool,
pub is_grid_coherent: bool,
pub l1_cached: bool,
pub l2_cached: bool,
pub max_access_size: u32,
pub alignment: u32,
}
pub fn get_address_space_properties(addrspace: u32) -> AddressSpaceProperties {
match addrspace {
address_space::GLOBAL => AddressSpaceProperties {
name: "global",
is_readable: true,
is_writable: true,
supports_atomics: true,
is_grid_coherent: false,
l1_cached: true,
l2_cached: true,
max_access_size: 16,
alignment: 1,
},
address_space::CONSTANT => AddressSpaceProperties {
name: "const",
is_readable: true,
is_writable: false,
supports_atomics: false,
is_grid_coherent: true,
l1_cached: true,
l2_cached: true,
max_access_size: 16,
alignment: 1,
},
address_space::SHARED => AddressSpaceProperties {
name: "shared",
is_readable: true,
is_writable: true,
supports_atomics: true,
is_grid_coherent: false,
l1_cached: false,
l2_cached: false,
max_access_size: 16,
alignment: 4,
},
address_space::LOCAL => AddressSpaceProperties {
name: "local",
is_readable: true,
is_writable: true,
supports_atomics: false,
is_grid_coherent: false,
l1_cached: true,
l2_cached: false,
max_access_size: 16,
alignment: 1,
},
address_space::PARAM => AddressSpaceProperties {
name: "param",
is_readable: true,
is_writable: false,
supports_atomics: false,
is_grid_coherent: true,
l1_cached: true,
l2_cached: true,
max_access_size: 16,
alignment: 1,
},
_ => AddressSpaceProperties {
name: "global",
is_readable: true,
is_writable: true,
supports_atomics: true,
is_grid_coherent: false,
l1_cached: true,
l2_cached: true,
max_access_size: 16,
alignment: 1,
},
}
}
pub struct NvptxTargetMachine {
pub target_triple: String,
pub sm_version: u32,
pub ptx_version: String,
pub is_64bit: bool,
pub opt_level: String,
pub features: Vec<String>,
pub data_layout: String,
pub instr_info: NvptxInstrInfo,
pub register_info: NvptxRegisterInfo,
pub max_threads_per_block: u32,
pub max_registers_per_thread: u32,
pub shared_memory_per_block: u32,
}
impl NvptxTargetMachine {
pub fn new(sm_version: u32, is_64bit: bool) -> Self {
let cap = ComputeCapability::from_sm(sm_version);
NvptxTargetMachine {
target_triple: format!("nvptx{}-nvidia-cuda", if is_64bit { "64" } else { "" }),
sm_version,
ptx_version: cap.ptx_version().to_string(),
is_64bit,
opt_level: "O2".to_string(),
features: Vec::new(),
data_layout: if is_64bit {
"e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n32:64".to_string()
} else {
"e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n32".to_string()
},
instr_info: NvptxInstrInfo::new(),
register_info: NvptxRegisterInfo,
max_threads_per_block: 1024,
max_registers_per_thread: 255,
shared_memory_per_block: cap.max_shared_memory(),
}
}
pub fn set_opt_level(&mut self, level: &str) {
self.opt_level = level.to_string();
}
pub fn enable_feature(&mut self, feature: &str) {
if !self.features.contains(&feature.to_string()) {
self.features.push(feature.to_string());
}
}
pub fn has_feature(&self, feature: &str) -> bool {
self.features.contains(&feature.to_string())
}
pub fn compute_capability(&self) -> ComputeCapability {
ComputeCapability::from_sm(self.sm_version)
}
pub fn supports(&self, feature: ComputeFeature) -> bool {
self.compute_capability().supports(feature)
}
pub fn get_data_layout(&self) -> &str {
&self.data_layout
}
pub fn get_target_triple(&self) -> &str {
&self.target_triple
}
pub fn is_64bit(&self) -> bool {
self.is_64bit
}
pub fn get_ptx_version(&self) -> &str {
&self.ptx_version
}
pub fn add_ir_passes(&self) -> Vec<&str> {
match self.opt_level.as_str() {
"O0" => vec!["mem2reg"],
"O1" => vec!["mem2reg", "instcombine", "simplifycfg", "dce"],
"O2" | "O3" => vec![
"mem2reg",
"instcombine",
"simplifycfg",
"gvn",
"dce",
"inline",
"licm",
"sccp",
"loop-unroll",
"slp-vectorize",
],
_ => vec!["mem2reg", "instcombine", "simplifycfg"],
}
}
pub fn create_mc_encoder(&self) -> NvptxMCEncoder {
NvptxMCEncoder::new(self.sm_version, self.is_64bit)
}
pub fn emit_module(&self, mf: &MachineFunction) -> String {
let mut encoder = self.create_mc_encoder();
encoder.encode_module(mf)
}
}
impl Default for NvptxTargetMachine {
fn default() -> Self {
Self::new(70, true)
}
}
#[derive(Debug, Clone)]
pub struct KernelArg {
pub name: String,
pub ty: PtxType,
pub addrspace: u32,
pub size: u32,
pub alignment: u32,
}
impl KernelArg {
pub fn new(name: &str, ty: PtxType) -> Self {
KernelArg {
name: name.to_string(),
ty,
addrspace: address_space::PARAM,
size: ty.size_bytes(),
alignment: 4,
}
}
pub fn emit_ptx(&self) -> String {
format!(
"\t.param {}{} {}",
self.ty.to_ptx(),
if self.alignment > 1 {
format!(".align {}", self.alignment)
} else {
String::new()
},
self.name
)
}
}
#[derive(Debug, Clone)]
pub struct KernelMetadata {
pub name: String,
pub params: Vec<KernelArg>,
pub return_type: Option<KernelArg>,
pub max_registers: u32,
pub dynamic_shared_memory: u32,
pub min_threads_per_block: u32,
pub max_threads_per_block: u32,
pub min_sm_version: u32,
}
impl KernelMetadata {
pub fn new(name: &str) -> Self {
KernelMetadata {
name: name.to_string(),
params: Vec::new(),
return_type: None,
max_registers: 64,
dynamic_shared_memory: 0,
min_threads_per_block: 1,
max_threads_per_block: 1024,
min_sm_version: 50,
}
}
pub fn add_param(&mut self, arg: KernelArg) {
self.params.push(arg);
}
pub fn set_return_type(&mut self, arg: KernelArg) {
self.return_type = Some(arg);
}
pub fn set_register_usage(&mut self, max_regs: u32) {
self.max_registers = max_regs;
}
pub fn set_shared_memory(&mut self, bytes: u32) {
self.dynamic_shared_memory = bytes;
}
pub fn emit_ptx(&self) -> String {
let mut out = String::new();
out.push_str(&format!("\t.entry {}(", self.name));
if !self.params.is_empty() {
out.push('\n');
for (i, p) in self.params.iter().enumerate() {
out.push_str(&p.emit_ptx());
if i < self.params.len() - 1 || self.return_type.is_some() {
out.push_str(",\n");
}
}
if let Some(ref ret) = self.return_type {
out.push_str(&ret.emit_ptx());
}
out.push('\n');
} else if let Some(ref ret) = self.return_type {
out.push_str(&ret.emit_ptx());
out.push('\n');
}
out.push(')');
if self.max_threads_per_block > 0 {
out.push('\n');
out.push_str(&format!(
"\t.maxntid\t{}, {}, {}",
self.max_threads_per_block, 1, 1
));
}
if self.max_registers > 0 {
out.push('\n');
out.push_str(&format!("\t.maxnreg\t{}", self.max_registers));
}
if self.dynamic_shared_memory > 0 {
out.push('\n');
out.push_str(&format!("\t.reqntid\t{}", self.dynamic_shared_memory));
}
if self.min_sm_version >= 20 {
out.push('\n');
out.push_str(&format!("\t.minnctapersm\t{}", self.min_sm_version * 10));
}
out.push('\n');
out
}
}
pub mod nvvm_metadata {
pub const ANNOTATIONS: &str = "nvvm.annotations";
pub const MAXNTID_X: &str = "maxntid_x";
pub const MAXNTID_Y: &str = "maxntid_y";
pub const MAXNTID_Z: &str = "maxntid_z";
pub const MAXNREG: &str = "maxnreg";
pub const REQNTID: &str = "reqntid";
pub const MINNCTAPERSM: &str = "minnctapersm";
pub const KERNEL_NAME: &str = "kernel";
pub const GLOBAL_VARIABLE: &str = "global";
}
pub fn emit_kernel(meta: &KernelMetadata, body: &str) -> String {
let mut out = meta.emit_ptx();
out.push_str("{\n");
out.push_str(body);
out.push_str("}\n");
out
}
pub fn kernel_metadata_from_function(
name: &str,
param_types: &[(String, PtxType)],
ret_type: Option<PtxType>,
attrs: &HashMap<String, String>,
) -> KernelMetadata {
let mut meta = KernelMetadata::new(name);
for (pname, pty) in param_types {
meta.add_param(KernelArg::new(pname, *pty));
}
if let Some(rty) = ret_type {
meta.set_return_type(KernelArg::new("retval", rty));
}
if let Some(val) = attrs.get("maxnreg") {
if let Ok(n) = val.parse() {
meta.max_registers = n;
}
}
if let Some(val) = attrs.get("maxntid") {
if let Ok(n) = val.parse() {
meta.max_threads_per_block = n;
}
}
if let Some(val) = attrs.get("reqntid") {
if let Ok(n) = val.parse() {
meta.dynamic_shared_memory = n;
}
}
meta
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct PtxIsaVersion {
pub major: u32,
pub minor: u32,
}
impl PtxIsaVersion {
pub const V4_0: PtxIsaVersion = PtxIsaVersion { major: 4, minor: 0 };
pub const V4_1: PtxIsaVersion = PtxIsaVersion { major: 4, minor: 1 };
pub const V4_2: PtxIsaVersion = PtxIsaVersion { major: 4, minor: 2 };
pub const V4_3: PtxIsaVersion = PtxIsaVersion { major: 4, minor: 3 };
pub const V5_0: PtxIsaVersion = PtxIsaVersion { major: 5, minor: 0 };
pub const V6_0: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 0 };
pub const V6_1: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 1 };
pub const V6_2: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 2 };
pub const V6_3: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 3 };
pub const V6_4: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 4 };
pub const V6_5: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 5 };
pub const V7_0: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 0 };
pub const V7_1: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 1 };
pub const V7_2: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 2 };
pub const V7_3: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 3 };
pub const V7_4: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 4 };
pub const V7_5: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 5 };
pub const V7_6: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 6 };
pub const V7_7: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 7 };
pub const V7_8: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 8 };
pub const V8_0: PtxIsaVersion = PtxIsaVersion { major: 8, minor: 0 };
pub const V8_1: PtxIsaVersion = PtxIsaVersion { major: 8, minor: 1 };
pub const V8_2: PtxIsaVersion = PtxIsaVersion { major: 8, minor: 2 };
pub fn from_sm(sm: u32) -> Self {
match sm {
10..=13 => PtxIsaVersion::V4_0,
20..=21 => PtxIsaVersion::V4_1,
30..=35 => PtxIsaVersion::V4_2,
37 => PtxIsaVersion::V4_3,
50..=53 => PtxIsaVersion::V5_0,
60..=62 => PtxIsaVersion::V6_0,
70..=72 => PtxIsaVersion::V7_0,
75 => PtxIsaVersion::V7_5,
80 => PtxIsaVersion::V7_8,
86 => PtxIsaVersion::V8_0,
89 => PtxIsaVersion::V8_1,
90..=99 => PtxIsaVersion::V8_2,
_ => PtxIsaVersion::V5_0,
}
}
pub fn to_version_string(&self) -> String {
format!("{}.{}", self.major, self.minor)
}
pub fn supports_wmma(&self) -> bool {
*self >= PtxIsaVersion::V6_0
}
pub fn supports_mma(&self) -> bool {
*self >= PtxIsaVersion::V7_0
}
pub fn supports_cp_async(&self) -> bool {
*self >= PtxIsaVersion::V7_8
}
pub fn supports_stmatrix(&self) -> bool {
*self >= PtxIsaVersion::V7_8
}
pub fn supports_dpx(&self) -> bool {
*self >= PtxIsaVersion::V8_0
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_address_space_to_ptx() {
assert_eq!(addrspace_to_ptx_space(address_space::GLOBAL), "global");
assert_eq!(addrspace_to_ptx_space(address_space::SHARED), "shared");
assert_eq!(addrspace_to_ptx_space(address_space::CONSTANT), "const");
assert_eq!(addrspace_to_ptx_space(address_space::LOCAL), "local");
assert_eq!(addrspace_to_ptx_space(address_space::PARAM), "param");
assert_eq!(addrspace_to_ptx_space(999), "global");
}
#[test]
fn test_address_space_properties() {
let props = get_address_space_properties(address_space::GLOBAL);
assert!(props.is_readable);
assert!(props.is_writable);
assert!(props.supports_atomics);
let shared = get_address_space_properties(address_space::SHARED);
assert!(shared.is_readable);
assert!(!shared.l2_cached);
let cnst = get_address_space_properties(address_space::CONSTANT);
assert!(cnst.is_readable);
assert!(!cnst.is_writable);
}
#[test]
fn test_target_machine_new() {
let tm = NvptxTargetMachine::new(70, true);
assert!(tm.is_64bit());
assert_eq!(tm.sm_version, 70);
assert!(tm.get_target_triple().contains("nvptx64"));
}
#[test]
fn test_target_machine_32bit() {
let tm = NvptxTargetMachine::new(50, false);
assert!(!tm.is_64bit());
assert!(!tm.get_target_triple().contains("64"));
}
#[test]
fn test_target_machine_features() {
let mut tm = NvptxTargetMachine::new(80, true);
tm.enable_feature("sm_80");
assert!(tm.has_feature("sm_80"));
assert!(!tm.has_feature("sm_70"));
}
#[test]
fn test_target_machine_opt_level() {
let mut tm = NvptxTargetMachine::new(60, true);
tm.set_opt_level("O3");
assert!(!tm.add_ir_passes().is_empty());
assert!(tm.add_ir_passes().len() > 6);
}
#[test]
fn test_compute_capability_support() {
let tm = NvptxTargetMachine::new(80, true);
assert!(tm.supports(ComputeFeature::TensorCores));
assert!(tm.supports(ComputeFeature::FP16));
assert!(tm.supports(ComputeFeature::BF16));
let tm2 = NvptxTargetMachine::new(50, true);
assert!(!tm2.supports(ComputeFeature::TensorCores));
assert!(!tm2.supports(ComputeFeature::BF16));
}
#[test]
fn test_default_target_machine() {
let tm = NvptxTargetMachine::default();
assert_eq!(tm.sm_version, 70);
assert!(tm.is_64bit());
}
#[test]
fn test_kernel_metadata_empty() {
let meta = KernelMetadata::new("my_kernel");
let ptx = meta.emit_ptx();
assert!(ptx.contains("my_kernel"));
assert!(ptx.contains(".entry"));
}
#[test]
fn test_kernel_metadata_with_params() {
let mut meta = KernelMetadata::new("add_kernel");
meta.add_param(KernelArg::new("a", PtxType::F32));
meta.add_param(KernelArg::new("b", PtxType::F32));
let ptx = meta.emit_ptx();
assert!(ptx.contains(".param .f32 a"));
assert!(ptx.contains(".param .f32 b"));
}
#[test]
fn test_kernel_metadata_with_return() {
let mut meta = KernelMetadata::new("compute");
meta.set_return_type(KernelArg::new("retval", PtxType::U64));
let ptx = meta.emit_ptx();
assert!(ptx.contains(".param .u64 retval"));
}
#[test]
fn test_emit_kernel_full() {
let mut meta = KernelMetadata::new("testkern");
meta.add_param(KernelArg::new("in", PtxType::S32));
meta.set_register_usage(32);
let body = "\tld.param.u32\t%r1, [in];\n\tret;\n";
let output = emit_kernel(&meta, body);
assert!(output.contains("testkern"));
assert!(output.contains("ld.param.u32"));
assert!(output.contains(".maxnreg\t32"));
}
#[test]
fn test_ptx_isa_version_from_sm() {
assert_eq!(PtxIsaVersion::from_sm(50), PtxIsaVersion::V5_0);
assert_eq!(PtxIsaVersion::from_sm(70), PtxIsaVersion::V7_0);
assert_eq!(PtxIsaVersion::from_sm(75), PtxIsaVersion::V7_5);
assert_eq!(PtxIsaVersion::from_sm(80), PtxIsaVersion::V7_8);
assert_eq!(PtxIsaVersion::from_sm(90), PtxIsaVersion::V8_2);
}
#[test]
fn test_ptx_isa_version_features() {
let v5 = PtxIsaVersion::V5_0;
assert!(!v5.supports_wmma());
assert!(!v5.supports_mma());
let v7 = PtxIsaVersion::V7_0;
assert!(v7.supports_wmma());
assert!(v7.supports_mma());
let v8 = PtxIsaVersion::V8_0;
assert!(v8.supports_cp_async());
assert!(v8.supports_dpx());
}
#[test]
fn test_kernel_metadata_from_function() {
let mut attrs = HashMap::new();
attrs.insert("maxnreg".to_string(), "64".to_string());
attrs.insert("maxntid".to_string(), "256".to_string());
let params = vec![
("x".to_string(), PtxType::F32),
("y".to_string(), PtxType::F32),
];
let meta = kernel_metadata_from_function("vecadd", ¶ms, None, &attrs);
assert_eq!(meta.name, "vecadd");
assert_eq!(meta.params.len(), 2);
assert_eq!(meta.max_registers, 64);
assert_eq!(meta.max_threads_per_block, 256);
}
#[test]
fn test_kernel_arg_emit_ptx() {
let arg = KernelArg::new("my_param", PtxType::U64);
let ptx = arg.emit_ptx();
assert!(ptx.contains(".param .u64"));
assert!(ptx.contains("my_param"));
}
#[test]
fn test_nvvm_metadata_constants() {
assert_eq!(nvvm_metadata::ANNOTATIONS, "nvvm.annotations");
assert_eq!(nvvm_metadata::KERNEL_NAME, "kernel");
}
#[test]
fn test_create_mc_encoder_from_tm() {
let tm = NvptxTargetMachine::new(75, true);
let enc = tm.create_mc_encoder();
assert_eq!(enc.sm_version, 75);
assert!(enc.is_64bit);
}
#[test]
fn test_emit_module_basic() {
let tm = NvptxTargetMachine::new(70, true);
let mut mf = MachineFunction::new("test_func");
let block = MachineBlock {
name: "entry".to_string(),
instructions: Vec::new(),
};
mf.blocks.push(block);
let output = tm.emit_module(&mf);
assert!(output.contains(".version"));
assert!(output.contains(".target"));
assert!(output.contains("test_func"));
}
}