PR Category

CINN

PR Types

New features

Description

一、is_index 设置
is_index 是判断 Expr 是否为 IndexExpr 的唯一标识，可以直接通过它判别，无需额外地递归判断 IRNode 类型。
本 PR 在关键节点上设置 is_index flag，并保证 is_index flag 在后端 IR 下标相关 IRNode 中自动蔓延。
共分为 4 类具体如下：

1. 直接设置 is_index 类（Expr 经过该类 Op 自动赋值 is_index）：

|- Buffer:: shape, strides, elem_offset
|- BufferRange:: ranges
|- Dim:: dim_expr
|- For:: min, extent, var
|- ScheduleBlock:: iter_vars
|- ScheduleBlockRealize:: iter_values
|- Store:: indices
|- Alloc:: extents
|- Load:: indices
|- Reduce:: reduce_axis
|- PlaceholderOp:: shape
|- ComputeOp:: axis, shape, reduce_axis

2. 需要额外条件才可设置 is_index 类（需要判断 Expr 类型，只有 IndexExpr 类型才设置 is_index）：

|- IfThenElse:: condition
|- Select:: condition

3. 基础运算类别（+、-、*、/....，当操作数为 IndexExpr 时，设置 is_index）：

|- Add、Sub、Mul、Div、Mod::make
|- Cast::make
|- Let::make

4. 特殊类别（替换、DimExpr -> Expr）：

|- MappingVarToExprMutator:: Visit
|- DimExprToIrExprVisitor
|- GenDefaultAxis
|- AutoSimplify
|- IRCopy

二、其余改动

1. IndexExpr 拓展 min、max、cast、load 类型
2. 新增 min、max 化简函数
3. Optim::Simplify 中加入 SimplifySelectMutator 化简 condition 为常量的 select，例如：

select(0<1, var_0[i], var_1[i]) ===> var_0[i]

4. cast::Make 时进行简化，当 cast 对象可直接改变类型时，不显式引入 cast op，例如：

cast(int32, IntImm) ===> IntImm with type int32

5. IRcompare 中当比较对象对 IndexExpr 时升格为 Index 判等

三、临时改动

1. IRCopy 中输出对象是否设置为 IndexExpr，除了判断输入是否为 IndexExpr 外，还需判断 Expr 类型，因为 AutoSimplify 中存在 Expr 的 inplace 修改，导致 Div 转换为 Frac，Frac 不能设置为 IndexExpr
2. AutoSimplify 中跳过 Load 内部化简，因为 Load::indices 已经被 IndexExpr 化简

Pcard-67164