MLIR笔记（2）

3. LVM有趣的代码

3.1. dyn_cast()与cast()

C++支持类型间的自动转换（如operator =声明的转换），但在转换的调用链里自动转换只能调用一次，这固然是避免给编译器带来过分的复杂性，但更重要的是允许自动转换接力调用几乎很难避免出现递归调用，而且调用链过长会很快失去控制，给人带来意想不到的结果。但是，C++原生的类型转换系统对于LLVM/MLIR来说局限性太大，因此，LLVM打造了自己的类型转换系统，它仿造C++的惯例，给出了自己的类型转换函数：cast、dyn_cast。显然，它们分别对应其他cast与dynamic_cast。它们与C++版本的最大不同在于：cast、dyn_cast都支持无限次自动转换（这在用户眼皮底下发生，但用户需要提供一些指定的方法）。在这个体系中，dyn_cast最有代表性，llvm里一共定义了3个版本，它们的区别在于函数参数类型：

331 template <class X, class Y>

332 LLVM_NODISCARD inline std::enable_if_t<

333 !is_simple_type<Y>::value, typename cast_retty<X, const Y>::ret_type>

334 dyn_cast(const Y &Val) {

335 return isa<X>(Val) ? cast<X>(Val) : nullptr;

336 }

337

338 template <class X, class Y>

339 LLVM_NODISCARD inline typename cast_retty<X, Y>::ret_type dyn_cast(Y &Val) {

340 return isa<X>(Val) ? cast<X>(Val) : nullptr;

341 }

342

343 template <class X, class Y>

344 LLVM_NODISCARD inline typename cast_retty<X, Y *>::ret_type dyn_cast(Y *Val) {

345 return isa<X>(Val) ? cast<X>(Val) : nullptr;

346 }

在这些函数里用到isa()与cast()，isa()是这样的：它检查参数的类型是否与模板参数里给出的类型一致。

141 template <class X, class Y> LLVM_NODISCARD inline bool isa(const Y &Val) { // 最终调用的函数

142 return isa_impl_wrap<X, const Y,

143 typename simplify_type<const Y>::SimpleType>::doit(Val);

144 }

145

146 template <typename First, typename Second, typename... Rest, typename Y>

147 LLVM_NODISCARD inline bool isa(const Y &Val) {

148 return isa<First>(Val) || isa<Second, Rest...>(Val);

149 }

实际上，isa()支持对某个类型列表，检查指定类型是否与其中一个类型一致。而两个特定类型比较的实现是由isa_impl_wrap这个类提供的：

116 template<typename To, typename From, typename SimpleFrom>

117 struct isa_impl_wrap {

118 // When From != SimplifiedType, we can simplify the type some more by using

119 // the simplify_type template.

120 static bool doit(const From &Val) {

121 return isa_impl_wrap<To, SimpleFrom,

122 typename simplify_type<SimpleFrom>::SimpleType>::doit(

123 simplify_type<const From>::getSimplifiedValue(Val));

124 }

125 };

126

127 template<typename To, typename FromTy>

128 struct isa_impl_wrap<To, FromTy, FromTy> {

129 // When From == SimpleType, we are as simple as we are going to get.

130 static bool doit(const FromTy &Val) {

131 return isa_impl_cl<To,FromTy>::doit(Val);

132 }

133 };

为了支持无限次自动转换，这里指出了3个类型：To、From、SimpleFrom。其中，被转换值从From类型转换到SimpleFrom类型，这个SimpleFrom类型还可能进一步转换为另一个SimpleFrom类型，直到From与SimpleFrom类型一致为止，这意味着From不能再“简化了”。这时我们使用128行的isa_impl_wrap特化定义。

首先，我们先看一下完成从From到SimpleFrom的simplify_type类。注意，如果存在llvm所不知道的From到SimpleFrom的转换，必须提供这个目的的simplify_type特化版本。如果不提供，将使用simplify_type的非特化版本（llvm对自己的数据结构定义了若干simplify_type特化版本）：

33 template<typename From> struct simplify_type {

34 using SimpleType = From; // The real type this represents...

36 // An accessor to get the real value...

37 static SimpleType &getSimplifiedValue(From &Val) { return Val; }

38 };

40 template<typename From> struct simplify_type<const From> {

41 using NonConstSimpleType = typename simplify_type<From>::SimpleType;

42 using SimpleType =

43 typename add_const_past_pointer<NonConstSimpleType>::type;

44 using RetType =

45 typename add_lvalue_reference_if_not_pointer<SimpleType>::type;

47 static RetType getSimplifiedValue(const From& Val) {

48 return simplify_type<From>::getSimplifiedValue(const_cast<From&>(Val));

49 }

50 };

可以看到，非特化版本的simplify_type通过getSimplifiedValue()给出的SimpleFrom就是值自己。

在simplify_type再也给不出“简化”时，isa_impl_wrap提供最后的判断，同样，它有许多特化版本，我们只看llvm的非定制版本好了：

55 template <typename To, typename From, typename Enabler = void>

56 struct isa_impl {

57 static inline bool doit(const From &Val) {

58 return To::classof(&Val); // <-- 最终的裁决者

59 }

60 };

62 /// Always allow upcasts, and perform no dynamic check for them.

63 template <typename To, typename From>

64 struct isa_impl<To, From, std::enable_if_t<std::is_base_of<To, From>::value>> {

65 static inline bool doit(const From &) { return true; }

66 };

68 template <typename To, typename From> struct isa_impl_cl {

69 static inline bool doit(const From &Val) {

70 return isa_impl<To, From>::doit(Val);

71 }

72 };

这里的调用关系始于68行的isa_impl_cl()，它所调用的isa_impl()也需要各个类提供自己的特化版本来实现特定的类型转换。在不提供特化版本时，如果两个类有继承关系，那么就会使用64行llvm的特化版本，无条件放行；否则，由上面58行To的classof()方法来判定它们是否一致。

在dyn_cast()中，当isa()认定两个类型是一致时，cast()就可以执行具体的转换了。类似的，cast()有4个重载版本：

251 template <class X, class Y>

252 inline std::enable_if_t<!is_simple_type<Y>::value,

253 typename cast_retty<X, const Y>::ret_type>

254 cast(const Y &Val) {

255 assert(isa<X>(Val) && "cast<Ty>() argument of incompatible type!");

256 return cast_convert_val<

257 X, const Y, typename simplify_type<const Y>::SimpleType>::doit(Val);

258 }

259

260 template <class X, class Y>

261 inline typename cast_retty<X, Y>::ret_type cast(Y &Val) {

262 assert(isa<X>(Val) && "cast<Ty>() argument of incompatible type!");

263 return cast_convert_val<X, Y,

264 typename simplify_type<Y>::SimpleType>::doit(Val);

265 }

266

267 template <class X, class Y>

268 inline typename cast_retty<X, Y *>::ret_type cast(Y *Val) {

269 assert(isa<X>(Val) && "cast<Ty>() argument of incompatible type!");

270 return cast_convert_val<X, Y*,

271 typename simplify_type<Y*>::SimpleType>::doit(Val);

272 }

273

274 template <class X, class Y>

275 inline typename cast_retty<X, std::unique_ptr<Y>>::ret_type

276 cast(std::unique_ptr<Y> &&Val) {

277 assert(isa<X>(Val.get()) && "cast<Ty>() argument of incompatible type!");

278 using ret_type = typename cast_retty<X, std::unique_ptr<Y>>::ret_type;

279 return ret_type(

280 cast_convert_val<X, Y *, typename simplify_type<Y *>::SimpleType>::doit(

281 Val.release()));

282 }

这里，首先由cast_retty::ret_type决定cast()的返回值类型：

212 template<class To, class From>

213 struct cast_retty {

214 using ret_type = typename cast_retty_wrap<

215 To, From, typename simplify_type<From>::SimpleType>::ret_type;

216 };

其中，cast_retty_wrap是这样的定义：

198 template<class To, class From, class SimpleFrom>

199 struct cast_retty_wrap {

200 // When the simplified type and the from type are not the same, use the type

201· // simplifier to reduce the type, then reuse cast_retty_impl to get the

202 // resultant type.

203 using ret_type = typename cast_retty<To, SimpleFrom>::ret_type;

204 };

205

206 template<class To, class FromTy>

207 struct cast_retty_wrap<To, FromTy, FromTy> {

208 // When the simplified type is equal to the from type, use it directly.

209 using ret_type = typename cast_retty_impl<To,FromTy>::ret_type;

210 };

同样，第二个版本的cast_retty_wrap是最终调用的版本，它使用的cast_retty_impl同样有多个特化版本（用户还可以自己定制），为了处理常量和指针等情况，我们看一下非特化版本：

169 template<class To, class From> struct cast_retty_impl {

170 using ret_type = To &; // Normal case, return Ty&

171 };

不出意料，cast()的返回值类型就是To类型（指针情况下使用引用类型）。确定了返回值类型后，cast()使用cast_convert_val::doit()执行转换：

221 template<class To, class From, class SimpleFrom> struct cast_convert_val {

222 // This is not a simple type, use the template to simplify it...

223 static typename cast_retty<To, From>::ret_type doit(From &Val) {

224 return cast_convert_val<To, SimpleFrom,

225 typename simplify_type<SimpleFrom>::SimpleType>::doit(

226 simplify_type<From>::getSimplifiedValue(Val));

227 }

228 };

229

230 template<class To, class FromTy> struct cast_convert_val<To,FromTy,FromTy> {

231 // This _is_ a simple type, just cast it.

232 static typename cast_retty<To, FromTy>::ret_type doit(const FromTy &Val) {

233 typename cast_retty<To, FromTy>::ret_type Res2

234 = (typename cast_retty<To, FromTy>::ret_type)const_cast<FromTy&>(Val);

235 return Res2;

236 }

237 };

对于我们最终调用的第二个版本来说，所谓的转换就是一个简单的C形式的强制转换，不过，由于前面的一系列检查，这个转换是安全的。

3.2. TrailingObjects

不同于C语言里比较原始的变长类型（即union），在LLVM里对变长类型有更好的支持。在需要使用变长类型时，只需要将llvm::TrailingObjects作为基类。这个类型定义的开头几行是这样的：

228 template <typename BaseTy, typename... TrailingTys>

229 class TrailingObjects : private trailing_objects_internal::TrailingObjectsImpl<

230 trailing_objects_internal::AlignmentCalcHelper<

231 TrailingTys...>::Alignment,

232 BaseTy, TrailingObjects<BaseTy, TrailingTys...>,

233 BaseTy, TrailingTys...> {

234

235 template <int A, typename B, typename T, typename P, typename... M>

236 friend class trailing_objects_internal::TrailingObjectsImpl;

237

238 template <typename... Tys> class Foo {};

239

240 typedef trailing_objects_internal::TrailingObjectsImpl<

241 trailing_objects_internal::AlignmentCalcHelper<TrailingTys...>::Alignment,

242 BaseTy, TrailingObjects<BaseTy, TrailingTys...>, BaseTy, TrailingTys...>

243 ParentType;

244 using TrailingObjectsBase = trailing_objects_internal::TrailingObjectsBase;

245

246 using ParentType::getTrailingObjectsImpl;

这里需要注意232行的TrailingObjects<BaseTy, TrailingTys...>，这其实是当前TrailingObjects的类型，它用途我们会在下面看到。

变长类型对象的大小由totalSizeToAlloc()方法给出，在创建对象时，它用于确定分配资源的数量。

340 template <typename... Tys>

341 static constexpr std::enable_if_t<

342 std::is_same<Foo<TrailingTys...>, Foo<Tys...>>::value, size_t>

343 totalSizeToAlloc(typename trailing_objects_internal::ExtractSecondType<

344 TrailingTys, size_t>::type... Counts) {

345 return sizeof(BaseTy) + ParentType::additionalSizeToAllocImpl(0, Counts...);

346 }

在上面243行定义的ParentType类型不出意料有两个特化版本（注意，ParentType也是当前TrailingObjects的父类）。第一个是全特化，它是特化产生的派生树的叶子类型，模板展开到此结束，开始回溯：

206 template <int Align, typename BaseTy, typename TopTrailingObj, typename PrevTy>

207 class TrailingObjectsImpl<Align, BaseTy, TopTrailingObj, PrevTy>

208 : public TrailingObjectsAligner<Align> {

209 protected:

210 // This is a dummy method, only here so the "using" doesn't fail --

211 // it will never be called, because this function recurses backwards

212 // up the inheritance chain to subclasses.

213 static void getTrailingObjectsImpl();

214

215 static constexpr size_t additionalSizeToAllocImpl(size_t SizeSoFar) {

216 return SizeSoFar;

217 }

218

219 template <bool CheckAlignment> static void verifyTrailingObjectsAlignment() {}

220 };

另一个则是偏特化，是整个模板展开过程的核心，它开头几行是这样的：

130 template <int Align, typename BaseTy, typename TopTrailingObj, typename PrevTy,

131 typename NextTy, typename... MoreTys>

132 class TrailingObjectsImpl<Align, BaseTy, TopTrailingObj, PrevTy, NextTy,

133 MoreTys...>

134 : public TrailingObjectsImpl<Align, BaseTy, TopTrailingObj, NextTy,

135 MoreTys...> {

136

137 typedef TrailingObjectsImpl<Align, BaseTy, TopTrailingObj, NextTy, MoreTys...>

138 ParentType;

139

140 struct RequiresRealignment {

141 static const bool value = alignof(PrevTy) < alignof(NextTy);

142 };

143

144 static constexpr bool requiresRealignment() {

145 return RequiresRealignment::value;

146 }

这里需要注意，对这两个版本来说，与它们相关的类型是模板参数PrevTy。另外，BaseTy都是变长类型里的第一个类型。我们关注的additionalSizeToAllocImpl()方法是：

193 static constexpr size_t additionalSizeToAllocImpl(

194 size_t SizeSoFar, size_t Count1,

195 typename ExtractSecondType<MoreTys, size_t>::type... MoreCounts) {

196 return ParentType::additionalSizeToAllocImpl( // <-- 调用父类的additionalSizeToAllocImpl()

197 (requiresRealignment() ? llvm::alignTo<alignof(NextTy)>(SizeSoFar)

198 : SizeSoFar) +

199 sizeof(NextTy) * Count1, // <-- 加上我们管的这部分

200 MoreCounts...); // <-- 剩下的交给父类处理

201 }

除了确定对象大小，访问指定对象的能力也是不可缺的，这由getTrailingObjects()提供，它有两个版本，一个返回const指针，另一个返回非const指针，除此之外实现上没有其他区别：

297 template <typename T> T *getTrailingObjects() {

298 verifyTrailingObjectsAssertions();

299 // Forwards to an impl function with overloads, since member

300 // function templates can't be specialized.

301 return this->getTrailingObjectsImpl(

302 static_cast<BaseTy *>(this), TrailingObjectsBase::OverloadToken<T>());

303 }

上面的方法没有特化版本，模板参数T在派生体系里选择最恰当的重载版本。我们以Operation为例：

29 class alignas(8) Operation final

29 : public llvm::ilist_node_with_parent<Operation, Block>,

30 private llvm::TrailingObjects<Operation, BlockOperand, Region,

31 detail::OperandStorage> {

30 ~ 31行将展开为这样的继承树，最底下是叶子节点，自底向上继承。与特定TrailingObjectsImpl相关的类型绿色高亮：

TrailingObjectsImpl<…, Operation , TrailingObjects, Operation, BlockOperand…>

TrailingObjectsImpl<…, Operation , TrailingObjects, BlockOperand, Region…>

TrailingObjectsImpl<…, Operation , TrailingObjects, Region, detail::OperandStorage>

TrailingObjectsImpl<…, Operation , TrailingObjects, detail::OperandStorage>

下面的getTrailingObjectsImpl()是由TrailingObjectsImpl定义的。注意返回类型NextTy是模板参数，是与当前TrailingObjectsImpl相关类型（PrevTy）的下一个类型。也就是说，下一个对象的位置由前一个对象返回，这也是合理的，因为只有找到前一个对象才能知道它后面的对象在哪里。

159 static NextTy *

160 getTrailingObjectsImpl(BaseTy *Obj,

161 TrailingObjectsBase::OverloadToken<NextTy>) {

162 auto *Ptr = TopTrailingObj::getTrailingObjectsImpl(

163 Obj, TrailingObjectsBase::OverloadToken<PrevTy>()) +

164 TopTrailingObj::callNumTrailingObjects(

165 Obj, TrailingObjectsBase::OverloadToken<PrevTy>());

166

167 if (requiresRealignment())

168 return reinterpret_cast<NextTy *>(alignAddr(Ptr, Align::Of<NextTy>()));

169 else

170 return reinterpret_cast<NextTy *>(Ptr);