Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 11 additions & 8 deletions cuda-oxide/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

18 changes: 10 additions & 8 deletions cuda-oxide/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -45,15 +45,16 @@ repository = "https://github.com/NVIDIA/cuda-rust"
# Examples are in: crates/rustc-codegen-cuda/examples/

[workspace.dependencies]
# Pliron MLIR framework, pinned to upstream struct-layout, GEP no-wrap, volatile,
# global-constant, inline-asm, and generic LLVM attribute semantics.
# Pliron MLIR framework, pinned to the published 0.19.0 release, retaining
# struct-layout, GEP no-wrap, volatile, global-constant, inline-asm, and generic
# LLVM attribute semantics.
# pliron-llvm uses default-features = false to keep llvm-sys (LLVM 23) out:
# cuda-oxide consumes the dialect modeling only, not the FFI bridge.
# `crates/rustc-codegen-cuda/Cargo.toml` MUST stay on this same rev, or pliron
# resolves to two distinct crates and types stop unifying.
pliron = { git = "https://github.com/pliron-org/pliron", rev = "06d99a3df938394583715aacf5eadb65b4791c09" }
pliron-derive = { git = "https://github.com/pliron-org/pliron", rev = "06d99a3df938394583715aacf5eadb65b4791c09" }
pliron-llvm = { git = "https://github.com/pliron-org/pliron", rev = "06d99a3df938394583715aacf5eadb65b4791c09", default-features = false, features = [
# `crates/rustc-codegen-cuda/Cargo.toml` MUST use the same registry source and
# exact release, or pliron resolves to distinct crates and types stop unifying.
pliron = { version = "=0.19.0" }
pliron-derive = { version = "=0.19.0" }
pliron-llvm = { version = "=0.19.0", default-features = false, features = [
"std",
] }

Expand Down Expand Up @@ -108,7 +109,8 @@ thiserror = "2"
linkme = "0.3"
anyhow = "1.0"
once_cell = "1.18"
combine = "4.6"
# Pliron 0.19 requires SourcePosition: Hash, introduced in combine 4.6.8.
combine = "4.6.8"
log = "0.4"
half = "2.7"
libloading = "0.8"
Expand Down
154 changes: 80 additions & 74 deletions cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,14 @@
//!
//! This module defines assertions, terminators, and branches for the MIR dialect.

use std::ops::Range;

use pliron::{
builtin::{
attributes::OperandSegmentSizesAttr,
op_interfaces::{
BranchOpInterface, IsTerminatorInterface, NResultsInterface, OperandSegmentInterface,
ATTR_KEY_OPERAND_SEGMENT_SIZES, BranchOpInterface, IsTerminatorInterface,
NResultsInterface, OperandSegmentInterface,
},
type_interfaces::FunctionTypeInterface,
types::IntegerType,
Expand Down Expand Up @@ -146,41 +150,23 @@ impl MirGotoOp {

impl Verify for MirGotoOp {
fn verify(&self, ctx: &Context) -> Result<(), Error> {
let op = &*self.get_operation().deref(ctx);
<Self as BranchOpInterface>::verify(self, ctx)
}
}

#[op_interface_impl]
impl BranchOpInterface for MirGotoOp {
fn verify_successor_operand_layout(&self, ctx: &Context) -> Result<(), Error> {
let op = self.get_operation().deref(ctx);
if op.get_num_successors() != 1 {
return verify_err!(op.loc(), "MirGotoOp must have exactly 1 successor");
}
let succ = op.get_successor(0);
let succ_block = succ.deref(ctx);

if op.get_num_operands() != succ_block.get_num_arguments() {
return verify_err!(
op.loc(),
"MirGotoOp operand count must match successor argument count"
);
}

for i in 0..succ_block.get_num_arguments() {
let arg = succ_block.get_argument(i);
let opd = op.get_operand(i);
if opd.get_type(ctx) != arg.get_type(ctx) {
return verify_err!(
op.loc(),
"MirGotoOp operand type mismatch with successor argument"
);
}
}

Ok(())
}
}

#[op_interface_impl]
impl BranchOpInterface for MirGotoOp {
fn successor_operands(&self, ctx: &Context, succ_idx: usize) -> Vec<Value> {
fn successor_operand_range(&self, ctx: &Context, succ_idx: usize) -> Range<usize> {
assert!(succ_idx == 0, "MirGotoOp has exactly one successor");
self.get_operation().deref(ctx).operands().collect()
0..self.get_operation().deref(ctx).get_num_operands()
}

fn add_successor_operand(&self, ctx: &mut Context, succ_idx: usize, operand: Value) -> usize {
Expand Down Expand Up @@ -218,13 +204,14 @@ impl BranchOpInterface for MirGotoOp {
///
/// # Verification
///
/// - Exactly three operand segments, with one condition operand.
/// - Condition must be `i1`.
/// - Must have exactly 2 successors.
/// - Operand types must match successor block argument types.
#[pliron_op(
name = "mir.cond_br",
format,
interfaces = [NResultsInterface<0>, IsTerminatorInterface, OperandSegmentInterface]
interfaces = [NResultsInterface<0>, IsTerminatorInterface]
)]
pub struct MirCondBranchOp;

Expand All @@ -237,20 +224,10 @@ impl MirCondBranchOp {

impl Verify for MirCondBranchOp {
fn verify(&self, ctx: &Context) -> Result<(), Error> {
let op = &*self.get_operation().deref(ctx);

if op.get_num_successors() != 2 {
return verify_err!(op.loc(), "MirCondBranchOp must have exactly 2 successors");
}
let true_block = op.get_successor(0).deref(ctx);
let false_block = op.get_successor(1).deref(ctx);

if op.get_num_operands() < 1 {
return verify_err!(
op.loc(),
"MirCondBranchOp must have at least 1 operand (condition)"
);
}
// Validate segment metadata before reading operands. The shared branch
// verifier also checks each segment against its target block arguments.
<Self as BranchOpInterface>::verify(self, ctx)?;
let op = self.get_operation().deref(ctx);
let cond = op.get_operand(0);
let cond_ty = cond.get_type(ctx);
let cond_ty_obj = cond_ty.deref(ctx);
Expand All @@ -263,55 +240,83 @@ impl Verify for MirCondBranchOp {
return verify_err!(op.loc(), "MirCondBranchOp condition must be integer type");
}

let expected_args = 1 + true_block.get_num_arguments() + false_block.get_num_arguments();
if op.get_num_operands() != expected_args {
Ok(())
}
}

#[op_interface_impl]
impl OperandSegmentInterface for MirCondBranchOp {
fn expected_num_segments(&self, _ctx: &Context) -> Option<usize> {
// Condition, true successor arguments, false successor arguments.
Some(3)
}

fn verify(op: &dyn Op, ctx: &Context) -> Result<(), Error> {
let op = op.get_operation().deref(ctx);
let Some(sizes) = op
.attributes
.get::<OperandSegmentSizesAttr>(&ATTR_KEY_OPERAND_SEGMENT_SIZES)
else {
return verify_err!(op.loc(), "MirCondBranchOp requires operand segment sizes");
};
if sizes.0.len() != 3 {
return verify_err!(
op.loc(),
"MirCondBranchOp operand count must match condition + successors arguments"
"MirCondBranchOp must have exactly 3 operand segments"
);
}

let mut op_idx = 1;
for i in 0..true_block.get_num_arguments() {
let arg = true_block.get_argument(i);
let opd = op.get_operand(op_idx);
if opd.get_type(ctx) != arg.get_type(ctx) {
return verify_err!(
op.loc(),
"MirCondBranchOp true successor argument type mismatch"
);
}
op_idx += 1;
}

for i in 0..false_block.get_num_arguments() {
let arg = false_block.get_argument(i);
let opd = op.get_operand(op_idx);
if opd.get_type(ctx) != arg.get_type(ctx) {
return verify_err!(
op.loc(),
"MirCondBranchOp false successor argument type mismatch"
);
}
op_idx += 1;
// Pliron 0.19's default verifier sums in u32 without checking overflow.
// Keep its range arithmetic valid even for untrusted segment metadata.
let Some(total) = sizes
.0
.iter()
.try_fold(0u32, |sum, size| sum.checked_add(*size))
else {
return verify_err!(
op.loc(),
"MirCondBranchOp operand segment total overflows u32"
);
};
if usize::try_from(total).ok() != Some(op.get_num_operands()) {
return verify_err!(
op.loc(),
"MirCondBranchOp operand segment total must match operand count"
);
}

Ok(())
}
}

#[op_interface_impl]
impl BranchOpInterface for MirCondBranchOp {
fn successor_operands(&self, ctx: &Context, succ_idx: usize) -> Vec<Value> {
fn verify_successor_operand_layout(&self, ctx: &Context) -> Result<(), Error> {
// Interface verification can run before the operation's own verifier.
// Check every precondition of segment_range here, before it is called.
<Self as OperandSegmentInterface>::verify(self, ctx)?;
let op = self.get_operation().deref(ctx);
if op.get_num_successors() != 2 {
return verify_err!(op.loc(), "MirCondBranchOp must have exactly 2 successors");
}
if self.segment_size(ctx, 0) != 1 {
return verify_err!(
op.loc(),
"MirCondBranchOp must have exactly 1 condition operand"
);
}
Ok(())
}

fn successor_operand_range(&self, ctx: &Context, succ_idx: usize) -> Range<usize> {
assert!(
succ_idx == 0 || succ_idx == 1,
"MirCondBranchOp has exactly two successors"
);
// Segment 0 = condition, segment 1 = true args, segment 2 = false args
self.get_segment(ctx, succ_idx + 1)
// Segment 0 = condition, segment 1 = true args, segment 2 = false args.
self.segment_range(ctx, succ_idx + 1)
}

fn add_successor_operand(&self, ctx: &mut Context, succ_idx: usize, operand: Value) -> usize {
assert!(succ_idx < 2, "MirCondBranchOp has exactly two successors");
self.push_to_segment(ctx, succ_idx + 1, operand)
}

Expand All @@ -321,6 +326,7 @@ impl BranchOpInterface for MirCondBranchOp {
succ_idx: usize,
opd_idx: usize,
) -> Value {
assert!(succ_idx < 2, "MirCondBranchOp has exactly two successors");
self.remove_from_segment(ctx, succ_idx + 1, opd_idx)
}
}
Expand Down
23 changes: 14 additions & 9 deletions cuda-oxide/crates/dialect-mir/src/ops/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use pliron::{
common_traits::Verify,
context::{Context, Ptr},
identifier::Identifier,
indented_block, input_err,
input_err,
irfmt::{
parsers::{spaced, type_parser},
printers::op::{region, typed_symb_op_header},
Expand Down Expand Up @@ -171,10 +171,9 @@ impl Printable for MirFuncOp {
.retain(|key, _| key != &*ATTR_KEY_MIR_FUNC_TYPE && key != &ATTR_KEY_SYM_NAME);

if !attributes_to_print_separately.0.is_empty() {
indented_block!(state, {
write!(f, "{}", indented_nl(state))?;
attributes_to_print_separately.fmt(ctx, state, f)?;
});
let _indent = state.indent();
write!(f, "{}", indented_nl(state))?;
attributes_to_print_separately.fmt(ctx, state, f)?;
}
write!(f, " ")?;
region(self).fmt(ctx, state, f)?;
Expand Down Expand Up @@ -202,19 +201,25 @@ impl Parsable for MirFuncOp {
vec![],
vec![],
vec![],
1,
// Region::parser creates the body region itself.
0,
);
let mut parser = (
spaced(token('@').with(Identifier::parser(()))).skip(spaced(token(':'))),
spaced(type_parser()),
spaced(AttributeDict::parser(())),
spaced(optional(AttributeDict::parser(()))),
spaced(optional(Region::parser(op))),
);
parser
.parse_stream(state_stream)
.map(|(fname, fty, attrs, _region)| -> OpObj {
.map(|(fname, fty, attrs, region)| -> OpObj {
let ctx = &mut state_stream.state.ctx;
op.deref_mut(ctx).attributes = attrs;
// Declarations still have one empty region, matching the
// function operation's OneRegionInterface contract.
if region.is_none() {
Operation::add_region(op, ctx);
}
op.deref_mut(ctx).attributes = attrs.unwrap_or_default();
let ty_attr = TypeAttr::new(fty);
let opop = MirFuncOp { op };
opop.set_symbol_name(ctx, fname);
Expand Down
Loading
Loading