[Refactor] Use switch-case instead of "static std::map" in taichi/lang_util.cpp

[yuanming-hu]

Where to refactor
taichi/lang_util.cpp has quite a few functions (e.g., data_type_name, data_type_format) that uses a static std::map to store input-output mapping.

taichi/taichi/lang_util.cpp

Lines 67 to 88 in a58a9af

	std::string data_type_name(DataType t) {
	static std::map<DataType, std::string> type_names;
	if (type_names.empty()) {
	#define REGISTER_DATA_TYPE(i, j) type_names[DataType::i] = #j;
	REGISTER_DATA_TYPE(f16, float16);
	REGISTER_DATA_TYPE(f32, float32);
	REGISTER_DATA_TYPE(f64, float64);
	REGISTER_DATA_TYPE(u1, int1);
	REGISTER_DATA_TYPE(i8, int8);
	REGISTER_DATA_TYPE(i16, int16);
	REGISTER_DATA_TYPE(i32, int32);
	REGISTER_DATA_TYPE(i64, int64);
	REGISTER_DATA_TYPE(u8, uint8);
	REGISTER_DATA_TYPE(u16, uint16);
	REGISTER_DATA_TYPE(u32, uint32);
	REGISTER_DATA_TYPE(u64, uint64);
	REGISTER_DATA_TYPE(gen, generic);
	REGISTER_DATA_TYPE(unknown, unknown);
	#undef REGISTER_DATA_TYPE
	}
	return type_names[t];
	}

What's wrong with the current design
When multiple threads are calling the function with the static map uninitialized, there will be parallel writes to the "std::map". This may cause a segmentation fault.

How to refactor

• Simply use C-syntax switch-case instead of "std::map"
• Still make use of macros to minimize changes

Example
The code above should be changed into something like

std::string data_type_name(DataType t) {
  switch (t) {
    
#define REGISTER_DATA_TYPE(i, j) \
  case DataType::i:              \
    return #j;

    REGISTER_DATA_TYPE(f16, float16);
    REGISTER_DATA_TYPE(f32, float32);
    REGISTER_DATA_TYPE(f64, float64);
    REGISTER_DATA_TYPE(u1, int1);
    REGISTER_DATA_TYPE(i8, int8);
    REGISTER_DATA_TYPE(i16, int16);
    REGISTER_DATA_TYPE(i32, int32);
    REGISTER_DATA_TYPE(i64, int64);
    REGISTER_DATA_TYPE(u8, uint8);
    REGISTER_DATA_TYPE(u16, uint16);
    REGISTER_DATA_TYPE(u32, uint32);
    REGISTER_DATA_TYPE(u64, uint64);
    REGISTER_DATA_TYPE(gen, generic);
    REGISTER_DATA_TYPE(unknown, unknown);
#undef REGISTER_DATA_TYPE
    default:
      TI_NOT_IMPLEMENTED
  }
}

Additional comments

This is a good starting point for newcomers. Please leave a note if you have contributed < 3 PRs and to take this :-)
Open a draft PR if you want early feedback. It is suggested to ask for a review from me after the first function is refactored before apply your modification to all functions in that file.

[aryansoman]

Hello! I've left one PR so far, and I'd be interested in taking this.

[yuanming-hu]

A lot of thanks to @aryansoman for refactoring most of the functions! Now we are almost there, the only function left is

taichi/taichi/lang_util.cpp

Lines 306 to 336 in c3f1a5d

	DataType promoted_type(DataType a, DataType b) {
	std::map<std::pair<DataType, DataType>, DataType> mapping;
	if (mapping.empty()) {
	#define TRY_SECOND(x, y) \
	mapping[std::make_pair(get_data_type<x>(), get_data_type<y>())] = \
	get_data_type<decltype(std::declval<x>() + std::declval<y>())>();
	#define TRY_FIRST(x) \
	TRY_SECOND(x, float32); \
	TRY_SECOND(x, float64); \
	TRY_SECOND(x, int8); \
	TRY_SECOND(x, int16); \
	TRY_SECOND(x, int32); \
	TRY_SECOND(x, int64); \
	TRY_SECOND(x, uint8); \
	TRY_SECOND(x, uint16); \
	TRY_SECOND(x, uint32); \
	TRY_SECOND(x, uint64);
	TRY_FIRST(float32);
	TRY_FIRST(float64);
	TRY_FIRST(int8);
	TRY_FIRST(int16);
	TRY_FIRST(int32);
	TRY_FIRST(int64);
	TRY_FIRST(uint8);
	TRY_FIRST(uint16);
	TRY_FIRST(uint32);
	TRY_FIRST(uint64);
	}
	return mapping[std::make_pair(a, b)];
	}

Suggested solution:

Let's not use switch case in this special case.

Note that std::map<std::pair<DataType, DataType>, DataType> is safe for parallel reading but not parallel writing. We are good as long as we initialize a class TypePromotionMapping, which contains the std::map and initializes everything before we do any query. This can be done if we create a global variable TypePromotionMapping type_promotion_mapping so that its initializer is always called before promoted_type (which then calls something like type_promotion_mapping::query(DataType, DataType).

[Hanke98]

Hi, I think it is not a difficult work, let me take it! :)

[yuanming-hu]

This is fully done. A lot of thanks to @Hanke98 and @aryansoman!