2023年4月

一、前提

  • pId需要传入用来确认第一级的父节点,而且pId可以为null。
  • 树实体类必须实现:TreeNode接口
  • MyTreeVo必须有这三个属性:id、pId、children
  • 可以根据不同需求,配置TreeNode和MyTreeVo中固定的属性

二、代码

  • 定义TreeNode接口
public interface TreeNode {
    String getId();
    String getpId();
    List getChildren();
}


  • 需要将pId作为参数传入,在方法中添加一个pId参数,用于确认第一级的父节点。
    -- 这个示例代码中,buildTree方法接收两个参数,一个是泛型类型的参数List,另一个是pId,用于确认第一级的父节点。在找到所有的根节点后,我们需要对每个根节点调用buildChildren方法,递归构建子节点。
    -- 以便支持传入的pId作为顶级节点
public class TreeBuilder {
    /**
     * @param nodes 需要构建的nodes,不要继承 TreeNode,且需要有 TreeNode 下的三个方法
     * @param pId 父id。为null则是所有pid为null的是顶级节点;不为null,则是从给的pId往下查询并构建树
     **/
    public static <T extends TreeNode> List<T> buildTree(List<T> nodes, String pId) {
        List<T> roots = new ArrayList<>();
        for (T node : nodes) {
            if (node.getId().equals(pId) || (node.getParentId() == null && (pId == null || pId.equals(node.getId())))) {
                roots.add(node);
            }
        }
        for (T root : roots) {
            buildChildren(root, nodes);
        }
        return roots;
    }

    private static <T extends TreeNode> void buildChildren(T node, List<T> nodes) {
        for (T child : nodes) {
            if (child.getParentId() != null && child.getParentId().equals(node.getId())) {
                node.getChildren().add(child);
                buildChildren(child, nodes);
            }
        }
    }
}


  • 构建到几级
    -- 如果再增加一个参数,构建到树的第几层该如何实现?
    -- 根据需求,如果增加一个参数来控制构建到树的第几层,可以在buildChildren方法中添加一个level参数,用于记录当前节点的层数。在递归构建子节点时,如果当前节点的层数小于指定的层数,则继续递归构建子节点,否则停止递归。以下是修改后的Java示例代码:
public class TreeBuilder {
     
     /**
      * @param level 级别 从1开始。1是最顶层
      **/
    public static <T extends TreeNode> List<T> buildTree(List<T> nodes, String pId, int level) {
        List<T> roots = new ArrayList<>();
        for (T node : nodes) {
            if (node.getId().equals(pId) || (node.getParentId() == null && (pId == null || pId.equals(node.getId())))) {
                roots.add(node);
            }
        }
        for (T root : roots) {
            buildChildren(root, nodes, level, 1);
        }
        return roots;
    }

    private static <T extends TreeNode> void buildChildren(T node, List<T> nodes, int level, int currentLevel) {
        if (currentLevel >= level) {
            return;
        }
        for (T child : nodes) {
            if (child.getParentId() != null && child.getParentId().equals(node.getId())) {
                if (node.getChildren() == null) {
                    node.setChildren(new ArrayList<>());
                }
                node.getChildren().add(child);
                buildChildren(child, nodes, level, currentLevel + 1);
            }
        }
    }
}


三、使用

1、实现TreeNode接口

public class MyTreeVo implements TreeNode {

    /**
     * 主键
     */
    private String id;

    /**
     * 父节点ID
     */
    private String pId;

    /**
     * 子级
     */
    private List<MyTreeVo> children = Lists.newArrayList();

    //其他属性……

    public List<MyTreeVo> getChildren() {
        return children;
    }

    public String getId() {
        return id;
    }

    public String getpId() {
        return pId;
    }

    //其他属性的getter、setter……

2、使用

-- pId可以传入null,也可以传入需要从哪个节点(X)开始构造的 X的id
-- pId比如可以传入3、样例中的 “二、噢噢噢噢”的id=“e6ee51485389495cb923a122be800012”。然后构建出来的,就是“二、噢噢噢噢”的下级树

List<MyTreeVo> tree = TreeUtilQz.buildTree(vos,null);
//tree就是构建好的树结构数据

3、样例

{
    "data": [
        {
            "id": "e6ee51485389495cb923a122be800011",
            "pId": "",
            "name": "一、钢管钢管",
            "children": [
                {
                    "id": "e6ee51485389495cb923a122be800014",
                    "pId": "e6ee51485389495cb923a122be800011",
                    "name": "(二)嘎嘎嘎嘎嘎",
                    "children": []
                },
                {
                    "id": "e6ee51485389495cb923a122be800013",
                    "pId": "e6ee51485389495cb923a122be800011",
                    "name": "(一)顶顶顶顶",
                    "children": []
                }
            ]
        },
        {
            "id": "e6ee51485389495cb923a122be800012",
            "pId": "",
            "name": "二、噢噢噢噢",
            "children": [
                {
                    "id": "e6ee51485389495cb923a122be800015",
                    "pId": "e6ee51485389495cb923a122be800012",
                    "name": "二的下级",
                    "children": [
                        {
                            "id": "e6ee51485389495cb923a122be800016",
                            "pId": "e6ee51485389495cb923a122be800015",
                            "name": "二的下级的下级",
                            "children": []
                        }
                    ]
                }
            ]
        }
    ]
}

1 WMMA (Warp-level Matrix Multiply Accumulate) API

对于计算能力在7.0及以上的CUDA设备,可以使用CUDA C++ API调用Tensor Core,支持形如D = AB + C的混合精度的矩阵乘运算。
template<typename Use, int m, int n, int k, typename T, typename Layout=void> class fragment;

void load_matrix_sync(fragment<...> &a, const T* mptr, unsigned ldm);
void load_matrix_sync(fragment<...> &a, const T* mptr, unsigned ldm, layout_t layout);
void store_matrix_sync(T* mptr, const fragment<...> &a, unsigned ldm, layout_t layout);
void fill_fragment(fragment<...> &a, const T& v);
void mma_sync(fragment<...> &d, const fragment<...> &a, const fragment<...> &b, const fragment<...> &c, bool satf=false);
  • fragment:Tensor Core数据存储类,支持matrix_a、matrix_b和accumulator
  • load_matrix_sync:Tensor Core数据加载API,支持将矩阵数据从global memory或shared memory加载到fragment
  • store_matrix_sync:Tensor Core结果存储API,支持将计算结果从fragment存储到global memory或shared memory
  • fill_fragment:fragment填充API,支持常数值填充
  • mma_sync:Tensor Core矩阵乘计算API,支持D = AB + C或者C = AB + C

2 示例

以m16n16k16为例,实现HGEMM:C = AB,其中矩阵A(M * K,row major)、B(K * N,col major)和C(M * N,row major)的精度均为FP16。首先我们看如何使用CUDA Core写HGEMM naive算法。

2.1 CUDA Core

按照每个线程计算矩阵C中的一个元素来构建naive kernel,首先确定当前线程处理矩阵C的元素坐标,再遍历K并直接从global memory中加载所需A、B矩阵元素到寄存器参与计算,最后将计算结果从寄存器直接写回矩阵C。所有block计算完成之后即可得到矩阵C。这个例子不能说简单,只能说技术含量不高,不过我们只是为了对比。
#define DIV_CEIL(x, y) (((x) + (y) - 1) / (y))

__global__ void naiveKernel(const half *__restrict__ A, const half *__restrict__ B, half *__restrict__ C, size_t M,
                                size_t N, size_t K) {
    size_t row = threadIdx.x + blockDim.x * blockIdx.x;
    size_t col = threadIdx.y + blockDim.y * blockIdx.y;
    if (row < M && col < N) {
        half tmp = 0.0;
        for (size_t i = 0; i < K; ++i) {
            tmp += A[row * K + i] * B[i + col * K];
        }
        C[row * N + col] = tmp;
    }
}

void hgemmNaive(half *A, half *B, half *C, size_t M, size_t N, size_t K) {
    dim3 block(16, 16);
    dim3 grid(DIV_CEIL(M, block.x), DIV_CEIL(N, block.y));

    naiveKernel<<<grid, block>>>(A, B, C, M, N, K);
}

2.2 Tensor Core

我们再来看如何用WMMA API来构建naive kernel,参考
cuda sample
。与CUDA Core naive不同的是,WMMA需要按照每个warp处理一个矩阵C的WMMA_M * WMMA_N大小的tile的思路来构建,因为Tensor Core的计算层级是warp级别,计算的矩阵元素也是二维的。接下来,与CUDA Core naive的处理思路一致,首先确定当前warp处理矩阵C的tile坐标,声明计算tilie所需的fragment,再以WMMA_K为步长遍历K并直接从global memory中加载所需A、B矩阵tile到fragment参与计算,最后将计算结果从fragment直接写回矩阵C。所有block计算完成之后即可得到矩阵C。
值得注意的是,load_matrix_sync和store_matrix_sync都是按stride访问矩阵元素。
#include <mma.h>

#define WARP_SIZE 32

#define WMMA_M 16
#define WMMA_N 16
#define WMMA_K 16

using namespace nvcuda;

__global__ void wmmaNaiveKernel(const half *__restrict__ A, const half *__restrict__ B, half *__restrict__ C, size_t M,
                                size_t N, size_t K) {
    size_t warpM = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE;
    size_t warpN = (blockIdx.y * blockDim.y + threadIdx.y);

    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, half, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, half> c_frag;

    wmma::fill_fragment(c_frag, 0.0f);

    for (size_t i = 0; i < K; i += WMMA_K) {
        size_t aCol = i;
        size_t aRow = warpM * WMMA_M;
        size_t bCol = warpN * WMMA_N;
        size_t bRow = i;

        if (aRow < M && aCol < K && bRow < K && bCol < N) {
            wmma::load_matrix_sync(a_frag, A + aCol + aRow * K, K);
            wmma::load_matrix_sync(b_frag, B + bRow + bCol * K, K);

            wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
        }
    }

    size_t cCol = warpN * WMMA_N;
    size_t cRow = warpM * WMMA_M;

    if (cRow < M && cCol < N) {
        wmma::store_matrix_sync(C + cCol + cRow * N, c_frag, N, wmma::mem_row_major);
    }
}

void hgemmWmmaNaive(half *A, half *B, half *C, size_t M, size_t N, size_t K) {
    dim3 block(128, 4);
    dim3 grid((M - 1) / (WMMA_M * block.x / WARP_SIZE) + 1, (N - 1) / (WMMA_N * block.y) + 1);

    wmmaNaiveKernel<<<grid, block>>>(A, B, C, M, N, K);
}

2.3 区别

从上述两个naive kernel的代码来看调用CUDA Core和Tensor Core的区别如下:
  • 计算层级:CUDA Core是线程级别,Tensor Core是warp级别
  • 计算维度:CUDA Core是一维逐点计算,Tensor Core是二维逐tile计算
  • 计算依赖:WMMA调用Tensor Core需要借助数据存储类fragment,CUDA Core不需要借助其他

3 底层代码

我们再对上述WMMA naive kernel做进一步探索,看一下它在RTX A6000(sm_86,CUDA 11.3)上对应的PTX和SASS。

3.1 PTX

dump出对应的PTX代码如下,好像不那么简单了。
.visible .entry _Z15wmmaNaiveKernelPK6__halfS1_PS_mmm(
.param .u64 _Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_0,
.param .u64 _Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_1,
.param .u64 _Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_2,
.param .u64 _Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_3,
.param .u64 _Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_4,
.param .u64 _Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_5
)
{
.reg .pred %p<8>;
.reg .b16 %rs<2>;
.reg .f32 %f<2>;
.reg .b32 %r<58>;
.reg .b64 %rd<28>;

ld.param.u64 %rd9, [_Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_0];
ld.param.u64 %rd10, [_Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_1];
ld.param.u64 %rd11, [_Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_2];
ld.param.u64 %rd14, [_Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_3];
ld.param.u64 %rd12, [_Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_4];
ld.param.u64 %rd13, [_Z15wmmaNaiveKernelPK6__halfS1_PS_mmm_param_5];
mov.u32 %r19, %ntid.x;
mov.u32 %r20, %ctaid.x;
mov.u32 %r21, %tid.x;
mad.lo.s32 %r22, %r20, %r19, %r21;
mov.u32 %r23, %ntid.y;
mov.u32 %r24, %ctaid.y;
mov.u32 %r25, %tid.y;
mad.lo.s32 %r26, %r24, %r23, %r25;
mov.f32 %f1, 0f00000000;

	{ cvt.rn.f16.f32 %rs1, %f1;}

	mov.b32 %r50, {%rs1, %rs1};
mul.wide.u32 %rd1, %r26, 16;
shr.u32 %r27, %r22, 1;
and.b32 %r28, %r27, 2147483632;
cvt.u64.u32 %rd2, %r28;
setp.lt.u64 %p2, %rd2, %rd14;
setp.lt.u64 %p3, %rd1, %rd12;
and.pred %p1, %p2, %p3;
setp.eq.s64 %p4, %rd13, 0;
mov.u32 %r51, %r50;
mov.u32 %r52, %r50;
mov.u32 %r53, %r50;
@%p4 bra $L__BB0_5;

mul.lo.s64 %rd3, %rd2, %rd13;
cvt.u32.u64 %r2, %rd13;
mul.lo.s64 %rd4, %rd1, %rd13;
cvta.to.global.u64 %rd5, %rd10;
cvta.to.global.u64 %rd6, %rd9;
mov.u64 %rd27, 0;
not.pred %p5, %p1;
mov.u32 %r51, %r50;
mov.u32 %r52, %r50;
mov.u32 %r53, %r50;

$L__BB0_2:
@%p5 bra $L__BB0_4;

add.s64 %rd16, %rd27, %rd3;
shl.b64 %rd17, %rd16, 1;
add.s64 %rd18, %rd6, %rd17;
wmma.load.a.sync.aligned.row.m16n16k16.global.f16 {%r29, %r30, %r31, %r32, %r33, %r34, %r35, %r36}, [%rd18], %r2;
add.s64 %rd19, %rd27, %rd4;
shl.b64 %rd20, %rd19, 1;
add.s64 %rd21, %rd5, %rd20;
wmma.load.b.sync.aligned.col.m16n16k16.global.f16 {%r37, %r38, %r39, %r40, %r41, %r42, %r43, %r44}, [%rd21], %r2;
wmma.mma.sync.aligned.row.col.m16n16k16.f16.f16 {%r53, %r52, %r51, %r50}, {%r29, %r30, %r31, %r32, %r33, %r34, %r35, %r36}, {%r37, %r38, %r39, %r40, %r41, %r42, %r43, %r44}, {%r53, %r52, %r51, %r50};

$L__BB0_4:
add.s64 %rd27, %rd27, 16;
setp.lt.u64 %p6, %rd27, %rd13;
@%p6 bra $L__BB0_2;

$L__BB0_5:
not.pred %p7, %p1;
@%p7 bra $L__BB0_7;

mul.lo.s64 %rd22, %rd2, %rd12;
add.s64 %rd23, %rd22, %rd1;
cvta.to.global.u64 %rd24, %rd11;
shl.b64 %rd25, %rd23, 1;
add.s64 %rd26, %rd24, %rd25;
cvt.u32.u64 %r45, %rd12;
wmma.store.d.sync.aligned.row.m16n16k16.global.f16 [%rd26], {%r53, %r52, %r51, %r50}, %r45;

$L__BB0_7:
ret;

}

不过我们主要关注WMMA相关的PTX指令,如下所示。可以看到这里正是Nvidia提供的WMMA PTX指令来调用Tensor Core,所以无论是使用WMMA API编程,还是使用WMMA PTX指令编程,底层差别不会太大。

wmma.load.a.sync.aligned.row.m16n16k16.global.f16
wmma.load.b.sync.aligned.col.m16n16k16.global.f16
wmma.mma.sync.aligned.row.col.m16n16k16.f16.f16
wmma.store.d.sync.aligned.row.m16n16k16.global.f16

3.2 SASS

进一步dump出对应的SASS代码,似乎也不简单。
      IMAD.MOV.U32 R1, RZ, RZ, c[0x0][0x28] 
      S2R R0, SR_CTAID.X 
      ISETP.NE.U32.AND P2, PT, RZ, c[0x0][0x188], PT 
      ULDC.64 UR4, c[0x0][0x118] 
      CS2R R8, SRZ 
      S2R R10, SR_CTAID.Y 
      ISETP.NE.AND.EX P2, PT, RZ, c[0x0][0x18c], PT, P2 
      S2R R5, SR_TID.Y 
      S2R R3, SR_TID.X 
      IMAD R10, R10, c[0x0][0x4], R5 
      IMAD R0, R0, c[0x0][0x0], R3 
      IMAD.WIDE.U32 R10, R10, 0x10, RZ 
      CS2R R2, SRZ 
      SHF.R.U32.HI R0, RZ, 0x1, R0 
      ISETP.GE.U32.AND P0, PT, R10, c[0x0][0x180], PT 
      LOP3.LUT R13, R0, 0x7ffffff0, RZ, 0xc0, !PT 
      ISETP.GE.U32.AND.EX P0, PT, R11, c[0x0][0x184], PT, P0 
      ISETP.LT.U32.AND P1, PT, R13, c[0x0][0x178], PT 
      ISETP.LT.U32.AND.EX P0, PT, RZ, c[0x0][0x17c], !P0, P1 
@!P2  BRA 0x7f1eaefc0160 
      BSSY B0, 0x7f1eaefc0160 
      IMAD.MOV.U32 R0, RZ, RZ, RZ 
      CS2R R8, SRZ 
      IMAD.MOV.U32 R15, RZ, RZ, RZ 
      IMAD.MOV.U32 R2, RZ, RZ, RZ 
      BSSY B1, 0x7f1eaefc0100 
@!P0  BRA 0x7f1eaefc00f0 
      S2R R16, SR_LANEID 
      IMAD R17, R11, c[0x0][0x188], RZ 
      IMAD.MOV.U32 R14, RZ, RZ, R0 
      IMAD.MOV.U32 R23, RZ, RZ, c[0x0][0x188] 
      IMAD.WIDE.U32 R6, R10, c[0x0][0x188], R14 
      SHF.R.U32.HI R12, RZ, 0x1, R23 
      IMAD R17, R10, c[0x0][0x18c], R17 
      LEA R21, P2, R6, c[0x0][0x168], 0x1 
      IMAD.WIDE.U32 R4, R13, c[0x0][0x188], R14 
      IMAD.IADD R7, R7, 0x1, R17 
      IMAD.MOV.U32 R17, RZ, RZ, RZ 
      IMAD R5, R13, c[0x0][0x18c], R5 
      LEA.HI.X R7, R6, c[0x0][0x16c], R7, 0x1, P2 
      SHF.R.U32.HI R19, RZ, 0x2, R16 
      LOP3.LUT R16, R16, 0x3, RZ, 0xc0, !PT 
      IMAD.WIDE.U32 R16, R19, R12, R16 
      LEA R19, P1, R4, c[0x0][0x160], 0x1 
      LEA.HI.X R5, R4, c[0x0][0x164], R5, 0x1, P1 
      LEA R18, P1, R16, R19, 0x2 
      LEA R20, P2, R16, R21, 0x2 
      LEA.HI.X R19, R16, R5, R17, 0x2, P1 
      LEA.HI.X R21, R16, R7, R17, 0x2, P2 
      IMAD.WIDE.U32 R16, R23, 0x10, R18 
      LDG.E R4, [R18.64] 
      IMAD.WIDE.U32 R22, R23, 0x10, R20 
      LDG.E R24, [R20.64] 
      LDG.E R25, [R20.64+0x10] 
      LDG.E R6, [R18.64+0x10] 
      LDG.E R5, [R16.64] 
      LDG.E R7, [R16.64+0x10] 
      LDG.E R26, [R22.64] 
      LDG.E R27, [R22.64+0x10] 
      WARPSYNC 0xffffffff 
      HMMA.16816.F16 R8, R4, R24, R8 
      HMMA.16816.F16 R2, R4, R26, R2 
      NOP 
      BSYNC B1 
      IADD3 R0, P1, R0, 0x10, RZ 
      IMAD.X R15, RZ, RZ, R15, P1 
      ISETP.GE.U32.AND P1, PT, R0, c[0x0][0x188], PT 
      ISETP.GE.U32.AND.EX P1, PT, R15, c[0x0][0x18c], PT, P1 
@!P1  BRA 0x7f1eaefbfe90 
      BSYNC B0 
@!P0  EXIT 
      S2R R4, SR_LANEID 
      IMAD.MOV.U32 R15, RZ, RZ, c[0x0][0x180] 
      WARPSYNC 0xffffffff 
      IMAD.WIDE.U32 R10, R13, c[0x0][0x180], R10 
      SHF.R.U32.HI R15, RZ, 0x1, R15 
      IMAD.MOV.U32 R5, RZ, RZ, RZ 
      LEA R7, P0, R10, c[0x0][0x170], 0x1 
      IMAD R11, R13, c[0x0][0x184], R11 
      LEA.HI.X R11, R10, c[0x0][0x174], R11, 0x1, P0 
      SHF.R.U32.HI R0, RZ, 0x2, R4 
      LOP3.LUT R4, R4, 0x3, RZ, 0xc0, !PT 
      IMAD.WIDE.U32 R4, R0, R15, R4 
      LEA R6, P0, R4, R7, 0x2 
      LEA.HI.X R7, R4, R11, R5, 0x2, P0 
      IMAD.WIDE.U32 R4, R15, 0x20, R6 
      STG.E [R6.64], R8 
      STG.E [R4.64], R9 
      STG.E [R6.64+0x10], R2 
      STG.E [R4.64+0x10], R3 
      EXIT 
      BRA 0x7f1eaefc02b0
      NOP
      NOP
      NOP
      NOP
      NOP
      NOP
      NOP
      NOP
      NOP
      NOP
      NOP
      NOP

我们依然主要关注WMMA相关的SASS指令,如下所示。可以发现WMMA161616在底层是通过两个HMMA16816指令实现,同样地,SASS指令也是Nvidia提供的另一种调用Tensor Core的编程方法。

HMMA.16816.F16

Nvidia Tensor Core初探
中提到Nvidia提供了四种调用Tensor Core的编程方法,这里提到了三种,还有一种是MMA PTX指令,其中MMA16816 PTX指令底层实现即是HMMA16816指令,后续会在MMA PTX相关文章中提及。

4 其他

4.1 HGEMM优化

学习WMMA API的目标在于调用Tensor Core优化HGEMM,相比于cublas,WMMA的性能究竟如何?

MASA技术团队联合成都.NET俱乐部,将在成都市举办一场.NET线下技术沙龙,为.NET开发者创造一次交流学习的契机,我们邀请到的几位技术大咖,将会围绕各自的主题向大家分享他们的技术心得。

本场沙龙
名额有限
,以报名
优先为准

时间:2023年4月15日13:30-17:30

地址:四川省成都市武侯区世纪城路198号世纪城新国际会展中心2区3楼(知域空间)【楼下可停车】【地铁一号线世纪城站附近】

流程安排

13:30-14:00 与会人员签到领取礼品

14:00-14:15 主持人开场并介绍到会嘉宾

14:15-14:45 嘉宾谷首道分享“小缓存,大学问,从理论到实践的缓存设计思路”及问答互动

14:50-15:20 嘉宾闫志刚分享“如何用技术改变生活”及问答互动

15:20-15:50 自由交流、茶歇休息、合照

15:50-16:20 嘉宾叶礼飞分享“敏捷开发之道:探索极限编程的实践之路”及问答互动

16:25-16:55 嘉宾周旭龙分享“西门子成都工厂的DevSecOps实践”及问答互动

17:00-17:30 抽奖

报名入口

扫码报名,与社区伙伴
面对面
交流

对广大的.NET开发者来说,每一次的技术交流都是一个提升自我能力、开拓知识面的机会,尤其是参与一场大咖云集的技术沙龙。我们希望通过举办这样的线下沙龙,让更多的.NET开发者了解.NET中国社区,加入.NET开源技术生态,找到自己的组织。

精品礼物

在这次的沙龙中,大家除了会得到超干货的技术分享外,我们也为大家准备了精品礼物,双面帆布袋、笔记本、MASA Stack贴纸、三合一伸缩充电线,以及神秘奖品,希望大家喜欢!

另外,这次的沙龙我们也考虑到外地的.NET开发者和因时间冲突无法参会者,所以我们将在MASA技术团队视频号、MASA技术团队B站同步开启直播,敬请观看。

主办方

协办方

最近比较忙,就少写两句,直接附上源代码,其中的细节点就不再赘述,如有疑问,请留言。


一共就是实现了两个函数,一个用于搜索特征码 (
SearchPattern
),一个用于生成特征码 (
GenerateFunctionSignature)

函数的参数和返回值:

1.
SearchPattern
接收
一个必要参数 hexStr(即要搜索的特征码),一个可选参数 num(最多返回多少个匹配的结果),
返回
一个存放所有符合条件的地址的
list

2.
GenerateFunctionSignature 接收
一个必要参数 addr(即要生成特征码的地址的起始位置),
返回
唯一的特征码
字符串
,例如输入地址:0x12345678,返回字符串 48 8B F2 4C 8B F1 E8 ?? ?? ?? ??

·
·
·

以下是
SearchPattern
的实现:

1 importida_bytes2 importida_ida3 
4 def SearchPattern(hexStr, num =0):5     '''
6 hexStr: 输入的特征码字符串,例:41 80 BE ?? ?? ?? ?? 0F 84 ?? ??7 num:    搜索到多少个结果时返回,为 0 时搜索全部匹配的地址8     '''
9     
10     #生成掩码
11     bMask = hexStr.replace('00', '01')12     bMask = bMask.replace('??', '00')13     bMask =bytes.fromhex(bMask)14     #print(bMask)
15     
16     #生成模式码
17     bPattern = hexStr.replace('??', '00')18     bPattern =bytes.fromhex(bPattern)19     #print(bPattern)
20 
21 
22     results =[]23     ea =ida_ida.inf_get_min_ea()24     
25     whileTrue:26         ea =ida_bytes.bin_search(27             ea + 1,28 ida_ida.inf_get_max_ea(),29 bPattern,30 bMask,31             1,32             ida_bytes.BIN_SEARCH_FORWARD| ida_bytes.BIN_SEARCH_NOBREAK|ida_bytes.BIN_SEARCH_NOSHOW)33         if ea ==ida_idaapi.BADADDR:34             break
35         else:36             #这里可以稍作修改,让返回值变为当前函数的首地址
37 results.append(hex(ea))38             if num != 0 and len(results) >=num:39                 break
40     #print("find {} result".format(len(results)))
41     
42     returnresults43 
44 
45 ## 测试
46 print(SearchPattern("48 8B C4 48 89 50 ??", 3))

·
·
·

以下是
GenerateFunctionSignature
的实现:

1 importida_bytes, ida_ua2 
3 defGenerateFunctionSignature(funcBase):4     #初始化要用到的变量
5     instruction =ida_ua.insn_t()6     offset =07     signature = ""
8     
9     #最多分析 100 条指令
10     for count in range(1, 101):11         ida_ua.decode_insn(instruction, funcBase +offset)12         offset +=instruction.size13         patternTemp = [0 for i inrange(0, instruction.size)]14         
15         
16         #遍历当前指令的全部操作数,并将部分操作数的机器码置为 ??
17         for op ininstruction.ops:18             if op.type ==o_void:19                 continue
20             #模糊位 ?? 的匹配条件:
21             #and op.type != ida_ua.o_phrase and op.type != ida_ua.o_displ
22             elif (op.type !=ida_ua.o_reg):23                 for index inrange(op.offb, instruction.size):24                     patternTemp[index] = "??"
25 
26 
27         #读入除模糊位之外的机器码
28         for index inrange(0, instruction.size):29             if patternTemp[index] == "??":30                 pass
31             else:32                 byteStr = format(ida_bytes.get_byte(instruction.ea + index), '02X')33                 patternTemp[index] =byteStr34             signature = signature + ' ' +patternTemp[index]35 
36         
37         #每分析 3 条指令,判断一次当前的 signature 是否唯一
38         if count%3 ==0:39             if len(SearchPattern(signature, 2)) == 1:40                 print('unique signature')41                 returnsignature42             else:43                 #print('not unique signature')
44                 continue
45             
46     returnNone47     
48 
49 ## 测试
50 print("pattern =", GenerateFunctionSignature(0x320E21F))