refactor: 移除 MaybeDyn 和复杂的 TensroLayout

Signed-off-by: YdrMaster <ydrml@hotmail.com>

Author: YdrMaster

operators/src/.clang-format

@@ -1,66 +1,29 @@
-# Generated from CLion C/C++ Code Style settings
+---
 BasedOnStyle: LLVM
-AccessModifierOffset: -4
-AlignAfterOpenBracket: Align
-# AlignConsecutiveAssignments: None
-AlignOperands: Align
-AllowAllArgumentsOnNextLine: false
-AllowAllConstructorInitializersOnNextLine: false
-AllowAllParametersOfDeclarationOnNextLine: false
-AllowShortBlocksOnASingleLine: Always
-AllowShortCaseLabelsOnASingleLine: false
-AllowShortFunctionsOnASingleLine: All
-AllowShortIfStatementsOnASingleLine: Always
-AllowShortLambdasOnASingleLine: All
-AllowShortLoopsOnASingleLine: true
-AlwaysBreakAfterReturnType: None
-AlwaysBreakTemplateDeclarations: No
-BreakBeforeBraces: Custom
+IndentWidth: 4                        # 缩进宽度，LLVM 默认值为 2，改为 4
+AccessModifierOffset: -4              # public/protected/private 访问控制符相对成员的偏移，与 IndentWidth 配合，LLVM 默认值为 -2
+AlignOperands: AlignAfterOperator     # 双目运算符的行间对齐，LLVM 默认值为 Align，改为带符号一起换行
+ColumnLimit: 0                        # 列宽限制，LLVM 默认值为 80，改为不限制
+AllowShortBlocksOnASingleLine: Always # 是否允许短块（单个语句的块）不换行，LLVM 默认值为 Never，改为允许
+AllowShortLoopsOnASingleLine: true    # 是否允许短循环不换行，LLVM 默认值为 false，改为允许
+InsertBraces: true                    # 是否在 if/for/while/switch 等语句后插入大括号，LLVM 默认值为 false，改为允许
+BreakBeforeBraces: Custom             # 大括号换行配置，LLVM 默认值为 LLVM，改为自定义以使 BraceWrapping 生效
 BraceWrapping:
   AfterCaseLabel: false
   AfterClass: false
   AfterControlStatement: Never
   AfterEnum: false
   AfterFunction: false
   AfterNamespace: false
+  AfterObjCDeclaration: false
+  AfterStruct: false
   AfterUnion: false
+  AfterExternBlock: false
   BeforeCatch: false
   BeforeElse: false
+  BeforeLambdaBody: false
+  BeforeWhile: false
   IndentBraces: false
-  SplitEmptyFunction: false
+  SplitEmptyFunction: true
   SplitEmptyRecord: true
-BreakBeforeBinaryOperators: None
-BreakBeforeTernaryOperators: true
-BreakConstructorInitializers: BeforeColon
-BreakInheritanceList: BeforeColon
-ColumnLimit: 0
-CompactNamespaces: true
-ContinuationIndentWidth: 4
-IndentCaseLabels: true
-IndentPPDirectives: None
-IndentWidth: 4
-KeepEmptyLinesAtTheStartOfBlocks: true
-MaxEmptyLinesToKeep: 2
-NamespaceIndentation: All
-ObjCSpaceAfterProperty: false
-ObjCSpaceBeforeProtocolList: true
-PointerAlignment: Right
-ReflowComments: false
-SpaceAfterCStyleCast: true
-SpaceAfterLogicalNot: false
-SpaceAfterTemplateKeyword: false
-SpaceBeforeAssignmentOperators: true
-SpaceBeforeCpp11BracedList: false
-SpaceBeforeCtorInitializerColon: true
-SpaceBeforeInheritanceColon: true
-SpaceBeforeParens: ControlStatements
-SpaceBeforeRangeBasedForLoopColon: true
-SpaceInEmptyParentheses: false
-SpacesBeforeTrailingComments: 0
-SpacesInAngles: false
-SpacesInCStyleCastParentheses: false
-SpacesInContainerLiterals: false
-SpacesInParentheses: false
-SpacesInSquareBrackets: false
-TabWidth: 4
-UseTab: Never
+  SplitEmptyNamespace: true

operators/src/add/args.rs

@@ -1,6 +1,6 @@
 use crate::{
-    get_static, rank_mismatch, shape_mismatch, shape_not_support, utils::type_distinct, ConstPtr,
-    Hardware, LaunchError, MutPtr, TensorLayout,
+    rank_mismatch, shape_mismatch, shape_not_support, utils::type_distinct, ConstPtr, Hardware,
+    LaunchError, MutPtr, TensorLayout,
 };
 use digit_layout::DigitLayout;
 use itertools::izip;
@@ -48,7 +48,7 @@ impl Scheme {
             ..
         } = args;
         // # 检查基本属性
-        let dt = type_distinct(&[c.dt(), a.dt(), b.dt()])?;
+        let dt = type_distinct(&[c.dt, a.dt, b.dt])?;
         let ndim = c.ndim();
         if a.ndim() != ndim || b.ndim() != ndim {
             return Err(rank_mismatch(format!(
@@ -68,17 +68,13 @@ impl Scheme {
         }
         let mut dims = Vec::with_capacity(ndim);
         for (&d, &da, &db, &sc, &sa, &sb) in izip!(
-            c.shape(),
-            a.shape(),
-            b.shape(),
+            c.shape_group(),
+            a.shape_group(),
+            b.shape_group(),
             c.strides(),
             a.strides(),
-            b.strides()
+            b.strides(),
         ) {
-            get_static! {
-                d  da db
-                sc sa sb
-            }
             if da != d || db != d {
                 return Err(shape_mismatch(format!(
                     "c: {:?}, a: {:?}, b: {:?}",

operators/src/add/cuda/add.cuh

@@ -1,4 +1,4 @@
-template<class Tdata>
+template <class Tdata>
 static __device__ void _add(
     Tdata *__restrict__ c,
     Tdata const *__restrict__ a,

operators/src/add_rows/args.rs

@@ -1,7 +1,7 @@
 use crate::{
     type_not_support,
     utils::{dim_distinct, rank_error, type_distinct},
-    ConstPtr, Hardware, LaunchError, MaybeDyn, MutPtr, TensorLayout,
+    ConstPtr, Hardware, LaunchError, MutPtr, TensorLayout,
 };
 use digit_layout::{DigitLayout, LayoutContent::Unsigned};
 use std::ptr::{null, null_mut};
@@ -37,10 +37,10 @@ impl<H: Hardware> Args<H> {
 pub(super) struct Meta {
     pub dt: DigitLayout,
     pub dt_idx: DigitLayout,
-    pub batch: MaybeDyn<usize>,
-    pub m: MaybeDyn<usize>,
-    pub n: MaybeDyn<usize>,
-    pub k: MaybeDyn<usize>,
+    pub batch: usize,
+    pub m: usize,
+    pub n: usize,
+    pub k: usize,
 }
 
 impl<H: Hardware> Args<H> {
@@ -52,30 +52,30 @@ impl<H: Hardware> Args<H> {
             ..
         } = self;
 
-        let dt = type_distinct(&[dst.dt(), src.dt()])?;
-        let dt_idx = idx.dt();
+        let dt = type_distinct(&[dst.dt, src.dt])?;
+        let dt_idx = idx.dt;
         if !matches!(dt_idx.decode(), Unsigned { .. }) {
             return Err(type_not_support(format!(
                 "data type {dt_idx} is not supported, must be unsigned integers"
             )));
         }
 
-        let &[batch, m, n] = dst.shape() else {
+        let &[batch, m, n] = &*dst.shape() else {
             return Err(rank_error("dst", 3, dst.ndim()));
         };
-        let &[k, n_] = src.shape() else {
+        let &[k, n_] = &*src.shape() else {
             return Err(rank_error("src", 2, src.ndim()));
         };
-        let &[batch_, m_] = idx.shape() else {
+        let &[batch_, m_] = &*idx.shape() else {
             return Err(rank_error("idx", 2, idx.ndim()));
         };
 
         Ok(Meta {
             dt,
             dt_idx,
-            batch: dim_distinct(&[batch, batch_])?,
-            m: dim_distinct(&[m, m_])?,
-            n: dim_distinct(&[n, n_])?,
+            batch: dim_distinct(&[batch, batch_]).expect("batch mismatch"),
+            m: dim_distinct(&[m, m_]).expect("m mismatch"),
+            n: dim_distinct(&[n, n_]).expect("n mismatch"),
             k,
         })
     }

operators/src/add_rows/common_cpu/mod.rs

@@ -1,5 +1,5 @@
 use super::{args::Meta, AddRows, Args};
-use crate::{common_cpu::Cpu, get_static, ByteOf, LaunchError, QueueAlloc, Unsigned};
+use crate::{common_cpu::Cpu, ByteOf, LaunchError, QueueAlloc, Unsigned};
 use digit_layout::types as ty;
 use half::f16;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
@@ -54,12 +54,6 @@ impl crate::Operator for Operator {
             unreachable!()
         };
 
-        get_static! {
-            b   m   n   k
-            bsd msd nsd
-            bsi msi nss kss
-        }
-
         let dst = *dst_base as usize;
         let src = *src_base as usize;
         let idx = *idx_base as usize;

operators/src/add_rows/cuda/add_rows.cuh

@@ -1,4 +1,4 @@
-template<class Tdata, class Tidx>
+template <class Tdata, class Tidx>
 static __device__ void add_rows(
     Tdata *__restrict__ dst,
     Tdata const *__restrict__ src,

operators/src/add_rows/cuda/mod.rs

@@ -2,7 +2,7 @@ use super::{AddRows, Args};
 use crate::{
     add_rows::args::Meta,
     cuda::{dt_name, Gpu, Handle, ModuleBox},
-    get_static, strides_not_support,
+    strides_not_support,
     utils::gcd,
     ByteOf, LaunchError, QueueAlloc, SchemeDiversity,
 };
@@ -63,13 +63,8 @@ impl crate::Operator for Operator {
             unreachable!()
         };
 
-        get_static! {
-            b   n   m
-            bsd msd nsd
-            bsi msi nss kss
-        }
-        let unit_dst = dst_layout.dt().nbytes() as isize;
-        let unit_idx = idx_layout.dt().nbytes() as isize;
+        let unit_dst = dst_layout.dt.nbytes() as isize;
+        let unit_idx = idx_layout.dt.nbytes() as isize;
         if nsd != unit_dst || nss != unit_dst || msi != unit_idx {
             return Err(strides_not_support(""));
         };
@@ -85,9 +80,7 @@ impl crate::Operator for Operator {
         let params = cuda::params![dst_base, src_base, idx_base, bsd, msd, kss, bsi];
         let block = gcd(self.max_threads_block, n);
         let dimx = n.div_ceil(block);
-        let key = SchemeKey {
-            dt: dst_layout.dt(),
-        };
+        let key = SchemeKey { dt: dst_layout.dt };
         let scheme = self
             .schemes
             .lock()
@@ -160,6 +153,7 @@ mod test {
     };
     use half::f16;
 
+    #[allow(clippy::too_many_arguments)]
     fn args<H: Hardware>(
         dt: DigitLayout,
         b: usize,

operators/src/all_reduce/args.rs

@@ -1,10 +1,8 @@
 use super::ReduceOp;
 use crate::{
-    dyn_not_support, rearrange, shape_mismatch, strides_not_support, utils::type_distinct,
-    Hardware, LaunchError, MaybeDyn,
+    rearrange, shape_mismatch, strides_not_support, utils::type_distinct, Hardware, LaunchError,
 };
 use digit_layout::DigitLayout;
-use ndarray_layout::ArrayLayout;
 
 pub struct Args<H: Hardware> {
     pub pair: rearrange::Args<H>,
@@ -35,15 +33,9 @@ impl<H: Hardware> Args<H> {
             ..
         } = self;
 
-        let dt = type_distinct(&[dst_layout.dt(), src_layout.dt()])?;
+        let dt = type_distinct(&[dst_layout.dt, src_layout.dt])?;
 
-        let Some(shape) = MaybeDyn::get_all(dst_layout.shape()) else {
-            return Err(dyn_not_support(""));
-        };
-        let Some(strides) = MaybeDyn::get_all(dst_layout.strides()) else {
-            return Err(dyn_not_support(""));
-        };
-        let dst = ArrayLayout::<2>::new(shape, strides, 0);
+        let dst = &dst_layout.layout;
         let &[dst] = dst
             .merge_be(0, dst.ndim())
             .ok_or(strides_not_support(""))?
@@ -52,13 +44,7 @@ impl<H: Hardware> Args<H> {
             unreachable!()
         };
 
-        let Some(shape) = MaybeDyn::get_all(src_layout.shape()) else {
-            return Err(dyn_not_support(""));
-        };
-        let Some(strides) = MaybeDyn::get_all(src_layout.strides()) else {
-            return Err(dyn_not_support(""));
-        };
-        let src = ArrayLayout::<2>::new(shape, strides, 0);
+        let src = &src_layout.layout;
         let &[src] = src
             .merge_be(0, src.ndim())
             .ok_or(strides_not_support(""))?

operators/src/attention/args.rs

@@ -1,7 +1,7 @@
 use crate::{
     fuesd_softmax::AttnMask,
     utils::{dim_distinct, rank_error, type_distinct},
-    ConstPtr, Hardware, LaunchError, MaybeDyn, MutPtr, TensorLayout,
+    ConstPtr, Hardware, LaunchError, MutPtr, TensorLayout,
 };
 use digit_layout::DigitLayout;
 
@@ -23,11 +23,11 @@ pub struct Args<H: Hardware> {
 
 pub(super) struct Meta {
     pub dt: DigitLayout,
-    pub nh: MaybeDyn<usize>,
-    pub nkvh: MaybeDyn<usize>,
-    pub seq: MaybeDyn<usize>,
-    pub att: MaybeDyn<usize>,
-    pub dh: MaybeDyn<usize>,
+    pub nh: usize,
+    pub nkvh: usize,
+    pub seq: usize,
+    pub att: usize,
+    pub dh: usize,
 }
 
 impl<H: Hardware> Args<H> {
@@ -40,26 +40,26 @@ impl<H: Hardware> Args<H> {
             ..
         } = self;
 
-        let &[nh_q, seq_q, dh_q] = q_layout.shape() else {
+        let &[nh_q, seq_q, dh_q] = &*q_layout.shape() else {
             return Err(rank_error("q", 3, q_layout.ndim()));
         };
-        let &[nkvh_k, att_k, dh_k] = k_layout.shape() else {
+        let &[nkvh_k, att_k, dh_k] = &*k_layout.shape() else {
             return Err(rank_error("k", 3, k_layout.ndim()));
         };
-        let &[nkvh_v, att_v, dh_v] = v_layout.shape() else {
+        let &[nkvh_v, att_v, dh_v] = &*v_layout.shape() else {
             return Err(rank_error("v", 3, v_layout.ndim()));
         };
-        let &[nh_o, seq_o, dh_o] = o_layout.shape() else {
+        let &[nh_o, seq_o, dh_o] = &*o_layout.shape() else {
             return Err(rank_error("o", 3, o_layout.ndim()));
         };
 
         Ok(Meta {
-            dt: type_distinct(&[q_layout.dt(), k_layout.dt(), v_layout.dt(), o_layout.dt()])?,
-            nh: dim_distinct(&[nh_q, nh_o])?,
-            nkvh: dim_distinct(&[nkvh_k, nkvh_v])?,
-            seq: dim_distinct(&[seq_q, seq_o])?,
-            att: dim_distinct(&[att_k, att_v])?,
-            dh: dim_distinct(&[dh_q, dh_k, dh_v, dh_o])?,
+            dt: type_distinct(&[q_layout.dt, k_layout.dt, v_layout.dt, o_layout.dt])?,
+            nh: dim_distinct(&[nh_q, nh_o]).expect("nh mismatch"),
+            nkvh: dim_distinct(&[nkvh_k, nkvh_v]).expect("nkvh mismatch"),
+            seq: dim_distinct(&[seq_q, seq_o]).expect("seq mismatch"),
+            att: dim_distinct(&[att_k, att_v]).expect("att mismatch"),
+            dh: dim_distinct(&[dh_q, dh_k, dh_v, dh_o]).expect("dh mismatch"),
         })
     }
 }

operators/src/attention/cuda.rs

@@ -6,6 +6,7 @@ mod test {
     use crate::{cuda::Gpu, ByteOf, Hardware, Operator as _, TensorLayout};
     use digit_layout::{types as ty, DigitLayout};
 
+    #[allow(clippy::too_many_arguments)]
     fn args<H: Hardware>(
         dt: DigitLayout,
         nh: usize,