MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / EmitTile

Method EmitTile

tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc:2040–2127  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

2038}
2039
2040void IrEmitterUnnested::EmitTile(
2041 const KernelMappingScheme& mapping_scheme,
2042 const IrArray::Index& tile_origin_index, const string& loop_name,
2043 KernelSupportLibrary* ksl, const ThreadIdInfo& thread_id_info,
2044 llvm::Value* tile_height, llvm::Value* tile_width,
2045 const IrEmitterUnnested::EmitElementFunction& emit_elem_function) {
2046 llvm::Type* index_ty = tile_width->getType();
2047 auto constant = [&](int64 val) {
2048 return llvm::ConstantInt::get(index_ty, val);
2049 };
2050 int64 num_threads_x = mapping_scheme.GetNumThreadsX();
2051 llvm::Value* num_threads_y = constant(mapping_scheme.GetNumThreadsY());
2052 int64 tile_size_x = mapping_scheme.GetTileSizeX();
2053
2054 int64 x_num_steps = tile_size_x / num_threads_x;
2055 llvm::Value* start_offset_x = GetStartOffsetX(
2056 mapping_scheme, thread_id_info.thread_id_x, index_ty, &b_);
2057
2058 // Using dilated mapping scheme, each thread steps with a stride of number
2059 // of threads.
2060 // Otherwise, the stride is one, but we multiply each offset by the limit of
2061 // number of steps which can be made.
2062 int64 step_x =
2063 mapping_scheme.GetIndexingOrder() == kLinearIndexingX ? 1 : num_threads_x;
2064 int64 vector_size = mapping_scheme.GetVectorSize();
2065
2066 IrArray::Index source_idx =
2067 tile_origin_index.AddOffsetToDim(start_offset_x, kDimX, &b_);
2068
2069 auto ceil_of_ratio = [&](llvm::Value* a, llvm::Value* b) {
2070 return b_.CreateUDiv(b_.CreateAdd(b_.CreateAdd(a, b), constant(-1)), b);
2071 };
2072
2073 // True iff all threads always execute all instructions in the tiling
2074 // dimension X.
2075 bool x_tile_fits =
2076 mapping_scheme.GetDimsInElems()[kDimX] % tile_size_x == 0 &&
2077 mapping_scheme.GetRowContiguous();
2078
2079 // The outer loop below is simply doing:
2080 //
2081 // for (int y_loc=thread_id_y; y_loc<tile_height; y_loc+=num_threads_y)
2082 //
2083 //
2084 // However, in order to avoid an LLVM optimization triggering the ptxas bug,
2085 // we write this loop in a convoluted way:
2086 //
2087 // y_bound = ceil_of_ratio(tile_height - thread_id_y, num_threads_y)
2088 // for (int y_indvar=0; y_indvar<y_bound; y_indvar+=1)
2089 // y_loc = thread_id_y + y_indvar * num_threads_y
2090 //
2091 // TODO(cheshire): Once ptxas is fixed and TF switches to it, remove the
2092 // workaround.
2093 ksl->For(
2094 loop_name + "_y_in_tile",
2095 /*start=*/constant(0),
2096 /*end=*/
2097 ceil_of_ratio(b_.CreateSub(tile_height, thread_id_info.thread_id_y),

Callers

nothing calls this directly

Calls 15

GetStartOffsetXFunction · 0.85
ceil_of_ratioFunction · 0.85
UnrollInnerTileLoopFunction · 0.85
GetNumThreadsXMethod · 0.80
GetNumThreadsYMethod · 0.80
GetTileSizeXMethod · 0.80
GetIndexingOrderMethod · 0.80
GetVectorSizeMethod · 0.80
AddOffsetToDimMethod · 0.80
GetDimsInElemsMethod · 0.80
GetRowContiguousMethod · 0.80
ForMethod · 0.80

Tested by

no test coverage detected