diff --git a/hercules_cg/src/cpu.rs b/hercules_cg/src/cpu.rs index 552dc3a30ea9f03813838c6c6d279f8f9e941de2..b15cf30106c76640016b1957fe9906ac01c74858 100644 --- a/hercules_cg/src/cpu.rs +++ b/hercules_cg/src/cpu.rs @@ -839,6 +839,8 @@ impl<'a> CPUContext<'a> { // // ((0 * s1 + p1) * s2 + p2) * s3 + p3 ... let elem_size = self.codegen_type_size(elem, body)?; + let elem_align = get_type_alignment(&self.types, elem); + let aligned_elem_size = Self::round_up_to(&elem_size, elem_align, body)?; let mut acc_offset = "0".to_string(); for (p, s) in zip(pos, dims) { let p = self.get_value(*p, false); @@ -848,7 +850,7 @@ impl<'a> CPUContext<'a> { } // Convert offset in # elements -> # bytes. - acc_offset = Self::multiply(&acc_offset, &elem_size, body)?; + acc_offset = Self::multiply(&acc_offset, &aligned_elem_size, body)?; acc_ptr = Self::gep(&acc_ptr, &acc_offset, body)?; collect_ty = elem; } @@ -910,6 +912,8 @@ impl<'a> CPUContext<'a> { // The size of an array is the size of the element multipled by // the dynamic constant bounds. let mut acc_size = self.codegen_type_size(elem, body)?; + acc_size = + Self::round_up_to(&acc_size, get_type_alignment(&self.types, elem), body)?; for dc in bounds { acc_size = Self::multiply(&acc_size, &format!("%dc{}", dc.idx()), body)?; } diff --git a/hercules_cg/src/gpu.rs b/hercules_cg/src/gpu.rs index 1434175660a261775352df1e715cf26947a3e531..25bbf1be2dabfe1ec88631edd7bc61b4a27ad42d 100644 --- a/hercules_cg/src/gpu.rs +++ b/hercules_cg/src/gpu.rs @@ -1877,7 +1877,11 @@ namespace cg = cooperative_groups; ")".repeat(array_indices.len()) )); let element_size = self.get_size(*element_type, None); - index_ptr.push_str(&format!(" * ({})", element_size)); + let element_align = self.get_alignment(*element_type); + index_ptr.push_str(&format!( + " * (({} + {} - 1 / {} * {}))", + element_size, element_align, element_align, element_align + )); type_id = *element_type; } } @@ -2049,7 +2053,15 @@ namespace cg = cooperative_groups; Type::Array(element_type, extents) => { assert!(num_fields.is_none()); let array_size = multiply_dcs(extents); - format!("{} * {}", self.get_size(*element_type, None), array_size) + let elem_align = self.get_alignment(type_id); + format!( + "(({} + {} - 1) / {} * {}) * {}", + self.get_size(*element_type, None), + elem_align, + elem_align, + elem_align, + array_size + ) } Type::Product(fields) => { let num_fields = num_fields.unwrap_or(fields.len()); diff --git a/hercules_cg/src/rt.rs b/hercules_cg/src/rt.rs index 3db0f16f35c711372841ced574c72cc50ee007ec..8fa0c09ee512e3f2e43c5280bc3cb6947bc31dc5 100644 --- a/hercules_cg/src/rt.rs +++ b/hercules_cg/src/rt.rs @@ -1111,6 +1111,13 @@ impl<'a> RTContext<'a> { // // ((0 * s1 + p1) * s2 + p2) * s3 + p3 ... let elem_size = self.codegen_type_size(elem); + let elem_align = get_type_alignment(&self.module.types, elem); + let aligned_elem_size = format!( + "(({} + {}) & !{})", + elem_size, + elem_align - 1, + elem_align - 1 + ); for (p, s) in zip(pos, dims) { let p = self.get_value(*p, bb, false); acc_offset = format!("{} * ", acc_offset); @@ -1119,7 +1126,7 @@ impl<'a> RTContext<'a> { } // Convert offset in # elements -> # bytes. - acc_offset = format!("({} * {})", acc_offset, elem_size); + acc_offset = format!("({} * {})", acc_offset, aligned_elem_size); collect_ty = elem; } } @@ -1192,6 +1199,13 @@ impl<'a> RTContext<'a> { // The size of an array is the size of the element multipled by // the dynamic constant bounds. let mut acc_size = self.codegen_type_size(elem); + let elem_align = get_type_alignment(&self.module.types, elem); + acc_size = format!( + "(({} + {}) & !{})", + acc_size, + elem_align - 1, + elem_align - 1 + ); for dc in bounds { acc_size = format!("{} * ", acc_size); self.codegen_dynamic_constant(*dc, &mut acc_size).unwrap();