| #include "platforms/xla/mosaic/dialect/llo/llo_builder.h" |
| |
| #include <cstdint> |
| #include <memory> |
| #include <optional> |
| #include <string_view> |
| |
| #include "platforms/xla/mosaic/dialect/llo/llo_dialect.h" |
| #include "platforms/xla/service/jellyfish/dma_strides.h" |
| #include "platforms/xla/service/jellyfish/execution_profiler.h" |
| #include "platforms/xla/service/jellyfish/execution_profiler_traceme.h" |
| #include "platforms/xla/service/jellyfish/llo_instruction.h" |
| #include "platforms/xla/service/jellyfish/llo_predicated_region.h" |
| #include "platforms/xla/service/jellyfish/llo_region.h" |
| #include "platforms/xla/service/jellyfish/llo_region_builder.h" |
| #include "platforms/xla/service/jellyfish/llo_value.h" |
| #include "platforms/xla/service/jellyfish/lowering/fusion_util.h" |
| #include "platforms/xla/service/jellyfish/lowering/net_util.h" |
| #include "platforms/xla/service/jellyfish/target.h" |
| #include "third_party/absl/log/check.h" |
| #include "third_party/absl/log/log.h" |
| #include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/STLExtras.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/SCF/IR/SCF.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/BuiltinTypes.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Types.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Value.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/Support/LLVM.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/Support/LogicalResult.h" |
| #include "third_party/tensorflow/compiler/xla/mlir/utils/type_util.h" |
| #include "third_party/tensorflow/compiler/xla/primitive_util.h" |
| #include "third_party/tensorflow/compiler/xla/shape_util.h" |
| |
| namespace mlir { |
| namespace llo { |
| |
| namespace { |
| |
| ::xla::jellyfish::LloValue* GetLloValue( |
| mlir::Value val, |
| const DenseMap<mlir::Value, ::xla::jellyfish::LloValue*>& value_map) { |
| auto it = value_map.find(val); |
| assert(it != value_map.end()); |
| return it->second; |
| } |
| |
| template <typename Op> |
| ::xla::jellyfish::LloRegionBuilder::CoreLocation GetCoreLocation( |
| BuilderContext* ctx, ::xla::jellyfish::LloRegionBuilder& b, Op op, |
| std::string_view tag, |
| const DenseMap<mlir::Value, ::xla::jellyfish::LloValue*>& value_map) { |
| ::xla::jellyfish::LloRegionBuilder::CoreLocation location{ |
| .chip_id = GetLloValue(op.getChipId(), value_map), |
| .core_index = GetLloValue(op.getCoreIndex(), value_map), |
| }; |
| if (!b.target().HasLimitedIciRouting()) { |
| return location; |
| } |
| return ::xla::jellyfish::net_util::GetLimitedIciRoutingTableIndex( |
| b, location, tag, ctx->program_shared_registry, |
| /*use_routing_table_indices=*/false, /*trap_if_invalid=*/false, |
| /*must_return_index=*/true); |
| } |
| |
| template <typename Op> |
| ::xla::jellyfish::LloRegionBuilder::CoreLocation GetOptionalCoreLocation( |
| BuilderContext* ctx, ::xla::jellyfish::LloRegionBuilder& b, Op op, |
| std::string_view tag, |
| const DenseMap<mlir::Value, ::xla::jellyfish::LloValue*>& value_map) { |
| if (!op.getChipId() && !op.getCoreIndex()) { |
| return {}; |
| } |
| if (op.getChipId() && op.getCoreIndex()) { |
| return GetCoreLocation(ctx, b, op, tag, value_map); |
| } |
| LOG(FATAL) |
| << "Either both or none of core_index and chip_id should be specified"; |
| } |
| |
| ::xla::jellyfish::DmaStrides<::xla::jellyfish::LloMemUnit> GetDmaStrides( |
| xla::jellyfish::LloRegionBuilder& b, ArrayRef<int32_t> src_strides, |
| ArrayRef<int32_t> dst_strides, OperandRange steps_per_stride, |
| const DenseMap<mlir::Value, ::xla::jellyfish::LloValue*>& value_map) { |
| using xla::jellyfish::Granule; |
| using xla::jellyfish::LloMemUnit; |
| CHECK_EQ(src_strides.size(), dst_strides.size()); |
| CHECK_EQ(src_strides.size(), steps_per_stride.size()); |
| xla::jellyfish::DmaStrides<LloMemUnit> result; |
| for (int i = 0; i < src_strides.size(); ++i) { |
| LloMemUnit steps; |
| if (i == src_strides.size() - 1) { |
| steps = |
| b.LloMemUnitFromBytes(GetLloValue(steps_per_stride[i], value_map)); |
| } else { |
| steps = LloMemUnit(GetLloValue(steps_per_stride[i], value_map), |
| Granule::kInvalid); |
| } |
| result.emplace_back(b.LloMemUnitFromBytes(b.SimmS32(src_strides[i])), |
| b.LloMemUnitFromBytes(b.SimmS32(dst_strides[i])), |
| steps); |
| } |
| return result; |
| } |
| |
| ::xla::PrimitiveType ToPrimitiveType(mlir::Type type, |
| bool throw_error_on_signless = true) { |
| if (type.isSignlessInteger() && throw_error_on_signless) { |
| LOG(FATAL) << "Signedness cannot be inferred"; |
| } |
| return ::xla::ConvertMlirTypeToPrimitiveType(type); |
| } |
| |
| #include "platforms/xla/mosaic/dialect/llo/llo_builder.cc.inc" |
| |
| } // namespace |
| |
| LogicalResult AppendBlock( |
| BuilderContext* ctx, ::xla::jellyfish::LloRegionBuilder& b, Block& block, |
| DenseMap<Value, ::xla::jellyfish::LloValue*>& value_map) { |
| for (mlir::Operation& any_op : block.without_terminator()) { |
| if (AppendInstruction(ctx, b, &any_op, value_map).failed()) { |
| return failure(); |
| } |
| } |
| return success(); |
| } |
| |
| LogicalResult AppendInstruction( |
| BuilderContext* ctx, ::xla::jellyfish::LloRegionBuilder& b, |
| Operation* raw_op, |
| DenseMap<Value, ::xla::jellyfish::LloValue*>& value_map) { |
| if (auto op = dyn_cast<llo::ConstantOp>(raw_op)) { |
| ::xla::jellyfish::LloValue* value = nullptr; |
| if (auto attr = dyn_cast<IntegerAttr>(op.getValue())) { |
| if (attr.getType().isUnsignedInteger(32)) { |
| value = b.SimmU32(attr.getUInt()); |
| } else if (attr.getType().isSignlessInteger(32)) { |
| value = b.SimmS32(attr.getInt()); |
| } else if (attr.getType().isSignlessInteger(1)) { |
| value = b.Pimm(attr.getInt()); |
| } |
| } else if (auto attr = dyn_cast<FloatAttr>(op.getValue())) { |
| if (attr.getType().isF32()) { |
| value = b.SimmF32(attr.getValueAsDouble()); |
| } |
| } else if (auto attr = dyn_cast<SplatElementsAttr>(op.getValue())) { |
| mlir::ShapedType ty = attr.getType(); |
| int64_t sublanes = b.target().SublaneCount(); |
| int64_t lanes = b.target().LaneCount(); |
| if (ty.getShape() == ArrayRef<int64_t>{sublanes, lanes}) { |
| if (attr.getElementType().isF32()) { |
| value = b.VimmF32(attr.getSplatValue<float>()); |
| } else if (attr.getElementType().isSignlessInteger(32)) { |
| value = b.VimmS32(attr.getSplatValue<int32_t>()); |
| } else if (attr.getElementType().isSignlessInteger(1)) { |
| // TODO(b/272487785): The shape check above might require adjustment. |
| value = b.Vmimm(attr.getSplatValue<bool>()); |
| } |
| } else if (ty.getShape() == ArrayRef<int64_t>{sublanes, lanes, 2}) { |
| if (attr.getElementType().isBF16()) { |
| uint16_t bits = attr.bitcast(IntegerType::get(op.getContext(), 16, |
| IntegerType::Unsigned)) |
| .getSplatValue<uint16_t>(); |
| value = b.VimmU32(bits | (bits << 16)); |
| } else if (attr.getElementType().isSignlessInteger(16)) { |
| uint16_t bits = attr.getSplatValue<uint16_t>(); |
| value = b.VimmU32(bits | (bits << 16)); |
| } else if (b.target().BitsPerVmregLaneAndSublane() >= 2 && |
| attr.getElementType().isSignlessInteger(1)) { |
| value = b.Vmimm(attr.getSplatValue<bool>()); |
| } |
| } else if (ty.getShape() == ArrayRef<int64_t>{sublanes, lanes, 4}) { |
| if (b.target().BitsPerVmregLaneAndSublane() >= 4 && |
| attr.getElementType().isSignlessInteger(1)) { |
| value = b.Vmimm(attr.getSplatValue<bool>()); |
| } else if (attr.getElementType().isSignlessInteger(8)) { |
| uint8_t bits = attr.getSplatValue<uint8_t>(); |
| value = b.VimmU32(bits | (bits << 8) | (bits << 16) | (bits << 24)); |
| } |
| } |
| } |
| if (value != nullptr) { |
| value_map[op.getResult()] = value; |
| return success(); |
| } |
| return failure(); |
| } |
| if (auto op = dyn_cast<llo::VectorBitcastOp>(raw_op)) { |
| value_map[op.getOutput()] = GetLloValue(op.getInput(), value_map); |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::ScalarBitcastOp>(raw_op)) { |
| value_map[op.getOutput()] = GetLloValue(op.getInput(), value_map); |
| return success(); |
| } |
| if (auto op = dyn_cast<mlir::scf::IfOp>(raw_op)) { |
| auto* pred = GetLloValue(op.getCondition(), value_map); |
| if (!op.getThenRegion().hasOneBlock()) { |
| return failure(); |
| } |
| auto& then_block = op.getThenRegion().getBlocks().front(); |
| if (!then_block.empty()) { |
| Operation* then_terminator = then_block.getTerminator(); |
| if (then_terminator->getNumOperands() != op.getNumResults()) { |
| return failure(); |
| } |
| for (int i = 0; i < op.getNumResults(); ++i) { |
| auto val = then_terminator->getOperand(i); |
| if (value_map.find(val) == value_map.end()) { |
| value_map[op.getResult(i)] = |
| val.getType().isa<VectorType>() ? b.VimmS32(0) : b.SimmS32(0); |
| } |
| } |
| xla::jellyfish::LloPredicatedRegion* subregion = b.Predicated(pred); |
| auto subbuilder = subregion->region_builder(); |
| if (AppendBlock(ctx, subbuilder, then_block, value_map).failed()) { |
| return failure(); |
| } |
| for (int i = 0; i < op.getNumResults(); ++i) { |
| auto& result = value_map[op.getResult(i)]; |
| auto* then_result = |
| GetLloValue(then_terminator->getOperand(i), value_map); |
| result = (result == nullptr) ? then_result |
| : subregion->fallthru_region_builder().Phi( |
| result, then_result); |
| } |
| } |
| if (op.getElseRegion().empty()) { |
| return success(); |
| } |
| if (!op.getElseRegion().hasOneBlock()) { |
| return failure(); |
| } |
| auto& else_block = op.getElseRegion().getBlocks().front(); |
| if (!else_block.empty()) { |
| xla::jellyfish::LloValue* neg_pred = b.Pneg(pred); |
| xla::jellyfish::LloPredicatedRegion* subregion = b.Predicated(neg_pred); |
| auto subbuilder = subregion->region_builder(); |
| if (AppendBlock(ctx, subbuilder, else_block, value_map).failed()) { |
| return failure(); |
| } |
| Operation* else_terminator = else_block.getTerminator(); |
| if (else_terminator->getNumOperands() != op.getNumResults()) { |
| return failure(); |
| } |
| for (int i = 0; i < op.getNumResults(); ++i) { |
| auto* result = GetLloValue(op.getResult(i), value_map); |
| auto* else_value = |
| GetLloValue(else_terminator->getOperand(i), value_map); |
| value_map[op.getResult(i)] = |
| subregion->fallthru_region_builder().Phi(result, else_value); |
| } |
| } |
| |
| return success(); |
| } |
| if (auto op = dyn_cast<mlir::scf::ForOp>(raw_op)) { |
| auto* lbd = GetLloValue(op.getLowerBound(), value_map); |
| auto* ubd = GetLloValue(op.getUpperBound(), value_map); |
| auto* step = GetLloValue(op.getStep(), value_map); |
| auto [loop, induction_variable] = |
| b.Loop(b.region()->hlo_instruction(), lbd, ubd, step); |
| value_map[op.getInductionVar()] = induction_variable; |
| // Add iter args to value_map. |
| for (auto [arg, init] : |
| llvm::zip_equal(op.getRegionIterArgs(), op.getInitArgs())) { |
| value_map[arg] = loop->header_builder().Phi(GetLloValue(init, value_map)); |
| } |
| CHECK(op.getRegion().hasOneBlock()); |
| xla::jellyfish::LloRegionBuilder body_builder(loop->body_builder()); |
| if (AppendBlock(ctx, body_builder, *op.getBody(), value_map).failed()) { |
| return op.emitOpError("Failed to build llo loop body in scf.for."); |
| } |
| // Update iter args after each loop and finally map result to iter arg. |
| Operation* terminator = op.getBody()->getTerminator(); |
| CHECK_EQ(terminator->getNumOperands(), op.getNumRegionIterArgs()); |
| for (int i = 0; i < op.getNumRegionIterArgs(); ++i) { |
| value_map[op.getRegionIterArg(i)]->PhiAppend( |
| GetLloValue(terminator->getOperand(i), value_map)); |
| value_map[op.getResult(i)] = |
| GetLloValue(op.getRegionIterArg(i), value_map); |
| } |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::GetBarrierSyncFlagOp>(raw_op)) { |
| if (!ctx->custom_barrier_sync_flag) { |
| return failure(); |
| } |
| value_map[op.getResult()] = ctx->custom_barrier_sync_flag; |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::GetIterationBoundOp>(raw_op)) { |
| if (op.getDim() >= ctx->dynamic_iteration_bounds.size()) { |
| return failure(); |
| } |
| value_map[op.getResult()] = ctx->dynamic_iteration_bounds[op.getDim()]; |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::RegionOp>(raw_op)) { |
| xla::jellyfish::LloRegionBuilder subbuilder( |
| b.region()->AddRegion(b.region()->hlo_instruction())); |
| if (!op.getRegion().hasOneBlock()) { |
| return failure(); |
| } |
| Block& block = op.getRegion().front(); |
| if (AppendBlock(ctx, subbuilder, block, value_map).failed()) { |
| return failure(); |
| } |
| Operation* terminator = block.getTerminator(); |
| if (terminator->getNumOperands() != op.getNumResults()) { |
| return failure(); |
| } |
| for (int i = 0; i < op.getNumResults(); ++i) { |
| value_map[op.getResult(i)] = |
| GetLloValue(terminator->getOperand(i), value_map); |
| } |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::TraceOp>(raw_op)) { |
| std::optional<xla::jellyfish::OnDeviceTraceMe> trace_guard; |
| if (xla::jellyfish::ExecutionProfiler* profiler = b.module()->profiler()) { |
| auto instr_id = |
| profiler->AllocateLloInstrumentationId(op.getMessage().str()); |
| if (!instr_id.ok()) { |
| return failure(); |
| } |
| trace_guard.emplace(b.region(), instr_id.value(), op.getLevel()); |
| } |
| if (!op.getRegion().hasOneBlock()) { |
| return failure(); |
| } |
| Block& block = op.getRegion().front(); |
| if (AppendBlock(ctx, b, block, value_map).failed()) { |
| return failure(); |
| } |
| Operation* terminator = block.getTerminator(); |
| if (terminator->getNumOperands() != op.getNumResults()) { |
| return failure(); |
| } |
| for (int i = 0; i < op.getNumResults(); ++i) { |
| value_map[op.getResult(i)] = |
| GetLloValue(terminator->getOperand(i), value_map); |
| } |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::TraceStartOp>(raw_op)) { |
| if (xla::jellyfish::ExecutionProfiler* profiler = b.module()->profiler()) { |
| auto instr_id = |
| profiler->AllocateLloInstrumentationId(op.getMessage().str()); |
| if (!instr_id.ok()) { |
| return failure(); |
| } |
| ctx->traceme_stack.push_back( |
| std::make_unique<xla::jellyfish::OnDeviceTraceMe>( |
| b.region(), instr_id.value(), op.getLevel())); |
| } |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::TraceStopOp>(raw_op)) { |
| if (b.module()->profiler()) { |
| ctx->traceme_stack.pop_back(); |
| } |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::LogicalDeviceIdOp>(raw_op)) { |
| const auto& module_config = b.hlo()->GetModule()->config(); |
| if (!module_config.has_static_device_assignment()) { |
| return failure(); |
| } |
| const auto& device_assignment = module_config.static_device_assignment(); |
| xla::jellyfish::LloValue* partition_id = |
| xla::jellyfish::net_util::GetPartitionId(b); |
| xla::jellyfish::LloValue* replica_id = |
| xla::jellyfish::net_util::GetReplicaId(b); |
| xla::jellyfish::LloValue* global_id = b.SaddS32( |
| b.SmulS32(replica_id, b.SimmS32(device_assignment.computation_count())), |
| partition_id); |
| value_map[op.getResult()] = global_id; |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::ScalarDivS32Op>(raw_op)) { |
| if (auto crhs = op.getRhs().getDefiningOp<llo::ConstantOp>()) { |
| auto divisor = cast<mlir::IntegerAttr>(crhs.getValue()).getInt(); |
| if (divisor == 0) { |
| goto generic_div; |
| } |
| value_map[op.getResult()] = |
| b.SdivS32(GetLloValue(op.getLhs(), value_map), divisor); |
| return success(); |
| } |
| generic_div: |
| value_map[op.getResult()] = |
| b.SdivS32General(GetLloValue(op.getLhs(), value_map), |
| GetLloValue(op.getRhs(), value_map)); |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::ScalarRemS32Op>(raw_op)) { |
| if (auto crhs = op.getRhs().getDefiningOp<llo::ConstantOp>()) { |
| auto divisor = cast<mlir::IntegerAttr>(crhs.getValue()).getInt(); |
| if (divisor == 0) { |
| goto generic_rem; |
| } |
| value_map[op.getResult()] = |
| b.SremS32(GetLloValue(op.getLhs(), value_map), divisor); |
| return success(); |
| } |
| generic_rem: |
| value_map[op.getResult()] = |
| b.SremS32General(GetLloValue(op.getLhs(), value_map), |
| GetLloValue(op.getRhs(), value_map)); |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::LogS32>(raw_op)) { |
| b.LogS32(GetLloValue(op.getOperand(), value_map), op.getTag().str()); |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::LogPred>(raw_op)) { |
| b.LogPred(GetLloValue(op.getOperand(), value_map), op.getTag().str()); |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::AllocaVmemOp>(raw_op)) { |
| int64_t num_words = op.getNumWords(); |
| int64_t word_size = b.target().VmemWordSizeBytes(); |
| auto addr = b.AllocateScopedVmem( |
| xla::ShapeUtil::MakeShape(::xla::U8, {num_words * word_size})); |
| value_map[op.getResult()] = addr; |
| return success(); |
| } |
| if (auto op = dyn_cast<llo::AllocaSmemOp>(raw_op)) { |
| int64_t num_words = op.getNumWords(); |
| int64_t word_size = b.target().SmemWordSizeBytes(); |
| auto addr = b.AllocateScopedSmem( |
| xla::ShapeUtil::MakeShape(::xla::U8, {num_words * word_size})); |
| value_map[op.getResult()] = addr; |
| return success(); |
| } |
| if (AppendInstructionGenerated(ctx, b, raw_op, value_map).succeeded()) { |
| return success(); |
| } |
| raw_op->emitOpError("Failed to translate operation"); |
| return failure(); |
| } |
| |
| } // namespace llo |
| } // namespace mlir |
