diff --git a/pytorch_blade/torch_blade/dynamo/__init__.py b/pytorch_blade/torch_blade/dynamo/__init__.py index 6f7a58bacd0..ef484a13b24 100644 --- a/pytorch_blade/torch_blade/dynamo/__init__.py +++ b/pytorch_blade/torch_blade/dynamo/__init__.py @@ -67,7 +67,6 @@ def _disc_compile(fx_g: fx.GraphModule, inps, use_ts=False, is_training=True) -> v = v.type new_kwargs[k] = v node.kwargs = new_kwargs - fx_g.graph.lint() fx_g.recompile() f = torch.jit.script(fx_g) diff --git a/tao_compiler/mlir/custom_ops/transpose_impl.cc b/tao_compiler/mlir/custom_ops/transpose_impl.cc index fb275a94f73..2f160d967a3 100644 --- a/tao_compiler/mlir/custom_ops/transpose_impl.cc +++ b/tao_compiler/mlir/custom_ops/transpose_impl.cc @@ -68,6 +68,7 @@ TAO_RAL_API("ral_transpose", "gpu", ral_transpose); TAO_RAL_API("ral_transpose", "gpu", ral_transpose); TAO_RAL_API("ral_transpose", "gpu", ral_transpose); TAO_RAL_API("ral_transpose", "gpu", ral_transpose); +TAO_RAL_API("ral_transpose", "gpu", ral_transpose); #endif } // namespace ral diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc index e8f0eb06ae3..1646fb8078d 100644 --- a/tao_compiler/mlir/disc/disc_compiler.cc +++ b/tao_compiler/mlir/disc/disc_compiler.cc @@ -601,7 +601,10 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) { pm.addNestedPass(createCSEPass()); pm.addNestedPass( createCanonicalizerPass(cano_rewrite_config, disablePatterns)); - pm.addNestedPass(disc_ral::createDiscMemRefCSEPass()); + // TODO(yancey): enable this pass after fix WAW issue in scalar reduction + // codegen template + if (!gpu_enabled) + pm.addNestedPass(disc_ral::createDiscMemRefCSEPass()); // convert linearizeOp/delinearizeOp to std dialect. pm.addNestedPass(disc_ral::createDiscConvertShapeToStandardPass()); pm.addNestedPass( diff --git a/tao_compiler/mlir/disc/tests/mlir_feature_test.cc b/tao_compiler/mlir/disc/tests/mlir_feature_test.cc index f38165cde1c..cb42088cb71 100644 --- a/tao_compiler/mlir/disc/tests/mlir_feature_test.cc +++ b/tao_compiler/mlir/disc/tests/mlir_feature_test.cc @@ -238,7 +238,6 @@ void addBoolFlags(EnvSettings& envSettings, const std::string& key) { } else { size_t original_size = envSettings.size(); for (int i = 0; i < original_size; ++i) { - envSettings[i][key].first = "false"; envSettings.push_back(envSettings[i]); envSettings[i][key].first = "true"; } diff --git a/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc b/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc old mode 100755 new mode 100644 index a3de94b8677..78e565edeac --- a/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc +++ b/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc @@ -489,6 +489,10 @@ struct TransposeConverter : public OpRewritePattern { LogicalResult matchAndRewrite(lmhlo::TransposeOp op, PatternRewriter& rewriter) const override { + if (auto fusion_op = + op.getOperation()->getParentOfType()) { + return failure(); + } auto permutation = op.getPermutation().getValues(); int rank = permutation.size(); if (rank != 2 && rank != 3) return failure(); diff --git a/tao_compiler/mlir/disc/transforms/element_type_converter.cc b/tao_compiler/mlir/disc/transforms/element_type_converter.cc index bee2b667a0e..d8c0d70d7da 100644 --- a/tao_compiler/mlir/disc/transforms/element_type_converter.cc +++ b/tao_compiler/mlir/disc/transforms/element_type_converter.cc @@ -117,8 +117,8 @@ struct ConvertReduceOpWithSmallWidthIntType int rank = ty.getRank(); int ndims_to_reduce = static_cast(dims_to_reduce.size()); - if (rank != 2 || ndims_to_reduce != 1) { - // Suppose that there are only rank-2 row/colunm reduction after + if (rank != 2) { + // Suppose that there are only rank-2 row/colunm/scalar reduction after // `canonicalize-reduction` pass. return failure(); } diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc index 2475f1d94fb..cf0cbeb314f 100644 --- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc +++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc @@ -175,6 +175,8 @@ StringRef fusionTypeToString(FusionType ft) { return "kRowReduction"; case FusionType::kColReduction: return "kColReduction"; + case FusionType::kScalarReduction: + return "kScalarReduction"; case FusionType::kInput: return "kInput"; case FusionType::kStitch: @@ -205,6 +207,8 @@ FusionType fusionTypeFromString(StringRef ft) { return FusionType::kRowReduction; } else if (ft == "kColReduction") { return FusionType::kColReduction; + } else if (ft == "kScalarReduction") { + return FusionType::kScalarReduction; } else if (ft == "kInput") { return FusionType::kInput; } else if (ft == "kStitch") { @@ -479,6 +483,16 @@ bool isRowReduction(Operation* op) { return true; } +bool isRank2ScalarReduction(Operation* op) { + auto reduce_op = dyn_cast(op); + if (!reduce_op || reduce_op.getDimensions().getNumElements() != 2) + return false; + if (auto ty = op->getOperand(2).getType().dyn_cast()) { + return ty.getRank() == 0; + } + return false; +} + // Returns true if this op is a rank-2 column reduction. bool isRank2ColReduction(Operation* op) { auto reduce_op = dyn_cast(op); @@ -554,10 +568,11 @@ bool initFusionPatternBase(ShapeAnalysis& shapeAnalysis, inferredFusionType = FusionType::kRowReduction; inferredDominantOp = op; } else if (isRank2ColReduction(op)) { - if (inferredFusionType != FusionType::kRowReduction) { - inferredFusionType = FusionType::kColReduction; - inferredDominantOp = op; - } + inferredFusionType = FusionType::kColReduction; + inferredDominantOp = op; + } else if (isRank2ScalarReduction(op)) { + inferredFusionType = FusionType::kScalarReduction; + inferredDominantOp = op; } else if (isFusible(op)) { // Ignore if already a kRowReduction or kColReduction, otherwise update // the fusion type to kLoop and dominant op to current op. This supposes @@ -750,6 +765,7 @@ FusionPattern::FusionPattern(lmhlo::FusionOp op, ShapeAnalysis* shape_analysis) FusionType fusionType = FusionType::kNone; auto deviceAttr = op->getAttrOfType(kDiscPlaceAssignment); auto fusionTypeAttr = op->getAttrOfType(kDiscFusionTypeAttrName); + if (fusionTypeAttr) { fusionType = fusionTypeFromString(fusionTypeAttr.getValue()); } @@ -773,6 +789,10 @@ FusionPattern::FusionPattern(lmhlo::FusionOp op, ShapeAnalysis* shape_analysis) FusionStrategy& strategy = getFusionStrategy(deviceAttr.getValue(), strategyStr); bool status = strategy.initFusionPattern(*shape_analysis, *this); + fusion_type_ = fusionType; + if (dominant_op_ == nullptr) { + llvm::dbgs() << "init fusion pattern failed, dominate_op is nullptr\n"; + } assert(status); (void)(status); } @@ -1451,10 +1471,11 @@ bool BaseCpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis, bool BaseGpuFusionStrategy::isFusible(Operation* op) { // Only rank-2 tensor -> rank-1 tensor reduction are supported now. if (isa(op) && - (!isRank2RowReduction(op) && !isRank2ColReduction(op))) + (!isRank2RowReduction(op) && !isRank2ColReduction(op) && + !isRank2ScalarReduction(op))) return false; - if (isa(op) && isRank2or3Transpose(op)) return false; + // if (isa(op) && isRank2or3Transpose(op)) return false; return BaseFusionStrategy::isFusible(op); } @@ -1481,8 +1502,12 @@ bool BaseGpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis, bool has_rank2_col_reduction = llvm::any_of(target.getOpList(), [](Operation* op) { return isRank2ColReduction(op); }); - - if (has_rank2_row_reduction && has_rank2_col_reduction) { + bool has_rank2_scalar_reduction = + llvm::any_of(target.getOpList(), + [](Operation* op) { return isRank2ScalarReduction(op); }); + int cnt = has_rank2_row_reduction + has_rank2_col_reduction + + has_rank2_scalar_reduction; + if (cnt >= 2) { return false; } diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.h b/tao_compiler/mlir/disc/transforms/fusion_utils.h index 06544547d72..d02857d1764 100644 --- a/tao_compiler/mlir/disc/transforms/fusion_utils.h +++ b/tao_compiler/mlir/disc/transforms/fusion_utils.h @@ -111,6 +111,7 @@ enum FusionType { // kInput fusion pattern and all reduce ops of the fused pattern are column // reduction kColReduction, + kScalarReduction, // kInput fusion pattern kInput, // Stitch Fusion pattern @@ -156,6 +157,8 @@ bool isRowReduction(Operation* op); // Returns true if this op is a rank-2 column reduction. bool isRank2ColReduction(Operation* op); +bool isRank2ScalarReduction(Operation* op); + // Returns true if this op is a rank-2 or rank-3 transpose bool isRank2or3Transpose(Operation* op); diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc b/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc index fed8edb2445..7690f7a177c 100644 --- a/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc +++ b/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc @@ -17,10 +17,10 @@ namespace disc_ral { ////////////////////// Stitch GPU FusionStrategy Implemenation ///////// //////////////////////////////////////////////////////////////////////// - bool findValidReductionOps(FusionPatternBase& target, SmallVectorImpl& row_reductions, - SmallVectorImpl& col_reductions) { + SmallVectorImpl& col_reductions, + SmallVectorImpl& scalar_reductions) { row_reductions.clear(); col_reductions.clear(); auto& op_list = target.getOpList(); @@ -28,7 +28,7 @@ bool findValidReductionOps(FusionPatternBase& target, if (!isa(op)) continue; if (isRank2RowReduction(op)) { row_reductions.push_back(op); - } else if (isRank2ColReduction(op)) { + } else if (isRank2ColReduction(op) || isRank2ScalarReduction(op)) { // Middle col-reduction is not supported currently. We may support it with // AStitch technique in the future. int num_input_operand = op->getNumOperands() - getNumResultOperands(op); @@ -41,7 +41,11 @@ bool findValidReductionOps(FusionPatternBase& target, } } } - col_reductions.push_back(op); + if (isRank2ScalarReduction(op)) { + scalar_reductions.push_back(op); + } else { + col_reductions.push_back(op); + } } else { // Non supported reduction type. return false; @@ -65,8 +69,13 @@ bool StitchGpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis, bool has_rank2_col_reduction = llvm::any_of(target.getOpList(), [](Operation* op) { return isRank2ColReduction(op); }); + bool has_rank2_scalar_reduction = + llvm::any_of(target.getOpList(), + [](Operation* op) { return isRank2ScalarReduction(op); }); - if (has_rank2_row_reduction && has_rank2_col_reduction) { + int cnt = has_rank2_row_reduction + has_rank2_col_reduction + + has_rank2_scalar_reduction; + if (cnt >= 2) { return false; } @@ -371,7 +380,9 @@ bool StitchGpuFusionStrategy::findFusionPatternTypeAndSubroot( SmallVector row_reductions; SmallVector col_reductions; - if (!findValidReductionOps(fusion_pattern, row_reductions, col_reductions)) { + SmallVector scalar_reductions; + if (!findValidReductionOps(fusion_pattern, row_reductions, col_reductions, + scalar_reductions)) { LLVM_DEBUG(llvm::dbgs() << "Check reduction ops failed."); return false; } @@ -440,7 +451,23 @@ bool StitchGpuFusionStrategy::findFusionPatternTypeAndSubroot( return true; } Value shape = getEffectiveShape(fusion_pattern, result); - return isRank2ColReduction(op) && + return (isRank2ColReduction(op)) && + shapeAnalysis.isShapeEqual(ref_shape, shape); + })) { + return false; + } + } else if (!scalar_reductions.empty()) { + fusion_type = FusionType::kScalarReduction; + dominant_op = scalar_reductions.back(); + Value ref = cast(dominant_op).getResultBuffer(); + Value ref_shape = getEffectiveShape(fusion_pattern, ref); + if (!llvm::all_of(results, [&](Value result) { + auto op = fusion_pattern.findLastWriter(result); + if (op == dominant_op) { + return true; + } + Value shape = getEffectiveShape(fusion_pattern, result); + return (isRank2ColReduction(op)) && shapeAnalysis.isShapeEqual(ref_shape, shape); })) { return false; diff --git a/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc index 93789480e5c..85d70cb2801 100755 --- a/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc +++ b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc @@ -1069,6 +1069,7 @@ void emitInitLoops(OpBuilder& b, ArrayRef col_reduction_ops) { Location loc = first_root->getLoc(); Value zero = b.create(loc, 0); Value one = b.create(loc, 1); + Value two = b.create(loc, 2); Value num_elements = emitNumElementsComputation(b, loc, first_root); Value var; @@ -1091,7 +1092,7 @@ void maybeEmitInitLoops(OpBuilder& b, b.setInsertionPoint(root_ops.back()); SmallVector col_reduction_ops; for (Operation* root_op : root_ops) { - if (isRank2ColReduction(root_op)) { + if (isRank2ColReduction(root_op) || isRank2ScalarReduction(root_op)) { col_reduction_ops.emplace_back(root_op); } } @@ -1260,6 +1261,333 @@ LogicalResult lowerWithScheduleRowReduction(ArrayRef, Operation*, return failure(); } +/** + * for (m = 0; m < block_nums; ++m) { + * for (n = 0; n < block_size; ++n) { + * __shared__ shm[block_size]; + * int tid = threadIdx.x; + * int i = blockIdx.x * block_size * 2 + tid; + * int grid_size = gridDim.x * block_size * 2; + * for (int j = i; j < n; j+=grid_size) { + * shm[tid] += inputs[j] + inputs[j + block_size]; + * } + * __syncthreads(); + * for (int stride = block_size / 2; stride > 0; stride /= 2) { + * if (tid < stride) { + * shm[tid] += shm[tid + stride]; + * } + * __syncthreads(); + * } + * if (tid < wrap_size) { + * auto val = shm[tid]; + * for (int stride = warp_size; stride > 0; stride /= 2) { + * val += __shfl_down_sync(0xffffffff, val, stride); + * } + * if (tid == 0) { + * atomicAdd(&g_output[0], shm[0]); + * } + * } + * } + * } + */ + +LogicalResult lowerWithScheduleParallelReduction( + ArrayRef root_ops, Operation* dominant_op, Block* parent, + const ShapeAnalysis* shape_analysis = nullptr, int vector_size = 1) { + if (!isRank2ScalarReduction(dominant_op)) { + return failure(); + } + // Create helper Values + SmallVector scalar_reduction_roots; + std::copy_if( + root_ops.begin(), root_ops.end(), + std::back_inserter(scalar_reduction_roots), + [](Operation* operation) { return isRank2ScalarReduction(operation); }); + auto root_op = scalar_reduction_roots.back(); + // after offline tunning, 12 elements_per_thread got the best performance for + // 256 threads per block + const int elements_per_thread = 12; + const int thread_per_block = 256; + Location loc = dominant_op->getLoc(); + OpBuilder b(root_ops.back()); + + Value lhs = dominant_op->getOperand(0); + Value rhs = dominant_op->getOperand(2); + Value zero = b.create(loc, 0); + Value one = b.create(loc, 1); + Value two = b.create(loc, 2); + Value total_elements = one; + for (auto root_op : scalar_reduction_roots) { + auto op_elements = b.create( + loc, b.create(loc, root_op->getOperand(0), zero), + b.create(loc, root_op->getOperand(0), one)); + total_elements = b.create(loc, total_elements, op_elements); + } + // num_blocks = ceil(total_elements/ block_size); + Value block_size = b.create(loc, thread_per_block); + Value elements_per_block = b.create( + loc, elements_per_thread * thread_per_block); + Value num_blocks = + b.create(loc, total_elements, elements_per_block); + + auto global_workgroup = b.create( + loc, SmallVector({zero}), SmallVector({num_blocks}), + SmallVector({one}), SmallVector({}), + /*bodyBuilderFn=*/nullptr); + b.setInsertionPointToStart(global_workgroup.getBody()); + auto local_workgroup = b.create( + loc, SmallVector({zero}), SmallVector({block_size}), + SmallVector({one}), SmallVector({}), + /*bodyBuilderFn=*/nullptr); + b.setInsertionPointToStart(local_workgroup.getBody()); + Value block_dim = + b.create(loc, b.getIndexType(), gpu::Dimension::x); + Value block_idx = + b.create(loc, b.getIndexType(), gpu::Dimension::x); + Value tid = + b.create(loc, b.getIndexType(), gpu::Dimension::x); + Value grid_dim = + b.create(loc, b.getIndexType(), gpu::Dimension::x); + + // i = blockIdx.x * block_size * 2 + tid; + Value i = b.create( + loc, + b.create( + loc, b.create(loc, block_idx, block_size), two), + tid); + + // grid_size = gridDim.x * block_size * 2; + Value grid_size = b.create( + loc, b.create(loc, grid_dim, block_size), two); + + Value n = b.create(loc, lhs, zero); + Value m = b.create(loc, lhs, one); + Value mn = b.create(loc, m, n); + + // acc: init_values[num_scalar_reductions] + SmallVector accum_factory( + scalar_reduction_roots.size()); + SmallVector init_values(scalar_reduction_roots.size()); + SmallVector init_values_types(scalar_reduction_roots.size()); + for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { + Operation* root_op = root_pair.value(); + int idx = root_pair.index(); + Value init_value = b.create( + loc, cast(root_op).getInitValues()[0]); + init_values[idx] = init_value; + init_values_types[idx] = init_value.getType(); + accum_factory[idx] = std::move(getFactory( + b, root_op->getLoc(), cast(root_op).getBody())); + } + // define SHM + std::map shared_mem_map; + for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { + Operation* root_op = root_pair.value(); + int idx = root_pair.index(); + const auto elemType = getLhloOpsElementType(root_op); + auto shared_mem = createSharedMemory(b, loc, thread_per_block, + getLhloOpsElementType(dominant_op)); + shared_mem_map[root_op] = shared_mem; + } + + { + // fused accumulation + SmallVector yield_values_for_if; + Value var_j = nullptr; + // for (; i < n; i += grid_size) + // acc += inputs[i] + inputs[i + grid_size]; + scf::ForOp for_op_k = + createLoopAndSetInsPt(b, loc, var_j, i, mn, grid_size, init_values); + for_op_k.getBody()->clear(); + b.setInsertionPointToStart(for_op_k.getBody()); + int scalar_red_root_op_idx = 0; + SmallVector load_index({var_j, zero}); + SmallVector load_index2( + {b.create(loc, var_j, block_size), zero}); + for (auto* root_op : root_ops) { + const auto elem_type = getLhloOpsElementType(root_op); + if (isRank2ScalarReduction(root_op)) { + auto lhs = root_op->getOperands().begin(); + Value data = createLoadOrUseCachedValue( + loc, &b, root_op, *lhs, load_index, b.saveInsertionPoint()); + + Value iter_value = + *(for_op_k.getRegionIterArgs().begin() + scalar_red_root_op_idx); + Value acc = (accum_factory[scalar_red_root_op_idx])(data, iter_value); + // if (i + block_size < n) sum += inputs[i + block_size]; + auto is_valid_load = b.create( + loc, arith::CmpIPredicate::slt, + b.create(loc, var_j, block_size), mn); + auto if_bound_op = b.create( + loc, /*resultTypes=*/TypeRange{elem_type}, is_valid_load, + /*hasElseRegion=*/true); + if_bound_op.getThenRegion().front().clear(); + b.setInsertionPointToStart(&if_bound_op.getThenRegion().front()); + Value data2 = createLoadOrUseCachedValue( + loc, &b, root_op, *lhs, load_index2, b.saveInsertionPoint()); + auto acc2 = (accum_factory[scalar_red_root_op_idx])(acc, data2); + b.create(loc, ValueRange({acc2})); + b.setInsertionPointToStart(&if_bound_op.getElseRegion().front()); + b.create(loc, ValueRange({acc})); + b.setInsertionPointAfter(if_bound_op); + yield_values_for_if.push_back(if_bound_op.getResults()[0]); + // yield_values_for_if.push_back(acc); + scalar_red_root_op_idx++; + } else { + auto dominant_shape = getShapeValues(&b, dominant_op->getOperand(0)); + Value linear_index = + calcLinearIndex(&b, loc, load_index, dominant_shape); + if (!succeeded( + lowerHelper(b, loc, root_op, linear_index, shape_analysis))) { + return failure(); + } + } + } + b.create(loc, yield_values_for_if); + b.setInsertionPointAfter(for_op_k); + for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { + Operation* root_op = root_pair.value(); + int idx = root_pair.index(); + b.create(loc, *(for_op_k.getResults().begin() + idx), + shared_mem_map[root_op], tid); + } + b.create(loc); + } + + { + SmallVector init_values = {}; + for (int stride = 128; stride >= 64; stride /= 2) { + Value stride_val = b.create(loc, stride); + // if (tid < stride) + scf::IfOp if_tid_valid_op = b.create( + loc, /*resultTypes*/ TypeRange{}, + b.create(loc, arith::CmpIPredicate::slt, tid, + stride_val), + /*hasElseRegion*/ false); + if_tid_valid_op.getThenRegion().front().clear(); + b.setInsertionPointToStart(&if_tid_valid_op.getThenRegion().front()); + SmallVector yield_values; + for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { + Operation* root_op = root_pair.value(); + int idx = root_pair.index(); + Value shm_val_1 = + b.create(loc, shared_mem_map[root_op], tid); + Value strid_tid = b.create(loc, tid, stride_val); + Value shm_val_2 = + b.create(loc, shared_mem_map[root_op], strid_tid); + Value sum = (accum_factory[idx])(shm_val_1, shm_val_2); + b.create(loc, sum, shared_mem_map[root_op], tid); + } + b.create(loc, yield_values); + b.setInsertionPointAfter(if_tid_valid_op); + b.create(loc); + } + } + + // warp reduce with shuffle_down + SmallVector shuffle_val(scalar_reduction_roots.size()); + SmallVector shuffle_type(scalar_reduction_roots.size()); + SmallVector shm_values(scalar_reduction_roots.size()); + { + for (const auto& root_pair : llvm::enumerate(scalar_reduction_roots)) { + Operation* root_op = root_pair.value(); + int idx = root_pair.index(); + Value shm_val = + b.create(loc, shared_mem_map[root_op], tid); + shm_values[idx] = shm_val; + shuffle_val[idx] = shm_val; + + const auto elem_type = getLhloOpsElementType(root_op); + Type shuffle_elem_type; + if (failed(getShuffleElemType(b, elem_type, &shuffle_elem_type))) { + return failure(); + } + shuffle_type[idx] = shuffle_elem_type; + } + + // for (offset = warpSize >> 1, offset > 0; offset >>= 1) { + // sum += __shfl_down(sum, offset); + // } + // if (tid < warpSize) { + // shm[tid] += shm[tid + warpSize]; + // sum = shm[tid]; + // sum = warp_reduce(sum) + // if (tid == 0) atomicAdd(&g_output[0], sum); + // } + Value warp_size = b.create(loc, 32); + scf::IfOp if_tid_valid_warp = b.create( + loc, shuffle_type, + b.create(loc, arith::CmpIPredicate::slt, tid, warp_size), + true); + + if_tid_valid_warp.getThenRegion().front().clear(); + if_tid_valid_warp.getElseRegion().front().clear(); + b.setInsertionPointToStart(&if_tid_valid_warp.getThenRegion().front()); + for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { + Operation* root_op = root_pair.value(); + int idx = root_pair.index(); + Value shm_val_1 = + b.create(loc, shared_mem_map[root_op], tid); + Value strid_tid = b.create(loc, tid, warp_size); + Value shm_val_2 = + b.create(loc, shared_mem_map[root_op], strid_tid); + Value sum = (accum_factory[idx])(shm_val_1, shm_val_2); + shuffle_val[idx] = sum; + } + // warp shuffle down + Value shuffle_mask = + b.create(loc, 0xFFFFFFFF, b.getIntegerType(32)); + for (int offset = kWarpSize / 2; offset > 0; offset /= 2) { + Value offset_val = + b.create(loc, offset, b.getIntegerType(32)); + for (const auto& root_pair : llvm::enumerate(scalar_reduction_roots)) { + Operation* root_op = root_pair.value(); + int idx = root_pair.index(); + auto result = emitWidthAdaptShuffle( + b, loc, shuffle_val[idx], shuffle_type[idx], offset_val, + shuffle_mask, gpu::ShuffleMode::DOWN); + shuffle_val[idx] = (accum_factory[idx])(shuffle_val[idx], result); + } + } + b.setInsertionPointToEnd(&if_tid_valid_warp.getThenRegion().front()); + b.create(loc, shuffle_val); + b.setInsertionPointToStart(&if_tid_valid_warp.getElseRegion().front()); + b.create(loc, init_values); + b.setInsertionPointAfter(if_tid_valid_warp); + Value shuffle_result = if_tid_valid_warp.getResults().front(); + // if (tid == 0) atomicAdd(&g_output[0], sum); + Value is_tid_zero_op = + b.create(loc, arith::CmpIPredicate::eq, tid, zero); + scf::IfOp if_tid_zero_op = + b.create(loc, /*resultTypes*/ TypeRange{}, is_tid_zero_op, + /*hasElseRegion*/ false); + if_tid_zero_op.getThenRegion().front().clear(); + b.setInsertionPointToStart(&if_tid_zero_op.getThenRegion().front()); + for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { + Operation* root_op = root_pair.value(); + int idx = root_pair.index(); + Value val = if_tid_valid_warp.getResults()[idx]; + Type root_element_type = getLhloOpsElementType(root_op); + + b.create( + loc, shuffle_type[idx], + getAtomicRMWKind(cast(root_op).getBody()), val, + root_op->getOperand(2), ValueRange({})); + } + b.create(loc, ValueRange({})); + b.setInsertionPointAfter(if_tid_zero_op); + } + b.setInsertionPointToEnd(local_workgroup.getBody()); + b.setInsertionPointAfter(global_workgroup); + if (parent == nullptr) { + for (Operation* root_op : root_ops) root_op->erase(); + } else { + assert(parent != nullptr && "Parent must be provided for fusion lowering"); + cleanUnusedLhloOps(parent); + } + return success(); +} + /* Row reduction with 1 round warp shuffle * * RowPerBlock = threads / warpSize; @@ -1894,7 +2222,12 @@ LogicalResult lowerWithScheduleRowReduction( b.create(loc, ValueRange({})); // remove the root_op if it has no other users except the memref - cleanUnusedLhloOps(parent); + if (parent == nullptr) { + for (Operation* root_op : root_ops) root_op->erase(); + } else { + assert(parent != nullptr && "Parent must be provided for fusion lowering"); + cleanUnusedLhloOps(parent); + } return success(); } @@ -2178,10 +2511,12 @@ LogicalResult lowerWithScheduleColReductionForRocm( b.create(loc, ValueRange({})); b.setInsertionPointAfter(if_col_valid); - // row_index_shuffle = sw_block_row_index_base * var_tile_h + local_row_index - // for (int stride = var_tile_h / 2; stride > 1; stride /= 2) { + // row_index_shuffle = sw_block_row_index_base * var_tile_h + + // local_row_index for (int stride = var_tile_h / 2; stride > 1; stride /= + // 2) { // __syncthreads(); - // if (local_row_index < stride && (row_index_shuffle + stride) < var_rows) + // if (local_row_index < stride && (row_index_shuffle + stride) < + // var_rows) // { // shm[n] += shm[stride * var_tile_w + n]; // } @@ -2626,8 +2961,8 @@ LogicalResult emitFirstRoundShuffleStitch( createAlignMemrefWithTile(b, result_buffer_shm, row_tile); } for (int i = 0; i < row_tile; i++) { - // The elements store in result shm buffer are in index order. (Note it - // is linear index.) + // The elements store in result shm buffer are in index order. (Note + // it is linear index.) b.create(loc, sum_vec[i], result_buffer_shm, block_row_offset[i]); } @@ -3053,8 +3388,8 @@ LogicalResult initSkeletonGrpsAndCloneOps( } } } else { - // For to-be-cloned ops, alloc for output and then clone the op with new - // outputs. + // For to-be-cloned ops, alloc for output and then clone the op with + // new outputs. SmallVector results; for (Value v : op->getOperands().drop_front(num_input_operand)) { Value new_operand = allocClonedValue(v); @@ -3266,8 +3601,8 @@ LogicalResult lowerWithScheduleStitch(lmhlo::FusionOp& fusion_op, } SmallVector outShapeValues = getShapeValues(&b, out_value); - // Deal with the case that tiled dims are not the same between result and - // sub-roots' input. + // Deal with the case that tiled dims are not the same between result + // and sub-roots' input. Value tiled_linear = nullptr; for (const auto& en : llvm::enumerate(outShapeValues)) { if (tile_info->second.tileSizes.count(en.index()) > 0) { @@ -3306,9 +3641,9 @@ LogicalResult lowerWithScheduleStitch(lmhlo::FusionOp& fusion_op, b.setInsertionPointAfter(if_row_in_bound); } - // Currently, a non-sub-root skeleton op will always be external only. We - // many need the following check for non-reduce sub-root someday. - // if (!external_only_roots.contains(skeleton)) { + // Currently, a non-sub-root skeleton op will always be external only. + // We many need the following check for non-reduce sub-root someday. if + // (!external_only_roots.contains(skeleton)) { // // TODO: use __threadfence_block instead. // b.create(loc); // } @@ -3654,8 +3989,8 @@ LogicalResult emitRowReduceThreadBlockV2( // generated code. for (int64_t i = 0; i < ops.size(); i++) { if (result_buffer_shms[i] != nullptr && !external_output_only[i]) { - // The elements store in result shm buffer are in index order. (Note - // it is linear index.) + // The elements store in result shm buffer are in index order. + // (Note it is linear index.) b.create(loc, warp_reduces[i], result_buffer_shms[i], warp_id); } @@ -3737,8 +4072,8 @@ LogicalResult emitRowReduceThreadBlockV2( b.setInsertionPointToStart( &if_lane_id_is_zero_2.getThenRegion().front()); { - // Do not fuse the loops in this block. It relies the loops for ILP of - // generated code. + // Do not fuse the loops in this block. It relies the loops for ILP + // of generated code. for (int64_t i = 0; i < ops.size(); i++) { if (result_buffer_shms[i] != nullptr && !external_output_only[i]) { b.create(loc, reduce_r2s[i], @@ -4159,6 +4494,7 @@ LogicalResult HandleGpuFusionOp(OpBuilder& b, Operation* fusion, r = lowerWithScheduleRowReduction( root_ops, dominant_op, fused_block, shape_analysis, vector_size); } + if (failed(r)) { return dominant_op->emitError() << "failed to lower row-reduction loops"; @@ -4188,12 +4524,21 @@ LogicalResult HandleGpuFusionOp(OpBuilder& b, Operation* fusion, r = lowerWithScheduleColReduction<512, 32>(root_ops, dominant_op, fused_block); } + if (failed(r)) { return dominant_op->emitError() << "failed to lower col-reduction loops"; } } break; - + case FusionType::kScalarReduction: { + auto kname = getFusionFullName(fusion_op); + LogicalResult r = lowerWithScheduleParallelReduction( + root_ops, dominant_op, fused_block); + if (failed(r)) { + return dominant_op->emitError() + << "failed to lower scalar-reduction loops"; + } + } break; case FusionType::kLoop: { const int vector_size = getVectorizeOrTileHint(dominant_op); if (isMemIntensiveOptExperimentalEnabled()) { @@ -4936,9 +5281,9 @@ LogicalResult lowerWithScheduleSparseSegmentReductionOpCPU( segment_count_memref = alloc.getResult(); { - auto for_op = - b.create(loc, /* lowerBound */ zero, - /* upperBound */ num_results, /* step */ one); + auto for_op = b.create(loc, /* lowerBound */ zero, + /* upperBound */ num_results, + /* step */ one); for_op.getBody()->clear(); b.setInsertionPointToStart(for_op.getBody()); Value i = for_op.getInductionVar(); @@ -5575,6 +5920,7 @@ LogicalResult lowerWithSchedulekInputCPU( b.create(loc, acc, out, outVars); b.setInsertionPointAfter(reductionForOp); + // remove the root_op if it has no other users except the memref for (Operation* root_op : root_ops) root_op->erase(); return success(); } @@ -5712,11 +6058,12 @@ struct DiscLhloLegalizeRootsToParallelLoops // TODO: We should put even single nodes into a fusion by fusion pass // Revisit this and walk lmhlo::FusionOp only after the revision done. func.walk([&](lmhlo::LmhloOp op) { - // Skip the embedded ops in lmhlo.fusion or lmhlo.reduce/scatter or - // lmhlo_disc.args_mutation + // Skip the embedded ops in lmhlo.fusion or lmhlo.reduce/scatter lmhlo::LmhloOp parent = op->getParentOfType(); - if (isa(op) || - parent && !isa(op)) { + if (isa(op)) { + return; + } + if (parent && !isa(op)) { return; } if (isFusionType(op) && @@ -5758,10 +6105,28 @@ struct DiscLhloLegalizeRootsToParallelLoops } for (Operation* op : gpu_non_fusion_worklist) { + if (isa(op) && isScalarReduction(op)) { + emitInitLoops(b, {op}); + if (failed(lowerWithScheduleParallelReduction({op}, op, nullptr, + &shape_analysis))) { + op->emitError() << "failed to lower non fusion reduction"; + signalPassFailure(); + return; + } + continue; + } + if (isa(op) && isRank2RowReduction(op)) { + if (failed(lowerWithScheduleRowReduction( + {op}, op, nullptr, &shape_analysis, 1))) { + op->emitError() << "failed to lower non fusion row reduction"; + signalPassFailure(); + return; + } + continue; + } // TODO(disc): single nodes with non kLoop schedule like ReduceOp // is not implemented yet. Currently ReduceOp is lowered with loop // schedule, which means for poor performance. - if (failed(lowerWithScheduleLoop({op}, op, nullptr, /*non_fusion=*/true, /*parallel_loop=*/true))) { diff --git a/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc b/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc index de1d4aa068c..400c170f765 100644 --- a/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc +++ b/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc @@ -47,8 +47,10 @@ struct ParallelLoopCollapsing if (fusion) { auto fusionTypeAttr = fusion->getAttrOfType(kDiscFusionTypeAttrName); - if (fusionTypeAttr && (fusionTypeAttr.getValue() == "kStitch" || - fusionTypeAttr.getValue() == "kTransform")) { + if (fusionTypeAttr && + (fusionTypeAttr.getValue() == "kStitch" || + fusionTypeAttr.getValue() == "kScalarReduction" || + fusionTypeAttr.getValue() == "kTransform")) { continue; } } diff --git a/tao_compiler/mlir/disc/transforms/parallel_loop_tiling.cc b/tao_compiler/mlir/disc/transforms/parallel_loop_tiling.cc index 551581be5a4..2584cc51f4e 100644 --- a/tao_compiler/mlir/disc/transforms/parallel_loop_tiling.cc +++ b/tao_compiler/mlir/disc/transforms/parallel_loop_tiling.cc @@ -65,7 +65,9 @@ struct ParallelLoopTiling // Do not deal with kStitch fusion. auto fusionTypeAttr = fusion->getAttrOfType(kDiscFusionTypeAttrName); - if (fusionTypeAttr && fusionTypeAttr.getValue() == "kStitch") { + if (fusionTypeAttr && + (fusionTypeAttr.getValue() == "kStitch" || + fusionTypeAttr.getValue() == "kScalarReduction")) { continue; } } diff --git a/tf_community b/tf_community index 1f7c9e80e6a..83cc10030c8 160000 --- a/tf_community +++ b/tf_community @@ -1 +1 @@ -Subproject commit 1f7c9e80e6a9eb786d960b4b84133d645a36e39c +Subproject commit 83cc10030c84c0538ba9455ea43431eb13374ec9