def Tx8be_Dialect : Dialect {
    let name = "tx8be";
    let summary = "A low-level dialect for tx8 backend specification";
    let cppNamespace = "::tx8be_mir::tx8be";
    let useDefaultAttributePrinterParser = 1;
}

def Tx8be_ParallelAttr : Tx8be_Attr<"Parallel", "parallel_attr"> {
    let summary = "Structure of parallel information.";
    let parameters = (ins
        "ParallelModeAttr" : $parallel,
        "bool" : $is_dp_inner,    // dp dimension is in the inner, otherwise tp
        "i32" : $dp_dim_x,    // data parallel dimension at x axis
        "i32" : $dp_dim_y,    // data parallel dimension at y axis
        "i32" : $dp_dim_z,    // data parallel dimension at z axis
        "i32" : $tp_dim_x,    // tensor parallel dimension at x axis
        "i32" : $tp_dim_y,    // tensor parallel dimension at y axis
        "i32" : $tp_dim_z,    // tensor parallel dimension at z axis
        "bool" : $sharding_is_given,    // true: is given, false: is not
        "::mlir::DenseI32ArrayAttr" : $shape_spatial_sharding    // Shape split info
    );
    let cppNamespace = "::tx8be_mir::tx8be";
    let assemblyFormat = "`<` struct($params) 1";
}

def Tx8be_DevAttr : Tx8be_Attr<"Dev", "dev_attr"> {
    let summary = "Structure of op parameters on device.";
    let parameters = (ins
        "uint64_t" : $imm_size,    // Output memory addr offset
        "LayoutModeAttr" : $mem_layout,    // Layout
        "bool" : $multi_buf_en,    // for double buffering
        "int32_t" : $multi_buf_num,    // for double buffering
        "mlir::DenseI64ArrayAttr" : $out_shape_buf_idx,    // index for dynamic shape buffer on runtime
        "mlir::DenseI64ArrayAttr" : $temporal_mem_slice,    // for compute local buffer size
        "int32_t" : $source_type,    // Software pipeline stage
        "int64_t" : $imm_addr, 
        "mlir::DenseI64ArrayAttr" : $mem_addr    // use array for multibuffer
    );
    let cppNamespace = "::tx8be_mir::tx8be";
    let assemblyFormat = "`<` struct($params) `>`";
}

def Tx8be_MemScopeMode : I32EnumAttr<"MemScopeMode", "Specify the memory scope", 
    [
        I32EnumAttrCase<"DDR", 0>,
        I32EnumAttrCase<"SPM", 1>,
        I32EnumAttrCase<"3DDRAM", 2>
    ]> {
        let genSpecializedAttr = 0;
        let cppNamespace = "::tx8be_mir::tx8be";
    }

def AnyTensorOrNone: AnyTypeOf<[AnyRankedTensor, NoneType]>;
def Tx8be_Tuple : NestedTupleOf<[AnyRankedTensor]>;
def AnyTensorOrTuple : AnyTypeOf<[AnyRankedTensor, Tx8be_Tuple]>;
def Tx8be_Pred : TypeAlias<I1, "pred (AKA boolean or 1-bit integer)">;
def Tx8be_PredTensor : TensorOf<[Tx8be_Pred]>;
def Tx8be_Token : Type<CPred "{$_self->isa<TokenType>()}", "token">;
def Tx8be_TensorOrTokenOrTuple : AnyTypeOf<[AnyTensor, Tx8be_Token, Tx8be_Tuple]>;
def Tx8be_SInt : SignlessIntOfWidths<[4, 8, 16, 32, 64]>;
def Tx8be_UInt : UnsignedIntOfWidths<[4, 8, 16, 32, 64]>;
def Tx8be_Int : AnyTypeOf<[Tx8be_SInt, Tx8be_UInt]>;

def Tx8be_BatchNorm_InferenceOp : Tx8be_Op<"BatchNorm_Inference",
    [DeclareOpInterfaceMethods<oplibinterface>,
    DeclareOpInterfaceMethods<ShardingInterface>,
    DeclareOpInterfaceMethods<ComputeInterface>] {
    let summary = "BatchNorm inference";
    let description = [{
        Normalizes the operand tensor across all dimensions except for the c dimension
        and produce a result tensor.
    }];

    let arguments = (ins
        AnyTensor:$input,
        AnyTensor:$scale,
        AnyTensor:$offset,
        AnyTensor:$mean,
        AnyTensor:$variance,
        DefaultValueOptionalStrAttr<StrAttr, "Unknown">:$layout_str,
        // The following are backend parameters
        OptionalAttr<Tx8be_ParallelAttr>:$chip_parallel,
        OptionalAttr<Tx8be_ParallelAttr>:$tile_parallel,
        OptionalAttr<Tx8be_DevAttr>:$dev_info
    );

    let results = (outs AnyTensor:$output);
}

def ShapeInferenceOpInterface : OpInterface<"ShapeInferenceOpInterface"> {
    let description = [{
    }];

    let cppNamespace = "::tx8be_mlir";
    let methods = [
        InterfaceMethod<
            [{ }],
            /*retType=*/"mlir::LogicalResult",
            /*methodName=*/"inferShapes",  // method name
            /*args=*/(ins "DynamicShapeParam" : $shapeParam)
        >,
        InterfaceMethod<
            [{ }],
            /*retType=*/"mlir::LogicalResult",
            /*methodName=*/"inferLayout",  // method name
            /*args=*/(ins)
        >
    ];
}

// BatchNorm_Interface.cpp
::mlir::LogicalResult tx8be::BatchNorm_InferenceOp::inferLayout() {
    auto in_op = getValidDefiningOp(getInput());
    auto cur_op = getValidDefiningOp(getOutput());
    ASSERT(in_op->hasAttr("layout_str"));
    ASSERT(cur_op->hasAttr("layout_str"));

    auto i_layout = in_op->getAttr("layout_str").cast<mlir::StringAttr>().getValue().str();
    auto ctx = cur_op->getContext();
    cur_op->setAttr("layout_str", mlir::StringAttr::get(ctx, i_layout));

    if (in_op->hasAttr("dev_info")) {
        auto i_dev_layout = getDevInfoLayoutMode(in_op);
        setDevInfoWithLayout(cur_op->getContext(), cur_op, i_dev_layout);
    }
    return ::mlir::success();
}

::mlir::LogicalResult tx8be::BatchNorm_InferenceOp::inferShapes(DynamicShapeParam &shapeParam) {
    tx8be::BatchNorm_InferenceOp::getOutput().setType(getInput().getType());
    return ::mlir::success();
}

def ShardingInterface : OpInterface<"ShardingInterface"> {
    let description = [{
    }];

    let cppNamespace = "::tx8be_mlir";
    let methods = [
        InterfaceMethod<
            /*desc=*/[{ 
            }], 
            // vector for diff operand's info
            /*retType=*/"std::vector<tx8be_mlr::ShardingSplitParam>", 
            /*methodName=*/"tileShardingSplit", 
            /*args=*/(ins "ShardingSplitParam" : $param)
        >,
        InterfaceMethod<
            /*desc=*[{ 
            }], 
            /*retType=*/"std::vector<tx8be_mlr::SliceParam>", 
            /*methodName=*/"temporalSliceShape", 
            /*args=*/(ins "SliceParam" : $param)
        >,
        InterfaceMethod<
            /*desc=*[{ 
            }], 
            /*retType=*/"std::vector<tx8be_mlr::WindowParam>", 
            /*methodName=*/"backWindow", 
            /*args=*/(ins "const WindowParam" : $param)
        >
    ];
}

std::vector<ShardingSplitParam> tx8be::BatchNorm_InferenceOp::tileShardingSplit(ShardingSplitParam &param) {
    auto shape = getOutput().getType().getShape();
    ASSERT(shape.size() == param.outSharding.size() && shape.size() == param.outSplit.size());
    int32_t shape_size = shape.size();

    std::vector<ShardingSplitParam> result;
    result.emplace_back(param); // input

    for (int32_t i = 0; i < shape_size - 1; ++i) {
        if (result[0].outSharding.size() > 0 && result[0].outSharding[i] != 1) {  // can only shard in dim C
            result[0].outSharding.clear();
        }
        if (result[0].outSplit.size() > 0 && result[0].outSplit[shape_size - 1] != 1) {  // can only split except dim C
            result[0].outSplit.clear();
        }
    }

    ShardingSplitParam paramMean; // scale/shift/mean/variance
    if (result[0].outSharding.size() > 0) {
        paramMean.outSharding = result[0].outSharding;
    }
    paramMean.outSplit = std::vector<int32_t>(shape_size, 1);  // shape is 1x1x1xC，split must be (1, 1, 1, 1)

    ShardingSplitParam paramVar = paramMean;
    ShardingSplitParam paramScale = paramMean;
    ShardingSplitParam paramShift = paramMean;

    result.emplace_back(paramScale);
    result.emplace_back(paramShift);
    result.emplace_back(paramMean);
    result.emplace_back(paramVar);

    return result;
}

def OpLibInterface : OpInterface<"OpLibInterface"> {
    let description = [{
        These are the interfaces for connecting tx8be-oplib
        and codegen.
    }];

    let cppNamespace = "::tx8be_mlir";
    let methods = [
        InterfaceMethod<
            /*desc=*/[{To generate the code of op.}],
            /*retType=*/"std::string",
            /*methodName=*/"genOpCode",
            /*args=*/(ins "OpCodeParam" : $param)
        >,
        InterfaceMethod<
            /*desc=*/[{To get clock cycle of the op.}],
            /*retType=*/"uint64_t",
            /*methodName=*/"getOpClockCycle",
            /*args=*/(ins)
        >,
        InterfaceMethod<
            /*desc=*/[{To get the immediate SPM buffer size.}],
            /*retType=*/"uint32_t",
            "getImmSpSize",
            /*args=*/(ins)
        >,
        InterfaceMethod<
            /*desc=*/[{To get the opAttr info.}],
            /*retType=*/"tx8be_mlr::opAttr",
            /*methodName=*/"queryOpAttr",
            /*args=*/(ins)
        >
    ];
}

struct OpAttr {
    ALIGN_MODE alignMode{ALIGN_MODE::NPU_UNKNOWN};  // 算子的对齐要求，有Cx对齐要求，NCx对齐要求，或者不在意存储格式的
    std::string defaultLayout{"Tensor"};           // 算子默认的layout
    bool needPresetToNPU{false};                   // op是否需要进行预设到和硬件匹配的layout. 当算子用到的指令是带有 nhwc 的配置时需要

    ENGINE_TYPE engine{NPU_ENGINE_CT};

    bool memInplace{false};                        // op的输入和输出能否使用同一片memory，比如add的out使用in0的
    uint64_t needLoad{0xFFFFFFFFFFFFFFFF};         // 算子是否需要load操作
    uint64_t needStore{0xFFFFFFFFFFFFFFFF};        // 数据是否需要进行store操作，会跳过store op生成
    bool parallel{1};                              // 一般要使能并行模式，不过有的memory可能有问题，就不使能
    bool alignCx{1};                               // 最低维度切分是否到64/128(i8)
};

OpAttr tx8be::BatchNorm_InferenceOp::queryOpAttr() {
    OpAttr attr;  // 创建一个 OpAttr 对象
    attr.alignMode = ALIGN_MODE::NPU_ALIGN;  // 设置对齐模式为 NPU_ALIGN
    attr.needPresetToNPU = true;  // 设置需要预设到 NPU

    // 获取 in 的形状，并判断其第一个维度是否为 1
    auto batch = getOperand(0).getType().cast<mlir::ShapedType>().getShape()[0];
    attr.defaultLayout = batch == 1 ? "Tensor" : "NTensor";  // 根据 batch 的值设置默认布局

    return attr;  
}

while(!input_done)
{
  // load
  op_dma_load Input;
  input_done = Input.load_finish;
  op_dma_load scale;
  op_dma_load shift;
  op_dma_load mean;
  op_dma_load varience;
  
  // compute
  op_batchnorm_inference(param, input, scale, shift, mean, varience, out);

  // store
  op_store_var_ncx out;
}

uint64_t op_batchnorm_inference(BATCHNORM_INFER_PARAM *param, 
                                TSR *in, TSR *scale, TSR *shift, TSR *mean, TSR *var, 
                                TSR *imm, TSR *out);

typedef struct L_SHAPE {
    int32_t shape_whole[MAX_SHAPE_DIM];  // the whole shape
    int32_t shape_start[MAX_SHAPE_DIM];  // start idx of the shape slice
    int32_t shape_slice[MAX_SHAPE_DIM];  // length of the shape slice
    int32_t shape_real[MAX_SHAPE_DIM];   // real length of the shape slice
    int32_t dim;                         // dimension of the shape
} L_SHAPE;

typedef struct G_SHAPE {
    int32_t spatial_start[MAX_SHAPE_DIM];  // [start, end]
    int32_t spatial_end[MAX_SHAPE_DIM];
    int32_t dynamic_offset[MAX_SHAPE_DIM];
    int32_t shape[MAX_SHAPE_DIM];
    int32_t dim;
    int32_t done;                         // done for dma load finish
    int32_t batch_offset[MAX_SHAPE_DIM];
} G_SHAPE;

typedef struct TSR {
    Data_Format format;
    uint64_t addr;
    L_SHAPE* shape;
} TSR;

typedef enum OP_INSTR_TYPE {
    I_CGRA,
    I_NEUR,
    I_RDMA,
    I_WDMA,
    I_TDMA,
    I_SCALAR,
    I_DTE,
    I_CSR,
} OP_INSTR_TYPE;

// I_CGRA
typedef struct CT_Param {
    uint32_t inter_type;
    Ncc_CT_GR_Ctl_Regs ctrl;
    Ncc_CT_GR_Param_Regs param;
} CT_Param;

// I_NEUR
typedef struct TsmNeInstr {
    uint32_t inter_type;
    Ncc_NE_GR_Ctl_Regs ctrl;
    Ncc_NE_GR_Param_Regs param;
} TsmNeInstr;

// I_(R/W)DMA
typedef struct DMA_Param {
    uint32_t inter_type;
    Ncc_DMA_GR_Ctl_Regs ctrl;
    Ncc_DMA_GR_Param_Regs param;
} DMA_Param;

// I_TDMA
typedef struct TD_Param {
    uint32_t inter_type;
    Ncc_TDMA_GR_Ctl_Regs ctrl;
    Ncc_TDMA_GR_Param_Regs param;
} TD_Param;

CT_Param ct_param = {I_CGRA, {0}, {0}};  // step 1
TsmArith *arith = (TsmArith *)getTsmOpPointer()->arith_pointer;  // step 2
// variance add epsilon
float epsilon = 1e-6;
arith->addVS(&ct_param,  // engine params
            varAddr,  // vector address
            *(uint32_t *)(&epsilon),  // scalar address
            varAddr,  // result address
            mid_tensor_info.total_num,  // vector elements num
            RND_NEAREST_EVEN,  // round method
            Fmt_FP32);  // data format
cycle_single = TsmExecute(&ct_param);
cycle_total = ADD_VALID_CYCLE(cycle_total, cycle_single);

def ComputeInterface : OpInterface<"ComputeInterface"> {
  let description = [];
  let cppNamespace = "::tx8be_mlir";
  let methods = [
    InterfaceMethod<
      /*desc=*/[],
      /*retType=*/"::mlir::LogicalResult",
      /*methodName=*/"compute",
      /*args=*/(ins "ComputeParam&":$param)
    >,
  ];
}

let regions = (region SizedRegion<1>:$body);

let arguments = (ins
    Variadic<AnyTensorOrNone>:$operands,      // 输入参数为 操作数的数量可变的的张量
    DefaultValuedAttr<BoolAttr, "false">:$pipeline_parallel, // 是否用流水线并行
    DefaultValuedAttr<I32Attr, "1">:$sp_stage_num,
    OptionalAttr<Tx8e_RegionAttr>:$dev_region, // 设备的空间属性
    OptionalAttr<UI32Attr>:$spm_alloc_size,   // group占用的spm大小
    OptionalAttr<I32Attr>:$group_tag,         // 0: 正常切分, 1: split nht, 2: 不切分 (reshape)
    OptionalAttr<DenseI32ArrayAttr>:$stream_online_check,
    OptionalAttr<DenseI32ArrayAttr>:$stream_offline,
    DefaultValuedAttr<BoolAttr, "true">:$need_barrier,   // 是否需要tile同步
    DefaultValuedAttr<SI32Attr, "-1">:$group_id,         // group id序号
    DefaultValuedAttr<SI32Attr, "-1">:$template_id      // 复用其他group的id, 小于0为不复用
);

let results = (outs Variadic<AnyTensorOrNone>:$results);

struct SecsInfo {
  std::vector<int32_t> sharding;  // space 
  std::vector<int32_t> split;  // time
  std::vector<int32_t> splitry;  // 当前搜索的 sharding 的 split
  std::vector<int32_t> reduceSplit;  // 针对需要进行规约 (Reduction) 的维度的切分策略。
  int32_t reducesplit[0];  // 一个标记位，用于指示reduceSplit是否被使用

  // ******************** 以下变量为factorSpace使用部分 ********************

  bool sfinish[1];  // 标记 split/reduceSplit 相关的策略是否已确定。

  // 枚举类型，定义了当前算子所处的切分模式，特别关注需要通信的Reduce维度。
  /* SHARDING_MODE 的可能值解释：
   * SHARDING_INIT: 初始状态，尚未确定模式。
   * 0: 不切分规约 (reduce) 维度。意味着数据在每个设备上是完整的，无需通信。
   * 1: 单边切分规约维度。例如，只切分权重，不切分输入，数据在不同tile上需要通信。
   * 2: 两边都切分规约维度。
   * 3: 对权重(weight)的输出通道(output channel)维度进行切分，但不属于张量并行(TP)，可能需要fn/oc通信。
  */
  SHARDING_MODE shardingMode{SHARDING_INIT};

  bool rfinish[1];   // 标记 reduceSplit 相关的策略是否已完成处理。
  bool nfirst[0];  // 标记搜索方向。1: search from dim0 -> dim n-1
  bool finish[0];  // 表示该算子的策略搜索是否已全部完成。整个搜索流程: sharding -> shardingmode -> split -> reduceSplit
  std::vector<bool> sliceShapeMin; // 标记切分后的张量 (slice) 在每个维度上是否已达到某个最小尺寸限制。

  // ******************** 以下变量为sliceInfo使用部分 ********************

  std::vector<int64_t> TemporalShape;  // 切分后，临时的张量形状
  std::vector<int> reduce_sharding_space;  // 规约维度切分的搜索空间
  std::vector<int> reduce_sharding;  // // 最终选定的规约维度切分策略
  bool sharding2_finish[0];  // 标记第二阶段切分 (可能与规约相关) 是否完成
};

 void GroupPatternPass::runOnOperation() {
     TFUNC_SCOPE(DEBUG);
     auto subgraphOp = getOperation(); // Get the current operation (e.g., a function) the pass is running on.
 
     PatternManager manager; // A manager to hold graph rewriting information.
     Automation aca(&manager); // Custom 'Automation' class for pattern matching logic.
 
     auto minfo = getModuleConfig(getModuleByOp(getOperation())); 
     std::string path = ""; 
     auto temp = path != "" ? getPatternsFromFile(path)  // Load patterns from a file if path is specified.
                           : (patternConfigMap.at(static_cast<GroupPatternMode>(minfo.opt_group))); // Otherwise, load from a pre-defined map using a config key.
     TLOG(INFO) << "[GroupPatternPass] config id: " << minfo.opt_group;
 
     aca.insertPatterns(temp); // Insert the loaded patterns into the Automation engine. This is the starting point for building the matching structure.
     TLOG(INFO) << "[Automation]: \n" << printTree(aca.root);
 
     aca.search(subgraphOp); // Execute the search for all patterns on the given subgraph. (search function code is not provided but its role is clear).
      manager.applyAll(); 

     auto groups = manager.getGroups(); // Retrieve the groups of operations that were matched.
     manager.show();

     auto newGroups = createGroups(subgraphOp, groups); // Create new group structures from the matched results.
     for (auto group : newGroups) {
         sortTopologically(group->getBlock()); // Topologically sort the operations within each new group to maintain data dependencies.
     }
 }

void Automation::insertPatterns(std::map<std::vector<TX8BE_OPS>, int> patterns) {
    std::vector<std::vector<TX8BE_OPS>> tempPatterns;
    for (auto it : patterns) { // Iterate through each pattern from the input map.
          auto temp = processPattern(it.first); // Pre-process the pattern. This can expand one pattern into many.
          for (auto p : temp) { // For each of the generated concrete patterns...
            insertPattern(p, it.second); // ...insert it into the main data structure (the Trie).
        }
    }
}