Skip to content

Commit

Permalink
temporary bring back some old definitions to avoid compiling error fr…
Browse files Browse the repository at this point in the history
…om XeGPUToSpirv, these changes will be removed by dewei after he merge his changes to XeGPUToSpirv
  • Loading branch information
chencha3 committed Sep 20, 2023
1 parent b1a630d commit 853014a
Show file tree
Hide file tree
Showing 2 changed files with 198 additions and 0 deletions.
178 changes: 178 additions & 0 deletions include/imex/Dialect/XeGPU/IR/XeGPUOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -587,4 +587,182 @@ def XeGPU_MfenceOp
}];
}

// ============================ begin for dewei ==================================== //
// The following old definitions pending removal by Dewei.
// They are temporary used here to avoid compiling errors of XeGPUToSPIRV and merge the PR
class XeGPU_OpWithOffsetSizesAndStrides<string mnemonic,
list<Trait> traits = []>
: XeGPU_Op<mnemonic, traits> {
code extraBaseClassDeclaration = [{
/// Returns the dynamic sizes for this init_tile operation if specified.
::mlir::Operation::operand_range getDynamicSizes() { return getSizes(); }

/// Return the list of Range (i.e. offset, size, stride). Each
/// Range entry contains either the dynamic value or a ConstantIndexOp
/// constructed with `b` at location `loc`.
::mlir::SmallVector<::mlir::Range, 8> getOrCreateRanges(
::mlir::OpBuilder &b, ::mlir::Location loc) {
return ::mlir::getOrCreateRanges(*this, b, loc);
}
}];
}

class XeGPU_Type<string name, string typeMnemonic, list<Trait> traits = [],
string baseCppClass = "::mlir::Type">
: TypeDef<XeGPUDialect, name, traits, baseCppClass> {
let mnemonic = typeMnemonic;
}

def XeGPU_Tile: XeGPU_Type<"Tile", "tile", [ShapedTypeInterface],
"::imex::xegpu::BaseTensorDescType">
{
let summary = "Tile Type representing a 2D tensor";

let parameters = (ins ArrayRefParameter<"int64_t">:$shape,
"::mlir::Type":$elementType);

let builders = [
TypeBuilderWithInferredContext<(ins
"::llvm::ArrayRef<int64_t>":$shape, "::mlir::Type":$elementType), [{
assert(shape.size()==2);
return $_get(elementType.getContext(), shape, elementType);
}]>,
TypeBuilderWithInferredContext<(ins
"int64_t":$dim0, "int64_t":$dim1, "::mlir::Type":$elementType), [{
llvm::SmallVector<int64_t, 2> shape{dim0, dim1};
assert(shape.size()==2);
return $_get(elementType.getContext(), shape, elementType);
}]>
];

let extraClassDeclaration = [{
using ::mlir::ShapedType::Trait<TileType>::clone;
using ::mlir::ShapedType::Trait<TileType>::getElementTypeBitWidth;
using ::mlir::ShapedType::Trait<TileType>::getRank;
using ::mlir::ShapedType::Trait<TileType>::getNumElements;
using ::mlir::ShapedType::Trait<TileType>::isDynamicDim;
using ::mlir::ShapedType::Trait<TileType>::hasStaticShape;
using ::mlir::ShapedType::Trait<TileType>::getNumDynamicDims;
using ::mlir::ShapedType::Trait<TileType>::getDimSize;
using ::mlir::ShapedType::Trait<TileType>::getDynamicDimIndex;

static mlir::LogicalResult parseShape(mlir::AsmParser &parser,
llvm::SmallVector<int64_t> &shape,
mlir::Type &type) {
llvm::SmallVector<int64_t> dimensions;
if (parser.parseDimensionList(dimensions))
return mlir::failure();

mlir::Type t;
if (parser.parseType(t))
return mlir::failure();

shape = std::move(dimensions);
type = std::move(t);
return mlir::success();
}

static void printShape(mlir::AsmPrinter &printer,
llvm::ArrayRef<int64_t> shape,
mlir::Type type) {
for (int64_t dim : shape) {
if (mlir::ShapedType::isDynamic(dim))
printer << '?';
else
printer << dim;
printer << 'x';
}
printer << type;
}
}];

let assemblyFormat = "`<` custom<Shape>($shape, $elementType) `>`";
}

def XeGPU_InitTileOp : XeGPU_OpWithOffsetSizesAndStrides<"init_tile", [
Pure, AttrSizedOperandSegments,
OffsetSizeAndStrideOpInterface,
ViewLikeOpInterface
]> {
let summary = "Tile operation";
let description = [{
}];

let arguments = (ins
AnyMemRef:$source,
Variadic<Index>:$offsets,
Variadic<Index>:$sizes,
Variadic<Index>:$strides,
DenseI64ArrayAttr:$static_offsets,
DenseI64ArrayAttr:$static_sizes,
DenseI64ArrayAttr:$static_strides
);
let results = (outs XeGPU_Tile:$result);

let assemblyFormat = [{
$source ``
custom<DynamicIndexList>($offsets, $static_offsets)
custom<DynamicIndexList>($sizes, $static_sizes)
custom<DynamicIndexList>($strides, $static_strides)
attr-dict `:` qualified(type($source)) `->` qualified(type($result))
}];

let extraClassDeclaration = extraBaseClassDeclaration # [{
/// Returns the type of the base memref operand.
::mlir::MemRefType getSourceType() {
return ::llvm::cast<::mlir::MemRefType>(getSource().getType());
}

/// The result of an init_tile is always a Tile of TileType.
TileType getType() {
return getResult().getType().cast<TileType>();
}

/// Return the expected rank of each of the`static_offsets`, `static_sizes`
/// and `static_strides` attributes.
std::array<unsigned, 3> getArrayAttrMaxRanks() {
unsigned rank = getSourceType().getRank();
return {rank, rank, rank};
}

/// Return the number of leading operands before the `offsets`, `sizes` and
/// and `strides` operands.
static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }

::mlir::Value getViewSource() { return getSource(); }

}];
}

def XeGPU_Load2DOp
: XeGPU_Op<"load_2d"> {
let summary = "loads a 2D block from global memory (represented by tile) to registers (represented by vector)";
let arguments = (ins
XeGPU_Tile: $tile,
OptionalAttr<I32Attr>: $vnni_axis,
OptionalAttr<BoolAttr>: $transpose
);
let results = (outs Builtin_Vector: $result);

let assemblyFormat = [{
$tile (`VNNI_AXIS` $vnni_axis^)? ` ` (`TRANSPOSE` $transpose^)?
attr-dict `:` qualified(type($tile)) `->` qualified(type($result))
}];
}

def XeGPU_Store2DOp
: XeGPU_Op<"store_2d", []> {
let summary = "stores a 2D block register region back to memory";
let arguments = (ins
XeGPU_Tile: $tile,
Builtin_Vector: $value
);

let assemblyFormat = [{
$tile`,`` `$value attr-dict `:` `(` qualified(type($tile)) `,` qualified(type($value)) `)`
}];
}
// ============================ end for dewei ==================================== //


#endif // _XeGPU_OPS_TD_INCLUDED_
20 changes: 20 additions & 0 deletions lib/Dialect/XeGPU/IR/XeGPUOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1084,6 +1084,26 @@ ::mlir::LogicalResult UpdateNDOffsetOp::verify() {
return ::mlir::success();
}


// =========================== begin for dewei ============================== //
// the following code block is pending removal by Dewei. It is from old definition,
// and is used to temporary avoid XeGPUToSPIRV compiling error.
bool BaseTensorDescType::hasRank() const { return true; }

llvm::ArrayRef<int64_t> BaseTensorDescType::getShape() const {
return cast<TileType>().getShape();
}

BaseTensorDescType BaseTensorDescType::cloneWith(std::optional<llvm::ArrayRef<int64_t>> shape,
Type elementType) const {
return TileType::get(shape.value_or(getShape()), elementType);
}

bool BaseTensorDescType::isValidElementType(Type type) {
return type.isIntOrIndexOrFloat() || type.isa<mlir::ComplexType>();
}
// =========================== end for dewei ============================== //

} // namespace xegpu
} // namespace imex

Expand Down

0 comments on commit 853014a

Please sign in to comment.