| 2038 | } |
| 2039 | |
| 2040 | void 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), |
nothing calls this directly
no test coverage detected