From 4b10aa0193326d31d22af31a19f4cd5bea3b85d2 Mon Sep 17 00:00:00 2001 From: yancey Date: Fri, 1 Mar 2024 14:12:44 +0800 Subject: [PATCH 01/10] support scalar reduction --- tao_compiler/mlir/disc/disc_compiler.cc | 2 +- .../mlir/disc/transforms/fusion_utils.cc | 38 +- .../mlir/disc/transforms/fusion_utils.h | 1 + .../transforms/fusion_utils_stitch_gpu.cc | 54 ++- .../lhlo_legalize_roots_to_loops.cc | 377 +++++++++++++++++- .../transforms/parallel_loop_collapsing.cc | 5 +- .../disc/transforms/parallel_loop_tiling.cc | 4 +- 7 files changed, 463 insertions(+), 18 deletions(-) diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc index e8f0eb06ae3..17770b0845c 100644 --- a/tao_compiler/mlir/disc/disc_compiler.cc +++ b/tao_compiler/mlir/disc/disc_compiler.cc @@ -601,7 +601,7 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) { pm.addNestedPass(createCSEPass()); pm.addNestedPass( createCanonicalizerPass(cano_rewrite_config, disablePatterns)); - pm.addNestedPass(disc_ral::createDiscMemRefCSEPass()); + // 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/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc index 2475f1d94fb..9614658f580 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,21 @@ bool isRowReduction(Operation* op) { return true; } +bool isRank2ScalarReduction(Operation* op) { + auto reduce_op = dyn_cast(op); + if (!reduce_op || reduce_op.getDimensions().getNumElements() != 1) + return false; + int rank = op->getOperand(2).getType().cast().getRank(); + // TODO(yancey): rewrite scalar reduction result to scalar tensor to avoid + // reshape to scalar tensor behand reduce op + Operation* reshapeOp = *op->getOperand(2).getUsers().begin(); + if (isa(reshapeOp) && + reshapeOp->getOperand(1).getType().cast().getRank() == 0) { + return true; + } + return false; +} + // Returns true if this op is a rank-2 column reduction. bool isRank2ColReduction(Operation* op) { auto reduce_op = dyn_cast(op); @@ -487,7 +506,8 @@ bool isRank2ColReduction(Operation* op) { int rank = op->getOperand(0).getType().cast().getRank(); auto dimensions = reduce_op.getDimensions().getValues(); - return ((*dimensions.begin() == 0) && (rank == 2)); + return ((*dimensions.begin() == 0) && (rank == 2)) && + !isRank2ScalarReduction(op); } // Return true if this op is a rank-2 transpose @@ -558,6 +578,9 @@ bool initFusionPatternBase(ShapeAnalysis& shapeAnalysis, 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 +773,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 +797,7 @@ FusionPattern::FusionPattern(lmhlo::FusionOp op, ShapeAnalysis* shape_analysis) FusionStrategy& strategy = getFusionStrategy(deviceAttr.getValue(), strategyStr); bool status = strategy.initFusionPattern(*shape_analysis, *this); + fusion_type_ = fusionType; assert(status); (void)(status); } @@ -1451,7 +1476,8 @@ 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))) // || isScalarReduction(op))) return false; if (isa(op) && isRank2or3Transpose(op)) return false; @@ -1481,8 +1507,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..b58bbeeefc5 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 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..75c4644e2fa 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,24 @@ namespace disc_ral { ////////////////////// Stitch GPU FusionStrategy Implemenation ///////// //////////////////////////////////////////////////////////////////////// - +bool isScalarReduction(Operation* op) { + auto reduce_op = dyn_cast(op); + if (!reduce_op || reduce_op.getDimensions().getNumElements() != 1) + return false; + int rank = op->getOperand(2).getType().cast().getRank(); + // TODO(yancey): rewrite scalar reduction result to scalar tensor to avoid + // reshape to scalar tensor behand reduce op + Operation* reshapeOp = *op->getOperand(2).getUsers().begin(); + if (reshapeOp && isa(reshapeOp) && + reshapeOp->getOperand(1).getType().cast().getRank() == 0) { + return true; + } + return false; +} 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 +42,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) || isScalarReduction(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 +55,11 @@ bool findValidReductionOps(FusionPatternBase& target, } } } - col_reductions.push_back(op); + if (isScalarReduction(op)) { + scalar_reductions.push_back(op); + } else { + col_reductions.push_back(op); + } } else { // Non supported reduction type. return false; @@ -65,8 +83,12 @@ 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 isScalarReduction(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 +393,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 +464,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..59be2d138ee 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 @@ -1259,6 +1259,361 @@ LogicalResult lowerWithScheduleRowReduction(ArrayRef, Operation*, int vector_size = 1) { return failure(); } +/* Row reduction with 1 round warp shuffle + * + * RowPerBlock = threads / warpSize; + * for (m = 0; m < rows; m += RowPerBlock) { + * for (n = 0; n < threads; ++n) { + * rowIdx = m + warpIdx; + * if (rowIdx < rows) { + * // intra-thread reduction + * sum = init_value; + * for (k = laneIdx; k < cols; k += warpSize) { + * sum += inputs[rowIdx][k]; + * } + * + * // inter-thread reduction via warp shuffle + * for (offset = warpSize / 2; offset > 0; offset /= 2) { + * sum += __shfl_xor(sum, offset); + * } + * + * // write to output + * if (laneIdx == 0) { + * outputs[rowIdx] = sum + * } + * } + * } + */ +/** + * 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 == 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 (!isRank2ColReduction(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 isRank2ColReduction(operation); }); + auto root_op = scalar_reduction_roots.back(); + const int thread_per_block = getCTASize(dominant_op); + ; + 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); + auto elemFloatType = getLhloOpsElementType(root_op).cast(); + Value zero_f = b.create( + loc, llvm::APFloat(0.0f), elemFloatType); // b.getF32Type()); + Value one_f = b.create( + loc, llvm::APFloat(1.0f), elemFloatType); // b.getF32Type()); + + // Start to emit. + Value num_blocks = b.create(loc, 1024); + Value block_size = b.create(loc, 256); + + 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); + local_workgroup.getBody()->clear(); + 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); + tid = b.create(loc, tid, block_dim); + // i = blockIdx.x * block_size * 2 + tid; + Value i = b.create( + loc, + b.create( + loc, b.create(loc, block_idx, block_dim), two), + tid); + + // grid_size = gridDim.x * block_size * 2; + Value grid_size = b.create( + loc, b.create(loc, grid_dim, block_dim), two); + + Value n = b.create(loc, lhs, zero); + Value m = b.create(loc, lhs, one); + Value mn = b.create(loc, m, n); + // __shared__ float shm[block_size]; + auto shared_mem = createSharedMemory(b, loc, thread_per_block, + getLhloOpsElementType(dominant_op)); + Value acc_value; + // fused accumulation + // acc: init_values[num_col_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; + for (auto* root_op : root_ops) { + if (isRank2ColReduction(root_op)) { + auto lhs = root_op->getOperands().begin(); + SmallVector load_index({i, zero}); + Value data = createLoadOrUseCachedValue( + loc, &b, root_op, *lhs, load_index, b.saveInsertionPoint()); + SmallVector load_index2( + {b.create(loc, i, block_dim), zero}); + Value data1 = createLoadOrUseCachedValue( + loc, &b, root_op, *lhs, load_index2, b.saveInsertionPoint()); + Value sum = (accum_factory[scalar_red_root_op_idx])(data, data1); + auto acc = (accum_factory[scalar_red_root_op_idx])( + *(for_op_k.getRegionIterArgs().begin() + scalar_red_root_op_idx), + sum); + yield_values_for_if.push_back(acc); + scalar_red_root_op_idx++; + } else if (isa(root_op)) { + auto dominant_shape = getShapeValues(&b, dominant_op->getOperand(0)); + Value linear_index = calcLinearIndex(&b, loc, i, dominant_shape); + auto root_shape = getShapeValues(&b, root_op->getOperand(0)); + auto mapped_index = + calcMultiDimIndex(&b, loc, linear_index, root_shape); + emitNotToVectorReduction(b, loc, root_op, mapped_index); + } else { + auto dominant_shape = getShapeValues(&b, dominant_op->getOperand(0)); + Value linear_index = calcLinearIndex(&b, loc, i, 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); + } + // acc_value = for_op_k.getResult(0); + // b.create(loc, acc_value, shared_mem, tid); + } + /* + { + Value init_value = b.create( + loc, cast(root_op).getInitValues()[0]); + SmallVector init_values{init_value}; + 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()); + auto lhs = dominant_op->getOperand(0); + SmallVector load_index({var_j, zero}); + Value data = createLoadOrUseCachedValue( + loc, &b, dominant_op, lhs, load_index, b.saveInsertionPoint()); + + SmallVector load_index2({b.create(loc, var_j, + block_dim), zero}); Value data1 = createLoadOrUseCachedValue( loc, &b, + dominant_op, lhs, load_index2, b.saveInsertionPoint()); + + Value sum = b.create(loc, data, data1); + acc = b.create(loc, acc, + *for_op_k.getRegionIterArgs().begin()); b.create(loc, + ValueRange{acc}); b.setInsertionPointAfter(for_op_k); acc_value = + for_op_k.getResult(0); b.create(loc, acc_value, shared_mem, + tid); + } + */ + { + Value var_j = nullptr; + SmallVector init_values = {}; + for (int stride = 128; stride > 16; stride /= 2) { + b.create(loc); + 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); + } + /* + SmallVector multidim_load_index({tid}); + ValueRange load_index(multidim_load_index); + + SmallVector multidim_load_index2({b.create(loc, + tid, stride_val)}); ValueRange load_index2(multidim_load_index2); + + Value in_data = b.create(loc, shared_mem, load_index); + Value in_data1 = b.create(loc, shared_mem, load_index2); + Value sum = b.create(loc, in_data, in_data1); + b.create(loc, sum, shared_mem, tid); + */ + b.create(loc); + b.create(loc, yield_values); + b.setInsertionPointAfter(if_tid_valid_op); + } + } + b.create(loc); + { + // warp reduce + // if (tid < 16) + scf::IfOp if_tid_valid_op = b.create( + loc, /*resultTypes*/ TypeRange{}, + b.create(loc, arith::CmpIPredicate::slt, tid, + b.create(loc, 32)), + /*hasElseRegion*/ false); + b.setInsertionPointToStart(&if_tid_valid_op.getThenRegion().front()); + SmallVector yield_values; + for (int stride = 16; stride > 0; stride /= 2) { + Value stride_val = b.create(loc, stride); + 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); + } + + // shm[tid] += shm[tid + stride]; + // Value idx2 = b.create(loc, tid, stride_val); + // Value shm_val_1 = b.create(loc, shared_mem, tid); + // Value shm_val_2 = b.create(loc, shared_mem, idx2); + // Value sum = b.create(loc, shm_val_1, shm_val_2); + // b.create(loc, sum, shared_mem_map[root_op], tid); + // b.create(loc); + } + b.setInsertionPointAfter(if_tid_valid_op); + } + b.create(loc); + + { + // if (tid == 0) + 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()); + SmallVector yield_values; + for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { + Operation* root_op = root_pair.value(); + int idx = root_pair.index(); + Value val = b.create(loc, shared_mem_map[root_op], tid); + Type root_element_type = getLhloOpsElementType(root_op); + b.create( + loc, root_element_type, + getAtomicRMWKind(cast(root_op).getBody()), val, + root_op->getOperand(2), ValueRange({zero})); + } + /* + Value val = b.create(loc, shared_mem, tid); + Type root_element_type = getLhloOpsElementType(root_op); + b.create( + loc, root_element_type, + getAtomicRMWKind(cast(root_op).getBody()), + val, root_op->getOperand(2), ValueRange({zero})); + */ + b.create(loc, yield_values); + + b.setInsertionPointAfter(if_tid_zero_op); + } + b.setInsertionPointToEnd(local_workgroup.getBody()); + b.create(loc, ValueRange({})); + 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 * @@ -4159,6 +4514,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 +4544,22 @@ 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); + llvm::errs() << "kScalarReduction <" << kname << ">\n"; + 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()) { @@ -5575,7 +5941,13 @@ LogicalResult lowerWithSchedulekInputCPU( b.create(loc, acc, out, outVars); b.setInsertionPointAfter(reductionForOp); - for (Operation* root_op : root_ops) root_op->erase(); + // remove the root_op if it has no other users except the memref + if (non_fusion) { + for (Operation* root_op : root_ops) root_op->erase(); + } else { + assert(parent != nullptr && "Parent must be provided for fusion lowering"); + cleanUnusedLhloOps(parent); + } return success(); } @@ -5761,7 +6133,6 @@ struct DiscLhloLegalizeRootsToParallelLoops // 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..b24bb33e27f 100644 --- a/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc +++ b/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc @@ -47,8 +47,9 @@ 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")) { 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; } } From d9ed7a1c2d5df696fd9888375aa20a2f53d71d5a Mon Sep 17 00:00:00 2001 From: yancey Date: Tue, 5 Mar 2024 15:53:44 +0800 Subject: [PATCH 02/10] update --- pytorch_blade/bazel_build.py | 16 ++-- pytorch_blade/torch_blade/dynamo/__init__.py | 3 +- .../mlir/disc_engine_conversion.py | 1 + scripts/python/tao_build.py | 6 +- .../mlir/custom_ops/transpose_impl.cc | 2 + .../transforms/disc_lower_to_library_call.cc | 2 +- .../mlir/disc/transforms/fusion_utils.cc | 53 +++++++---- .../lhlo_legalize_roots_to_loops.cc | 89 +------------------ 8 files changed, 54 insertions(+), 118 deletions(-) diff --git a/pytorch_blade/bazel_build.py b/pytorch_blade/bazel_build.py index e41ebaba6d3..32f2aea10dc 100644 --- a/pytorch_blade/bazel_build.py +++ b/pytorch_blade/bazel_build.py @@ -68,9 +68,9 @@ def __init__(self, *args, **kwargs): "@org_disc_compiler//mlir/custom_ops:libdisc_custom_ops.so", "//pytorch_blade:libtorch_blade.so", "//pytorch_blade:_torch_blade.so", - "//tests/mhlo/torch-mlir-opt:torch-mlir-opt", - "//tests/torchscript:shape_analysis_tool", - "//tests/torch-disc-pdll:torch-disc-pdll", + #"//tests/mhlo/torch-mlir-opt:torch-mlir-opt", + #"//tests/torchscript:shape_analysis_tool", + #"//tests/torch-disc-pdll:torch-disc-pdll", ] torch_major_version, torch_minor_version = self.torch_version.split(".")[:2] @@ -264,15 +264,17 @@ def test(self): env["GCC_HOST_COMPILER_PATH"] = env.get("GCC_HOST_COMPILER_PATH", which("gcc")) self.test_suites = [ - "//tests/mhlo/...", - "//pytorch_blade:torch_blade_test_suite", - "//tests/torch-disc-pdll/tests/...", + "@org_disc_compiler//mlir/ral:collective_ops_test", + #"//tests/mhlo/...", + #"//pytorch_blade:torch_blade_test_suite", + #"//tests/torch-disc-pdll/tests/...", ] if (self.torch_major_version, self.torch_minor_version) > (1, 6): # torchscript graph ir parser changed after torch 1.6. # We will not test torchscript graph ir before torch 1.6 - self.test_suites.append("//tests/torchscript/...") + #self.test_suites.append("//tests/torchscript/...") + pass test_cmd = " ".join( [self.shell_setting, self.test_cmd] diff --git a/pytorch_blade/torch_blade/dynamo/__init__.py b/pytorch_blade/torch_blade/dynamo/__init__.py index 6f7a58bacd0..48c9840eca8 100644 --- a/pytorch_blade/torch_blade/dynamo/__init__.py +++ b/pytorch_blade/torch_blade/dynamo/__init__.py @@ -67,9 +67,10 @@ 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 - + print(fx_g.graph, flush=True) fx_g.graph.lint() fx_g.recompile() + f = torch.jit.script(fx_g) torch._C._jit_pass_remove_mutation(f.graph) if not is_training: diff --git a/pytorch_blade/torch_blade/mlir/disc_engine_conversion.py b/pytorch_blade/torch_blade/mlir/disc_engine_conversion.py index 49a83870d3d..abbbf449d19 100644 --- a/pytorch_blade/torch_blade/mlir/disc_engine_conversion.py +++ b/pytorch_blade/torch_blade/mlir/disc_engine_conversion.py @@ -237,4 +237,5 @@ def fusion_block(block): with tools.trust_tracing_shape(): fusion_block(graph) + print(graph, flush=True) _disc_engine_conversion(c_module) diff --git a/scripts/python/tao_build.py b/scripts/python/tao_build.py index 0de6f474eaa..5fb5f485f1a 100755 --- a/scripts/python/tao_build.py +++ b/scripts/python/tao_build.py @@ -327,11 +327,11 @@ def bazel_build(target, flag=""): flag = build_tao_compiler_add_flags_platform_alibaba(root, args, flag) - bazel_build(TARGET_TAO_COMPILER_MAIN, flag=flag) + #bazel_build(TARGET_TAO_COMPILER_MAIN, flag=flag) bazel_build(TARGET_DISC_OPT, flag=flag) # TODO:(fl237079) Support disc_replay for rocm version - if not args.rocm and not args.dcu: - bazel_build(TARGET_DISC_REPLAY, flag=flag) + #if not args.rocm and not args.dcu: + # bazel_build(TARGET_DISC_REPLAY, flag=flag) execute( "cp -f -p {}/tao/third_party/ptxas/10.2/ptxas ./bazel-bin/decoupling/".format( root diff --git a/tao_compiler/mlir/custom_ops/transpose_impl.cc b/tao_compiler/mlir/custom_ops/transpose_impl.cc index fb275a94f73..0087bd9b682 100644 --- a/tao_compiler/mlir/custom_ops/transpose_impl.cc +++ b/tao_compiler/mlir/custom_ops/transpose_impl.cc @@ -68,6 +68,8 @@ 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); +>>>>>>> update #endif } // namespace ral 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 index a3de94b8677..266ac1825f9 100755 --- a/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc +++ b/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc @@ -494,7 +494,7 @@ struct TransposeConverter : public OpRewritePattern { if (rank != 2 && rank != 3) return failure(); // only rewriter custom library when switch 1 and 2 dimensions of // a 3d tensor, that means permute = [0, 2, 1] - if (rank == 3 && (permutation[1] != 2 || permutation[2] != 1)) + if (rank == 3 && (permutation[1] != 2 && permutation[2] != 1)) return failure(); bool on_gpu = placement_utils::isGpuMemRef(op->getOperand(0)); // TODO: support other device diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc index 9614658f580..d88d0fc9b47 100644 --- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc +++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc @@ -487,13 +487,28 @@ bool isRank2ScalarReduction(Operation* op) { auto reduce_op = dyn_cast(op); if (!reduce_op || reduce_op.getDimensions().getNumElements() != 1) return false; - int rank = op->getOperand(2).getType().cast().getRank(); - // TODO(yancey): rewrite scalar reduction result to scalar tensor to avoid - // reshape to scalar tensor behand reduce op - Operation* reshapeOp = *op->getOperand(2).getUsers().begin(); - if (isa(reshapeOp) && - reshapeOp->getOperand(1).getType().cast().getRank() == 0) { - return true; + auto isRank0Tensor = [](Value v) -> bool { + return v.getType().cast().getRank() == 0; + }; + // TODO(yancey): it's a temporary solution to match scalar reduction, we need + // to erase the reshape op after scalar reduction, the result buffer of scalar + // reduction should be a scalar tensor instead of a <1xf32>tensor + { + Operation* reshapeOp = *op->getOperand(2).getUsers().begin(); + if (isa(reshapeOp) && isRank0Tensor(reshapeOp->getOperand(1))) { + return true; + } + } + { + Operation* convertOp = *op->getOperand(2).getUsers().begin(); + if (isa(convertOp)) { + auto resultBuffer = + convertOp->getOperand(convertOp->getNumOperands() - 1); + for (auto user : resultBuffer.getUsers()) { + if (isa(user) && isRank0Tensor(user->getOperand(1))) + return true; + } + } } return false; } @@ -1480,7 +1495,7 @@ bool BaseGpuFusionStrategy::isFusible(Operation* op) { !isRank2ScalarReduction(op))) // || isScalarReduction(op))) return false; - if (isa(op) && isRank2or3Transpose(op)) return false; + // if (isa(op) && isRank2or3Transpose(op)) return false; return BaseFusionStrategy::isFusible(op); } @@ -1515,18 +1530,18 @@ bool BaseGpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis, if (cnt >= 2) { return false; } - - if (has_rank2_col_reduction) { - const auto& results = target.getResults(); - auto ref_shape = getEffectiveShape(target, results[0]); - if (llvm::any_of(results, [&](Value result) { - auto op = target.findLastWriter(result); - return isa(op); - })) { - return false; + /* + if (has_rank2_col_reduction) { + const auto& results = target.getResults(); + auto ref_shape = getEffectiveShape(target, results[0]); + if (llvm::any_of(results, [&](Value result) { + auto op = target.findLastWriter(result); + return isa(op); + })) { + return false; + } } - } - + */ return BaseFusionStrategy::tryFuse(shapeAnalysis, lhs, rhs, target); } 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 59be2d138ee..b0da0c6e572 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 @@ -1259,31 +1259,7 @@ LogicalResult lowerWithScheduleRowReduction(ArrayRef, Operation*, int vector_size = 1) { return failure(); } -/* Row reduction with 1 round warp shuffle - * - * RowPerBlock = threads / warpSize; - * for (m = 0; m < rows; m += RowPerBlock) { - * for (n = 0; n < threads; ++n) { - * rowIdx = m + warpIdx; - * if (rowIdx < rows) { - * // intra-thread reduction - * sum = init_value; - * for (k = laneIdx; k < cols; k += warpSize) { - * sum += inputs[rowIdx][k]; - * } - * - * // inter-thread reduction via warp shuffle - * for (offset = warpSize / 2; offset > 0; offset /= 2) { - * sum += __shfl_xor(sum, offset); - * } - * - * // write to output - * if (laneIdx == 0) { - * outputs[rowIdx] = sum - * } - * } - * } - */ + /** * for (m = 0; m < block_nums; ++m) { * for (n = 0; n < block_size; ++n) { @@ -1457,38 +1433,7 @@ LogicalResult lowerWithScheduleParallelReduction( b.create(loc, *(for_op_k.getResults().begin() + idx), shared_mem_map[root_op], tid); } - // acc_value = for_op_k.getResult(0); - // b.create(loc, acc_value, shared_mem, tid); } - /* - { - Value init_value = b.create( - loc, cast(root_op).getInitValues()[0]); - SmallVector init_values{init_value}; - 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()); - auto lhs = dominant_op->getOperand(0); - SmallVector load_index({var_j, zero}); - Value data = createLoadOrUseCachedValue( - loc, &b, dominant_op, lhs, load_index, b.saveInsertionPoint()); - - SmallVector load_index2({b.create(loc, var_j, - block_dim), zero}); Value data1 = createLoadOrUseCachedValue( loc, &b, - dominant_op, lhs, load_index2, b.saveInsertionPoint()); - - Value sum = b.create(loc, data, data1); - acc = b.create(loc, acc, - *for_op_k.getRegionIterArgs().begin()); b.create(loc, - ValueRange{acc}); b.setInsertionPointAfter(for_op_k); acc_value = - for_op_k.getResult(0); b.create(loc, acc_value, shared_mem, - tid); - } - */ { Value var_j = nullptr; SmallVector init_values = {}; @@ -1515,18 +1460,6 @@ LogicalResult lowerWithScheduleParallelReduction( Value sum = (accum_factory[idx])(shm_val_1, shm_val_2); b.create(loc, sum, shared_mem_map[root_op], tid); } - /* - SmallVector multidim_load_index({tid}); - ValueRange load_index(multidim_load_index); - - SmallVector multidim_load_index2({b.create(loc, - tid, stride_val)}); ValueRange load_index2(multidim_load_index2); - - Value in_data = b.create(loc, shared_mem, load_index); - Value in_data1 = b.create(loc, shared_mem, load_index2); - Value sum = b.create(loc, in_data, in_data1); - b.create(loc, sum, shared_mem, tid); - */ b.create(loc); b.create(loc, yield_values); b.setInsertionPointAfter(if_tid_valid_op); @@ -1535,7 +1468,7 @@ LogicalResult lowerWithScheduleParallelReduction( b.create(loc); { // warp reduce - // if (tid < 16) + // if (tid < 32) scf::IfOp if_tid_valid_op = b.create( loc, /*resultTypes*/ TypeRange{}, b.create(loc, arith::CmpIPredicate::slt, tid, @@ -1557,14 +1490,6 @@ LogicalResult lowerWithScheduleParallelReduction( b.create(loc, sum, shared_mem_map[root_op], tid); b.create(loc); } - - // shm[tid] += shm[tid + stride]; - // Value idx2 = b.create(loc, tid, stride_val); - // Value shm_val_1 = b.create(loc, shared_mem, tid); - // Value shm_val_2 = b.create(loc, shared_mem, idx2); - // Value sum = b.create(loc, shm_val_1, shm_val_2); - // b.create(loc, sum, shared_mem_map[root_op], tid); - // b.create(loc); } b.setInsertionPointAfter(if_tid_valid_op); } @@ -1591,16 +1516,7 @@ LogicalResult lowerWithScheduleParallelReduction( getAtomicRMWKind(cast(root_op).getBody()), val, root_op->getOperand(2), ValueRange({zero})); } - /* - Value val = b.create(loc, shared_mem, tid); - Type root_element_type = getLhloOpsElementType(root_op); - b.create( - loc, root_element_type, - getAtomicRMWKind(cast(root_op).getBody()), - val, root_op->getOperand(2), ValueRange({zero})); - */ b.create(loc, yield_values); - b.setInsertionPointAfter(if_tid_zero_op); } b.setInsertionPointToEnd(local_workgroup.getBody()); @@ -4552,7 +4468,6 @@ LogicalResult HandleGpuFusionOp(OpBuilder& b, Operation* fusion, } break; case FusionType::kScalarReduction: { auto kname = getFusionFullName(fusion_op); - llvm::errs() << "kScalarReduction <" << kname << ">\n"; LogicalResult r = lowerWithScheduleParallelReduction( root_ops, dominant_op, fused_block); if (failed(r)) { From 05a7668ab1bffc901c745f3fc0cc1f11e9f040e4 Mon Sep 17 00:00:00 2001 From: yancey Date: Tue, 5 Mar 2024 18:58:19 +0800 Subject: [PATCH 03/10] update --- tao_compiler/mlir/custom_ops/transpose_impl.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/tao_compiler/mlir/custom_ops/transpose_impl.cc b/tao_compiler/mlir/custom_ops/transpose_impl.cc index 0087bd9b682..2f160d967a3 100644 --- a/tao_compiler/mlir/custom_ops/transpose_impl.cc +++ b/tao_compiler/mlir/custom_ops/transpose_impl.cc @@ -69,7 +69,6 @@ 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); ->>>>>>> update #endif } // namespace ral From 81cbd146e8375bc3751c62734130943aa4dcf537 Mon Sep 17 00:00:00 2001 From: yancey Date: Fri, 8 Mar 2024 15:23:03 +0800 Subject: [PATCH 04/10] update --- pytorch_blade/bazel_build.py | 14 ++++---- pytorch_blade/torch_blade/dynamo/__init__.py | 3 +- .../mlir/disc_engine_conversion.py | 1 - scripts/python/tao_build.py | 6 ++-- tao_compiler/mlir/disc/disc_compiler.cc | 2 ++ .../transforms/disc_lower_to_library_call.cc | 6 +++- .../disc/transforms/element_type_converter.cc | 4 +-- .../mlir/disc/transforms/fusion_utils.cc | 34 +++++-------------- .../mlir/disc/transforms/fusion_utils.h | 2 ++ .../transforms/fusion_utils_stitch_gpu.cc | 19 +++++++---- .../lhlo_legalize_roots_to_loops.cc | 9 +++-- 11 files changed, 46 insertions(+), 54 deletions(-) mode change 100755 => 100644 tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc diff --git a/pytorch_blade/bazel_build.py b/pytorch_blade/bazel_build.py index 32f2aea10dc..57c3fc7b4bc 100644 --- a/pytorch_blade/bazel_build.py +++ b/pytorch_blade/bazel_build.py @@ -68,9 +68,9 @@ def __init__(self, *args, **kwargs): "@org_disc_compiler//mlir/custom_ops:libdisc_custom_ops.so", "//pytorch_blade:libtorch_blade.so", "//pytorch_blade:_torch_blade.so", - #"//tests/mhlo/torch-mlir-opt:torch-mlir-opt", - #"//tests/torchscript:shape_analysis_tool", - #"//tests/torch-disc-pdll:torch-disc-pdll", + "//tests/mhlo/torch-mlir-opt:torch-mlir-opt", + "//tests/torchscript:shape_analysis_tool", + "//tests/torch-disc-pdll:torch-disc-pdll", ] torch_major_version, torch_minor_version = self.torch_version.split(".")[:2] @@ -265,15 +265,15 @@ def test(self): self.test_suites = [ "@org_disc_compiler//mlir/ral:collective_ops_test", - #"//tests/mhlo/...", - #"//pytorch_blade:torch_blade_test_suite", - #"//tests/torch-disc-pdll/tests/...", + "//tests/mhlo/...", + "//pytorch_blade:torch_blade_test_suite", + "//tests/torch-disc-pdll/tests/...", ] if (self.torch_major_version, self.torch_minor_version) > (1, 6): # torchscript graph ir parser changed after torch 1.6. # We will not test torchscript graph ir before torch 1.6 - #self.test_suites.append("//tests/torchscript/...") + self.test_suites.append("//tests/torchscript/...") pass test_cmd = " ".join( diff --git a/pytorch_blade/torch_blade/dynamo/__init__.py b/pytorch_blade/torch_blade/dynamo/__init__.py index 48c9840eca8..0ed4c19f6b5 100644 --- a/pytorch_blade/torch_blade/dynamo/__init__.py +++ b/pytorch_blade/torch_blade/dynamo/__init__.py @@ -67,10 +67,9 @@ 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 - print(fx_g.graph, flush=True) fx_g.graph.lint() fx_g.recompile() - + f = torch.jit.script(fx_g) torch._C._jit_pass_remove_mutation(f.graph) if not is_training: diff --git a/pytorch_blade/torch_blade/mlir/disc_engine_conversion.py b/pytorch_blade/torch_blade/mlir/disc_engine_conversion.py index abbbf449d19..49a83870d3d 100644 --- a/pytorch_blade/torch_blade/mlir/disc_engine_conversion.py +++ b/pytorch_blade/torch_blade/mlir/disc_engine_conversion.py @@ -237,5 +237,4 @@ def fusion_block(block): with tools.trust_tracing_shape(): fusion_block(graph) - print(graph, flush=True) _disc_engine_conversion(c_module) diff --git a/scripts/python/tao_build.py b/scripts/python/tao_build.py index 5fb5f485f1a..0de6f474eaa 100755 --- a/scripts/python/tao_build.py +++ b/scripts/python/tao_build.py @@ -327,11 +327,11 @@ def bazel_build(target, flag=""): flag = build_tao_compiler_add_flags_platform_alibaba(root, args, flag) - #bazel_build(TARGET_TAO_COMPILER_MAIN, flag=flag) + bazel_build(TARGET_TAO_COMPILER_MAIN, flag=flag) bazel_build(TARGET_DISC_OPT, flag=flag) # TODO:(fl237079) Support disc_replay for rocm version - #if not args.rocm and not args.dcu: - # bazel_build(TARGET_DISC_REPLAY, flag=flag) + if not args.rocm and not args.dcu: + bazel_build(TARGET_DISC_REPLAY, flag=flag) execute( "cp -f -p {}/tao/third_party/ptxas/10.2/ptxas ./bazel-bin/decoupling/".format( root diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc index 17770b0845c..7d51ea0b0c7 100644 --- a/tao_compiler/mlir/disc/disc_compiler.cc +++ b/tao_compiler/mlir/disc/disc_compiler.cc @@ -601,6 +601,8 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) { pm.addNestedPass(createCSEPass()); pm.addNestedPass( createCanonicalizerPass(cano_rewrite_config, disablePatterns)); + // TODO(yancey): enable this pass after fix WAW issue in scalar reduction + // codegen template // pm.addNestedPass(disc_ral::createDiscMemRefCSEPass()); // convert linearizeOp/delinearizeOp to std dialect. pm.addNestedPass(disc_ral::createDiscConvertShapeToStandardPass()); 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 266ac1825f9..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,12 +489,16 @@ 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(); // only rewriter custom library when switch 1 and 2 dimensions of // a 3d tensor, that means permute = [0, 2, 1] - if (rank == 3 && (permutation[1] != 2 && permutation[2] != 1)) + if (rank == 3 && (permutation[1] != 2 || permutation[2] != 1)) return failure(); bool on_gpu = placement_utils::isGpuMemRef(op->getOperand(0)); // TODO: support other device 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 d88d0fc9b47..2beeb6e5fd3 100644 --- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc +++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc @@ -485,30 +485,10 @@ bool isRowReduction(Operation* op) { bool isRank2ScalarReduction(Operation* op) { auto reduce_op = dyn_cast(op); - if (!reduce_op || reduce_op.getDimensions().getNumElements() != 1) + if (!reduce_op || reduce_op.getDimensions().getNumElements() != 2) return false; - auto isRank0Tensor = [](Value v) -> bool { - return v.getType().cast().getRank() == 0; - }; - // TODO(yancey): it's a temporary solution to match scalar reduction, we need - // to erase the reshape op after scalar reduction, the result buffer of scalar - // reduction should be a scalar tensor instead of a <1xf32>tensor - { - Operation* reshapeOp = *op->getOperand(2).getUsers().begin(); - if (isa(reshapeOp) && isRank0Tensor(reshapeOp->getOperand(1))) { - return true; - } - } - { - Operation* convertOp = *op->getOperand(2).getUsers().begin(); - if (isa(convertOp)) { - auto resultBuffer = - convertOp->getOperand(convertOp->getNumOperands() - 1); - for (auto user : resultBuffer.getUsers()) { - if (isa(user) && isRank0Tensor(user->getOperand(1))) - return true; - } - } + if (auto ty = op->getOperand(2).getType().dyn_cast()) { + return ty.getRank() == 0; } return false; } @@ -521,8 +501,7 @@ bool isRank2ColReduction(Operation* op) { int rank = op->getOperand(0).getType().cast().getRank(); auto dimensions = reduce_op.getDimensions().getValues(); - return ((*dimensions.begin() == 0) && (rank == 2)) && - !isRank2ScalarReduction(op); + return (*dimensions.begin() == 0) && (rank == 2); } // Return true if this op is a rank-2 transpose @@ -813,6 +792,9 @@ FusionPattern::FusionPattern(lmhlo::FusionOp op, ShapeAnalysis* shape_analysis) 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); } @@ -1492,7 +1474,7 @@ bool BaseGpuFusionStrategy::isFusible(Operation* op) { // Only rank-2 tensor -> rank-1 tensor reduction are supported now. if (isa(op) && (!isRank2RowReduction(op) && !isRank2ColReduction(op) && - !isRank2ScalarReduction(op))) // || isScalarReduction(op))) + !isRank2ScalarReduction(op))) return false; // if (isa(op) && isRank2or3Transpose(op)) return false; diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.h b/tao_compiler/mlir/disc/transforms/fusion_utils.h index b58bbeeefc5..d02857d1764 100644 --- a/tao_compiler/mlir/disc/transforms/fusion_utils.h +++ b/tao_compiler/mlir/disc/transforms/fusion_utils.h @@ -157,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 75c4644e2fa..69f3f1b8931 100644 --- a/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc +++ b/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc @@ -18,9 +18,13 @@ namespace disc_ral { ////////////////////// Stitch GPU FusionStrategy Implemenation ///////// //////////////////////////////////////////////////////////////////////// bool isScalarReduction(Operation* op) { - auto reduce_op = dyn_cast(op); - if (!reduce_op || reduce_op.getDimensions().getNumElements() != 1) - return false; + if (auto reduce_op = dyn_cast(op)) { + llvm::dbgs() << "reduce op:" << *reduce_op << "\n"; + return reduce_op->getOperand(2).getType().cast().getRank() == 0; + } + return false; + // if (!reduce_op || reduce_op.getDimensions().getNumElements() != 1) + // return false; int rank = op->getOperand(2).getType().cast().getRank(); // TODO(yancey): rewrite scalar reduction result to scalar tensor to avoid // reshape to scalar tensor behand reduce op @@ -42,7 +46,7 @@ bool findValidReductionOps(FusionPatternBase& target, if (!isa(op)) continue; if (isRank2RowReduction(op)) { row_reductions.push_back(op); - } else if (isRank2ColReduction(op) || isScalarReduction(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); @@ -55,7 +59,7 @@ bool findValidReductionOps(FusionPatternBase& target, } } } - if (isScalarReduction(op)) { + if (isRank2ScalarReduction(op)) { scalar_reductions.push_back(op); } else { col_reductions.push_back(op); @@ -83,8 +87,9 @@ 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 isScalarReduction(op); }); + 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; 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 b0da0c6e572..014da0be968 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 @@ -1285,7 +1285,7 @@ LogicalResult lowerWithScheduleRowReduction(ArrayRef, Operation*, LogicalResult lowerWithScheduleParallelReduction( ArrayRef root_ops, Operation* dominant_op, Block* parent, const ShapeAnalysis* shape_analysis = nullptr, int vector_size = 1) { - if (!isRank2ColReduction(dominant_op)) { + if (!isRank2ScalarReduction(dominant_op)) { return failure(); } // Create helper Values @@ -1293,10 +1293,9 @@ LogicalResult lowerWithScheduleParallelReduction( std::copy_if( root_ops.begin(), root_ops.end(), std::back_inserter(scalar_reduction_roots), - [](Operation* operation) { return isRank2ColReduction(operation); }); + [](Operation* operation) { return isRank2ScalarReduction(operation); }); auto root_op = scalar_reduction_roots.back(); const int thread_per_block = getCTASize(dominant_op); - ; Location loc = dominant_op->getLoc(); OpBuilder b(root_ops.back()); @@ -1393,7 +1392,7 @@ LogicalResult lowerWithScheduleParallelReduction( b.setInsertionPointToStart(for_op_k.getBody()); int scalar_red_root_op_idx = 0; for (auto* root_op : root_ops) { - if (isRank2ColReduction(root_op)) { + if (isRank2ScalarReduction(root_op)) { auto lhs = root_op->getOperands().begin(); SmallVector load_index({i, zero}); Value data = createLoadOrUseCachedValue( @@ -1514,7 +1513,7 @@ LogicalResult lowerWithScheduleParallelReduction( b.create( loc, root_element_type, getAtomicRMWKind(cast(root_op).getBody()), val, - root_op->getOperand(2), ValueRange({zero})); + root_op->getOperand(2), ValueRange({})); } b.create(loc, yield_values); b.setInsertionPointAfter(if_tid_zero_op); From ecf7705f36ae2ac748ee8a8e4d907951c9022f17 Mon Sep 17 00:00:00 2001 From: yancey Date: Fri, 8 Mar 2024 15:25:57 +0800 Subject: [PATCH 05/10] update --- pytorch_blade/bazel_build.py | 2 -- pytorch_blade/torch_blade/dynamo/__init__.py | 1 - 2 files changed, 3 deletions(-) diff --git a/pytorch_blade/bazel_build.py b/pytorch_blade/bazel_build.py index 57c3fc7b4bc..e41ebaba6d3 100644 --- a/pytorch_blade/bazel_build.py +++ b/pytorch_blade/bazel_build.py @@ -264,7 +264,6 @@ def test(self): env["GCC_HOST_COMPILER_PATH"] = env.get("GCC_HOST_COMPILER_PATH", which("gcc")) self.test_suites = [ - "@org_disc_compiler//mlir/ral:collective_ops_test", "//tests/mhlo/...", "//pytorch_blade:torch_blade_test_suite", "//tests/torch-disc-pdll/tests/...", @@ -274,7 +273,6 @@ def test(self): # torchscript graph ir parser changed after torch 1.6. # We will not test torchscript graph ir before torch 1.6 self.test_suites.append("//tests/torchscript/...") - pass test_cmd = " ".join( [self.shell_setting, self.test_cmd] diff --git a/pytorch_blade/torch_blade/dynamo/__init__.py b/pytorch_blade/torch_blade/dynamo/__init__.py index 0ed4c19f6b5..ef484a13b24 100644 --- a/pytorch_blade/torch_blade/dynamo/__init__.py +++ b/pytorch_blade/torch_blade/dynamo/__init__.py @@ -69,7 +69,6 @@ def _disc_compile(fx_g: fx.GraphModule, inps, use_ts=False, is_training=True) -> node.kwargs = new_kwargs fx_g.graph.lint() fx_g.recompile() - f = torch.jit.script(fx_g) torch._C._jit_pass_remove_mutation(f.graph) if not is_training: From 3b5a598893ebf973890952263c2548ea9c734eed Mon Sep 17 00:00:00 2001 From: yancey Date: Tue, 12 Mar 2024 11:34:46 +0800 Subject: [PATCH 06/10] fix ut --- tao_compiler/mlir/disc/disc_compiler.cc | 3 +- .../mlir/disc/transforms/fusion_utils.cc | 30 +++++++++---------- .../transforms/fusion_utils_stitch_gpu.cc | 18 ----------- .../lhlo_legalize_roots_to_loops.cc | 7 +---- .../transforms/parallel_loop_collapsing.cc | 3 +- 5 files changed, 19 insertions(+), 42 deletions(-) diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc index 7d51ea0b0c7..1646fb8078d 100644 --- a/tao_compiler/mlir/disc/disc_compiler.cc +++ b/tao_compiler/mlir/disc/disc_compiler.cc @@ -603,7 +603,8 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) { createCanonicalizerPass(cano_rewrite_config, disablePatterns)); // TODO(yancey): enable this pass after fix WAW issue in scalar reduction // codegen template - // pm.addNestedPass(disc_ral::createDiscMemRefCSEPass()); + 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/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc index 2beeb6e5fd3..cf0cbeb314f 100644 --- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc +++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc @@ -501,7 +501,7 @@ bool isRank2ColReduction(Operation* op) { int rank = op->getOperand(0).getType().cast().getRank(); auto dimensions = reduce_op.getDimensions().getValues(); - return (*dimensions.begin() == 0) && (rank == 2); + return ((*dimensions.begin() == 0) && (rank == 2)); } // Return true if this op is a rank-2 transpose @@ -568,10 +568,8 @@ 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; @@ -1512,18 +1510,18 @@ bool BaseGpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis, if (cnt >= 2) { return false; } - /* - if (has_rank2_col_reduction) { - const auto& results = target.getResults(); - auto ref_shape = getEffectiveShape(target, results[0]); - if (llvm::any_of(results, [&](Value result) { - auto op = target.findLastWriter(result); - return isa(op); - })) { - return false; - } + + if (has_rank2_col_reduction) { + const auto& results = target.getResults(); + auto ref_shape = getEffectiveShape(target, results[0]); + if (llvm::any_of(results, [&](Value result) { + auto op = target.findLastWriter(result); + return isa(op); + })) { + return false; } - */ + } + return BaseFusionStrategy::tryFuse(shapeAnalysis, lhs, rhs, target); } 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 69f3f1b8931..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,24 +17,6 @@ namespace disc_ral { ////////////////////// Stitch GPU FusionStrategy Implemenation ///////// //////////////////////////////////////////////////////////////////////// -bool isScalarReduction(Operation* op) { - if (auto reduce_op = dyn_cast(op)) { - llvm::dbgs() << "reduce op:" << *reduce_op << "\n"; - return reduce_op->getOperand(2).getType().cast().getRank() == 0; - } - return false; - // if (!reduce_op || reduce_op.getDimensions().getNumElements() != 1) - // return false; - int rank = op->getOperand(2).getType().cast().getRank(); - // TODO(yancey): rewrite scalar reduction result to scalar tensor to avoid - // reshape to scalar tensor behand reduce op - Operation* reshapeOp = *op->getOperand(2).getUsers().begin(); - if (reshapeOp && isa(reshapeOp) && - reshapeOp->getOperand(1).getType().cast().getRank() == 0) { - return true; - } - return false; -} bool findValidReductionOps(FusionPatternBase& target, SmallVectorImpl& row_reductions, SmallVectorImpl& col_reductions, 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 014da0be968..3576ace5c1a 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 @@ -5856,12 +5856,7 @@ LogicalResult lowerWithSchedulekInputCPU( b.setInsertionPointAfter(reductionForOp); // remove the root_op if it has no other users except the memref - if (non_fusion) { - for (Operation* root_op : root_ops) root_op->erase(); - } else { - assert(parent != nullptr && "Parent must be provided for fusion lowering"); - cleanUnusedLhloOps(parent); - } + for (Operation* root_op : root_ops) root_op->erase(); return success(); } diff --git a/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc b/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc index b24bb33e27f..400c170f765 100644 --- a/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc +++ b/tao_compiler/mlir/disc/transforms/parallel_loop_collapsing.cc @@ -49,7 +49,8 @@ struct ParallelLoopCollapsing fusion->getAttrOfType(kDiscFusionTypeAttrName); if (fusionTypeAttr && (fusionTypeAttr.getValue() == "kStitch" || - fusionTypeAttr.getValue() == "kScalarReduction")) { + fusionTypeAttr.getValue() == "kScalarReduction" || + fusionTypeAttr.getValue() == "kTransform")) { continue; } } From 3a3de95d97e46527a25099b431a8cc8cbb2646ed Mon Sep 17 00:00:00 2001 From: YanXu Date: Tue, 19 Mar 2024 20:24:44 +0800 Subject: [PATCH 07/10] fix ut --- .../lhlo_legalize_roots_to_loops.cc | 51 +++++++++++-------- tf_community | 2 +- 2 files changed, 31 insertions(+), 22 deletions(-) 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 3576ace5c1a..70766a0eeda 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 @@ -1091,7 +1091,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); } } @@ -1295,7 +1295,7 @@ LogicalResult lowerWithScheduleParallelReduction( std::back_inserter(scalar_reduction_roots), [](Operation* operation) { return isRank2ScalarReduction(operation); }); auto root_op = scalar_reduction_roots.back(); - const int thread_per_block = getCTASize(dominant_op); + const int thread_per_block = 256; Location loc = dominant_op->getLoc(); OpBuilder b(root_ops.back()); @@ -1305,10 +1305,10 @@ LogicalResult lowerWithScheduleParallelReduction( Value one = b.create(loc, 1); Value two = b.create(loc, 2); auto elemFloatType = getLhloOpsElementType(root_op).cast(); - Value zero_f = b.create( - loc, llvm::APFloat(0.0f), elemFloatType); // b.getF32Type()); - Value one_f = b.create( - loc, llvm::APFloat(1.0f), elemFloatType); // b.getF32Type()); + Value zero_f = b.create( + loc, b.getFloatAttr(getLhloOpsElementType(root_op), 0)); + Value one_f = b.create( + loc, b.getFloatAttr(getLhloOpsElementType(root_op), 1)); // Start to emit. Value num_blocks = b.create(loc, 1024); @@ -1334,8 +1334,8 @@ LogicalResult lowerWithScheduleParallelReduction( b.create(loc, b.getIndexType(), gpu::Dimension::x); Value grid_dim = b.create(loc, b.getIndexType(), gpu::Dimension::x); - tid = b.create(loc, tid, block_dim); - // i = blockIdx.x * block_size * 2 + tid; + // tid = b.create(loc, tid, block_dim); + // i = blockIdx.x * block_size * 2 + tid; Value i = b.create( loc, b.create( @@ -1349,11 +1349,7 @@ LogicalResult lowerWithScheduleParallelReduction( Value n = b.create(loc, lhs, zero); Value m = b.create(loc, lhs, one); Value mn = b.create(loc, m, n); - // __shared__ float shm[block_size]; - auto shared_mem = createSharedMemory(b, loc, thread_per_block, - getLhloOpsElementType(dominant_op)); - Value acc_value; - // fused accumulation + // acc: init_values[num_col_reductions] SmallVector accum_factory( scalar_reduction_roots.size()); @@ -1394,14 +1390,29 @@ LogicalResult lowerWithScheduleParallelReduction( for (auto* root_op : root_ops) { if (isRank2ScalarReduction(root_op)) { auto lhs = root_op->getOperands().begin(); - SmallVector load_index({i, zero}); + SmallVector load_index({var_j, zero}); Value data = createLoadOrUseCachedValue( loc, &b, root_op, *lhs, load_index, b.saveInsertionPoint()); - SmallVector load_index2( - {b.create(loc, i, block_dim), zero}); + Value index2 = b.create(loc, var_j, block_dim); + // if (i + grid_size < n) + scf::IfOp if_tid_valid_op = b.create( + loc, /*resultTypes*/ init_values_types, + b.create(loc, arith::CmpIPredicate::slt, index2, n), + /*hasElseRegion*/ true); + if_tid_valid_op.getThenRegion().front().clear(); + if_tid_valid_op.getElseRegion().front().clear(); + b.setInsertionPointToStart(&if_tid_valid_op.getThenRegion().front()); + SmallVector load_index2({index2, zero}); Value data1 = createLoadOrUseCachedValue( loc, &b, root_op, *lhs, load_index2, b.saveInsertionPoint()); - Value sum = (accum_factory[scalar_red_root_op_idx])(data, data1); + b.setInsertionPointToEnd(&if_tid_valid_op.getThenRegion().front()); + b.create(loc, data1); + b.setInsertionPointToStart(&if_tid_valid_op.getElseRegion().front()); + b.create(loc, zero_f); + b.setInsertionPointAfter(if_tid_valid_op); + Value sum = (accum_factory[scalar_red_root_op_idx])( + data, if_tid_valid_op.getResults().front()); + auto acc = (accum_factory[scalar_red_root_op_idx])( *(for_op_k.getRegionIterArgs().begin() + scalar_red_root_op_idx), sum); @@ -1425,7 +1436,7 @@ LogicalResult lowerWithScheduleParallelReduction( } b.create(loc, yield_values_for_if); b.setInsertionPointAfter(for_op_k); - + b.create(loc); for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { Operation* root_op = root_pair.value(); int idx = root_pair.index(); @@ -1434,7 +1445,6 @@ LogicalResult lowerWithScheduleParallelReduction( } } { - Value var_j = nullptr; SmallVector init_values = {}; for (int stride = 128; stride > 16; stride /= 2) { b.create(loc); @@ -1464,7 +1474,6 @@ LogicalResult lowerWithScheduleParallelReduction( b.setInsertionPointAfter(if_tid_valid_op); } } - b.create(loc); { // warp reduce // if (tid < 32) @@ -1508,7 +1517,7 @@ LogicalResult lowerWithScheduleParallelReduction( for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { Operation* root_op = root_pair.value(); int idx = root_pair.index(); - Value val = b.create(loc, shared_mem_map[root_op], tid); + Value val = b.create(loc, shared_mem_map[root_op], zero); Type root_element_type = getLhloOpsElementType(root_op); b.create( loc, root_element_type, 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 From 833fbc516f92efe52ed035b901b569c6d6fb8d01 Mon Sep 17 00:00:00 2001 From: YanXu Date: Wed, 20 Mar 2024 20:22:15 +0800 Subject: [PATCH 08/10] fix ut --- .../mlir/disc/tests/mlir_feature_test.cc | 1 - .../lhlo_legalize_roots_to_loops.cc | 24 ++++++++----------- 2 files changed, 10 insertions(+), 15 deletions(-) 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/lhlo_legalize_roots_to_loops.cc b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc index 70766a0eeda..bfe7b9acdd4 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 @@ -1304,13 +1304,6 @@ LogicalResult lowerWithScheduleParallelReduction( Value zero = b.create(loc, 0); Value one = b.create(loc, 1); Value two = b.create(loc, 2); - auto elemFloatType = getLhloOpsElementType(root_op).cast(); - Value zero_f = b.create( - loc, b.getFloatAttr(getLhloOpsElementType(root_op), 0)); - Value one_f = b.create( - loc, b.getFloatAttr(getLhloOpsElementType(root_op), 1)); - - // Start to emit. Value num_blocks = b.create(loc, 1024); Value block_size = b.create(loc, 256); @@ -1334,8 +1327,7 @@ LogicalResult lowerWithScheduleParallelReduction( b.create(loc, b.getIndexType(), gpu::Dimension::x); Value grid_dim = b.create(loc, b.getIndexType(), gpu::Dimension::x); - // tid = b.create(loc, tid, block_dim); - // i = blockIdx.x * block_size * 2 + tid; + // i = blockIdx.x * block_size * 2 + tid; Value i = b.create( loc, b.create( @@ -1394,6 +1386,8 @@ LogicalResult lowerWithScheduleParallelReduction( Value data = createLoadOrUseCachedValue( loc, &b, root_op, *lhs, load_index, b.saveInsertionPoint()); Value index2 = b.create(loc, var_j, block_dim); + Value iter_value = + *(for_op_k.getRegionIterArgs().begin() + scalar_red_root_op_idx); // if (i + grid_size < n) scf::IfOp if_tid_valid_op = b.create( loc, /*resultTypes*/ init_values_types, @@ -1405,17 +1399,19 @@ LogicalResult lowerWithScheduleParallelReduction( SmallVector load_index2({index2, zero}); Value data1 = createLoadOrUseCachedValue( loc, &b, root_op, *lhs, load_index2, b.saveInsertionPoint()); + data1 = (accum_factory[scalar_red_root_op_idx])(iter_value, data1); b.setInsertionPointToEnd(&if_tid_valid_op.getThenRegion().front()); b.create(loc, data1); b.setInsertionPointToStart(&if_tid_valid_op.getElseRegion().front()); - b.create(loc, zero_f); + b.create(loc, iter_value); + // loc, cast(root_op).getInitValues().front()); b.setInsertionPointAfter(if_tid_valid_op); - Value sum = (accum_factory[scalar_red_root_op_idx])( + Value acc = (accum_factory[scalar_red_root_op_idx])( data, if_tid_valid_op.getResults().front()); - auto acc = (accum_factory[scalar_red_root_op_idx])( - *(for_op_k.getRegionIterArgs().begin() + scalar_red_root_op_idx), - sum); + // acc = (accum_factory[scalar_red_root_op_idx])( + // *(for_op_k.getRegionIterArgs().begin() + scalar_red_root_op_idx), + // acc); yield_values_for_if.push_back(acc); scalar_red_root_op_idx++; } else if (isa(root_op)) { From cedce1569566cd65f8f87ea9e70c37858e40cb6a Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Tue, 31 Dec 2024 00:44:07 +0800 Subject: [PATCH 09/10] update --- .../lhlo_legalize_roots_to_loops.cc | 288 +++++++++++------- 1 file changed, 184 insertions(+), 104 deletions(-) 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 bfe7b9acdd4..c39c9329bd2 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, 1); Value num_elements = emitNumElementsComputation(b, loc, first_root); Value var; @@ -1273,15 +1274,23 @@ LogicalResult lowerWithScheduleRowReduction(ArrayRef, Operation*, * __syncthreads(); * for (int stride = block_size / 2; stride > 0; stride /= 2) { * if (tid < stride) { - * shm[tid] += shm[tid + stride]; + * shm[tid] += shm[tid + stride]; + * } * __syncthreads(); * } - * if (tid == 0) { - * atomicAdd(&g_output[0], shm[0]); + * 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) { @@ -1295,6 +1304,9 @@ LogicalResult lowerWithScheduleParallelReduction( 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()); @@ -1304,8 +1316,19 @@ LogicalResult lowerWithScheduleParallelReduction( Value zero = b.create(loc, 0); Value one = b.create(loc, 1); Value two = b.create(loc, 2); - Value num_blocks = b.create(loc, 1024); - Value block_size = b.create(loc, 256); + 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}), @@ -1316,9 +1339,7 @@ LogicalResult lowerWithScheduleParallelReduction( loc, SmallVector({zero}), SmallVector({block_size}), SmallVector({one}), SmallVector({}), /*bodyBuilderFn=*/nullptr); - local_workgroup.getBody()->clear(); b.setInsertionPointToStart(local_workgroup.getBody()); - Value block_dim = b.create(loc, b.getIndexType(), gpu::Dimension::x); Value block_idx = @@ -1327,22 +1348,23 @@ LogicalResult lowerWithScheduleParallelReduction( 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_dim), two), + 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_dim), two); + 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_col_reductions] + // acc: init_values[num_scalar_reductions] SmallVector accum_factory( scalar_reduction_roots.size()); SmallVector init_values(scalar_reduction_roots.size()); @@ -1379,51 +1401,42 @@ LogicalResult lowerWithScheduleParallelReduction( 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(); - SmallVector load_index({var_j, zero}); Value data = createLoadOrUseCachedValue( loc, &b, root_op, *lhs, load_index, b.saveInsertionPoint()); - Value index2 = b.create(loc, var_j, block_dim); + Value iter_value = *(for_op_k.getRegionIterArgs().begin() + scalar_red_root_op_idx); - // if (i + grid_size < n) - scf::IfOp if_tid_valid_op = b.create( - loc, /*resultTypes*/ init_values_types, - b.create(loc, arith::CmpIPredicate::slt, index2, n), - /*hasElseRegion*/ true); - if_tid_valid_op.getThenRegion().front().clear(); - if_tid_valid_op.getElseRegion().front().clear(); - b.setInsertionPointToStart(&if_tid_valid_op.getThenRegion().front()); - SmallVector load_index2({index2, zero}); - Value data1 = createLoadOrUseCachedValue( + 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()); - data1 = (accum_factory[scalar_red_root_op_idx])(iter_value, data1); - b.setInsertionPointToEnd(&if_tid_valid_op.getThenRegion().front()); - b.create(loc, data1); - b.setInsertionPointToStart(&if_tid_valid_op.getElseRegion().front()); - b.create(loc, iter_value); - // loc, cast(root_op).getInitValues().front()); - b.setInsertionPointAfter(if_tid_valid_op); - Value acc = (accum_factory[scalar_red_root_op_idx])( - data, if_tid_valid_op.getResults().front()); - - // acc = (accum_factory[scalar_red_root_op_idx])( - // *(for_op_k.getRegionIterArgs().begin() + scalar_red_root_op_idx), - // acc); - yield_values_for_if.push_back(acc); + 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 if (isa(root_op)) { - auto dominant_shape = getShapeValues(&b, dominant_op->getOperand(0)); - Value linear_index = calcLinearIndex(&b, loc, i, dominant_shape); - auto root_shape = getShapeValues(&b, root_op->getOperand(0)); - auto mapped_index = - calcMultiDimIndex(&b, loc, linear_index, root_shape); - emitNotToVectorReduction(b, loc, root_op, mapped_index); } else { auto dominant_shape = getShapeValues(&b, dominant_op->getOperand(0)); - Value linear_index = calcLinearIndex(&b, loc, i, dominant_shape); + Value linear_index = + calcLinearIndex(&b, loc, load_index, dominant_shape); if (!succeeded( lowerHelper(b, loc, root_op, linear_index, shape_analysis))) { return failure(); @@ -1432,18 +1445,18 @@ LogicalResult lowerWithScheduleParallelReduction( } b.create(loc, yield_values_for_if); b.setInsertionPointAfter(for_op_k); - b.create(loc); 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 > 16; stride /= 2) { - b.create(loc); + 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( @@ -1465,66 +1478,106 @@ LogicalResult lowerWithScheduleParallelReduction( Value sum = (accum_factory[idx])(shm_val_1, shm_val_2); b.create(loc, sum, shared_mem_map[root_op], tid); } - b.create(loc); 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()); { - // warp reduce - // if (tid < 32) - scf::IfOp if_tid_valid_op = b.create( - loc, /*resultTypes*/ TypeRange{}, - b.create(loc, arith::CmpIPredicate::slt, tid, - b.create(loc, 32)), - /*hasElseRegion*/ false); - b.setInsertionPointToStart(&if_tid_valid_op.getThenRegion().front()); - SmallVector yield_values; - for (int stride = 16; stride > 0; stride /= 2) { - Value stride_val = b.create(loc, stride); - for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { + 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(); - 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); + 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.setInsertionPointAfter(if_tid_valid_op); - } - b.create(loc); - - { - // if (tid == 0) + 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()); - SmallVector yield_values; for (auto root_pair : llvm::enumerate(scalar_reduction_roots)) { Operation* root_op = root_pair.value(); int idx = root_pair.index(); - Value val = b.create(loc, shared_mem_map[root_op], zero); + Value val = if_tid_valid_warp.getResults()[idx]; Type root_element_type = getLhloOpsElementType(root_op); + b.create( - loc, root_element_type, + loc, shuffle_type[idx], getAtomicRMWKind(cast(root_op).getBody()), val, root_op->getOperand(2), ValueRange({})); } - b.create(loc, yield_values); + b.create(loc, ValueRange({})); b.setInsertionPointAfter(if_tid_zero_op); } b.setInsertionPointToEnd(local_workgroup.getBody()); - b.create(loc, ValueRange({})); b.setInsertionPointAfter(global_workgroup); if (parent == nullptr) { for (Operation* root_op : root_ops) root_op->erase(); @@ -2169,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(); } @@ -2453,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]; // } @@ -2901,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]); } @@ -3328,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); @@ -3541,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) { @@ -3581,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); // } @@ -3929,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); } @@ -4012,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], @@ -5221,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(); @@ -5998,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) && @@ -6044,6 +6105,25 @@ 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. From befe9324ec2dd74524c9ab9e06e81917f3a71085 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Tue, 31 Dec 2024 00:51:27 +0800 Subject: [PATCH 10/10] update --- .../mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 c39c9329bd2..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,7 +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, 1); + Value two = b.create(loc, 2); Value num_elements = emitNumElementsComputation(b, loc, first_root); Value var;