This commit is contained in:
@@ -1,8 +1,14 @@
|
||||
#pragma once
|
||||
|
||||
#include <optional>
|
||||
|
||||
#include "llvm/ADT/ArrayRef.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
#include "llvm/ADT/STLFunctionalExtras.h"
|
||||
|
||||
#include "mlir/IR/BuiltinTypes.h"
|
||||
#include "mlir/IR/Builders.h"
|
||||
#include "mlir/IR/Value.h"
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/Support/LogicalResult.h"
|
||||
|
||||
@@ -59,6 +65,39 @@ forEachContiguousDestinationChunk(llvm::ArrayRef<int64_t> destShape,
|
||||
llvm::function_ref<mlir::LogicalResult(llvm::ArrayRef<int64_t>, int64_t, int64_t)>
|
||||
callback);
|
||||
|
||||
struct FragmentAssemblyCopy {
|
||||
mlir::Value source;
|
||||
mlir::RankedTensorType sourceType;
|
||||
unsigned hostTargetIndex = 0;
|
||||
int64_t lane = 0;
|
||||
int64_t sourceByteOffset = 0;
|
||||
int64_t hostByteOffset = 0;
|
||||
int64_t byteSize = 0;
|
||||
};
|
||||
|
||||
struct FragmentAssemblyCopyRun {
|
||||
mlir::Value source;
|
||||
mlir::RankedTensorType sourceType;
|
||||
unsigned hostTargetIndex = 0;
|
||||
int64_t count = 0;
|
||||
int64_t sourceStepBytes = 0;
|
||||
int64_t hostStepBytes = 0;
|
||||
int64_t byteSize = 0;
|
||||
mlir::SmallVector<int64_t, 8> sourceStartBytesByLane;
|
||||
mlir::SmallVector<int64_t, 8> hostStartBytesByLane;
|
||||
};
|
||||
|
||||
mlir::FailureOr<mlir::SmallVector<FragmentAssemblyCopyRun, 8>>
|
||||
groupFragmentAssemblyCopyRuns(llvm::ArrayRef<FragmentAssemblyCopy> copies, uint32_t laneCount = 1);
|
||||
|
||||
mlir::FailureOr<mlir::Value> emitFragmentAssemblyCopyRuns(mlir::IRRewriter& rewriter,
|
||||
mlir::Location loc,
|
||||
llvm::ArrayRef<FragmentAssemblyCopyRun> runs,
|
||||
mlir::Value hostTarget,
|
||||
mlir::Operation* anchor,
|
||||
std::optional<mlir::Value> laneArg = std::nullopt,
|
||||
mlir::Value baseHostOffset = {});
|
||||
|
||||
inline mlir::tensor::EmptyOp
|
||||
createEmptyTensorFromShaped(mlir::IRRewriter& rewriter, mlir::Location loc, mlir::ShapedType shapedType) {
|
||||
return mlir::tensor::EmptyOp::create(rewriter, loc, shapedType.getShape(), shapedType.getElementType());
|
||||
|
||||
Reference in New Issue
Block a user