robba
Validate Operations / validate-operations (push) Has been cancelled

This commit is contained in:
NiccoloN
2026-06-29 12:22:33 +02:00
parent 78e97f9fd8
commit e8f09fd67f
8 changed files with 1376 additions and 492 deletions
@@ -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());