Skip to content

Commit 2534ac7

Browse files
akroviakovakroviakov
committed
add initial conversions for xevm block ops
1 parent 238706b commit 2534ac7

File tree

9 files changed

+1032
-8
lines changed

9 files changed

+1032
-8
lines changed

‎include/gc/Conversion/Passes.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
#ifndef GC_CONVERSION_PASSES_H
1010
#define GC_CONVERSION_PASSES_H
1111

12-
#include "gc/Conversion/XeVMToLLVM.h"
12+
#include "gc/Conversion/XeVMToLLVM/XeVMToLLVM.h"
1313

1414
namespace mlir {
1515

‎include/gc/Dialect/LLVMIR/XeVMOps.td

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,15 @@ def XeVM_Dialect : Dialect {
1919
let name = "xevm";
2020
let cppNamespace = "::mlir::xevm";
2121
let dependentDialects = ["LLVM::LLVMDialect"];
22+
23+
let extraClassDeclaration = [{
24+
/// Get the name for the attribute used to specify cache control
25+
/// decorations.
26+
static constexpr ::llvm::StringRef getCacheControlsAttrName() {
27+
return ::llvm::StringLiteral("xevm.DecorationCacheControlINTEL");
28+
}
29+
}];
30+
2231
let useDefaultAttributePrinterParser = 1;
2332
}
2433

@@ -161,6 +170,52 @@ def XeVM_BlockStore2dOp : XeVM_Op<"blockstore2d">,
161170
let hasVerifier = 1;
162171
}
163172

173+
def XeVM_BlockPrefetch2dOp : XeVM_Op<"blockprefetch2d">,
174+
Arguments<(ins
175+
Arg<LLVM_AnyPointer, "", [MemRead]>:$ptr,
176+
I32:$base_width,
177+
I32:$base_height,
178+
I32:$base_pitch,
179+
I32:$x,
180+
I32:$y,
181+
I32Attr:$elem_size_in_bits,
182+
I32Attr:$tile_width,
183+
I32Attr:$tile_height,
184+
I32Attr:$v_blocks,
185+
DefaultValuedAttr<XeVM_L1LoadCacheControl, "::mlir::xevm::L1LoadCacheControl::DEFAULT">:$l1_cache_control,
186+
DefaultValuedAttr<XeVM_L3LoadCacheControl, "::mlir::xevm::L3LoadCacheControl::DEFAULT">:$l3_cache_control
187+
)> {
188+
189+
let summary = "2D block prefetch";
190+
191+
let description = [{
192+
The `xevm.blockprefetch2d` operation prefetches a two dimensional tile
193+
from a larger matrix residing in memory. The parameters are:
194+
$ptr - the base address of the matrix containing the tile to prefetch
195+
$base_width, $base_height, $base_pitch - the shape of the matrix
196+
$x, $y, $tile_width, $tile_height - the starting offsets and shape of tile to prefetch
197+
$elem_size_in_bits - the size in bits of the matrix element
198+
- 32 for f32, bf32
199+
- 16 for f16, int16, bf16
200+
- 8 for int8, int4, int2
201+
$v_blocks - number of tiles to prefetch
202+
$cache_control - an enumerator that sets the L1 and L3 cache behaviour
203+
204+
Notes:
205+
- coordinate is provided in elements, while width and pitch are provided in bytes.
206+
}];
207+
208+
let assemblyFormat = [{
209+
operands ` ` `{` `elem_size_in_bits` `=` $elem_size_in_bits `,` `tile_width` `=` $tile_width `,`
210+
`tile_height` `=` $tile_height `,` `v_blocks` `=` $v_blocks `,` `l1_cache_control` `=` $l1_cache_control `,`
211+
`l3_cache_control` `=` $l3_cache_control `}`
212+
attr-dict `:` `(` type(operands) `)`
213+
}];
214+
215+
let hasVerifier = 1;
216+
}
217+
218+
164219
def XeVM_TargetAttr : XeVM_Attr<"XeVMTarget", "target"> {
165220
let description = [{
166221
GPU target attribute for controlling compilation of targets. All

‎lib/gc/Conversion/XeVMToLLVM/XeVMToLLVM.cpp

Lines changed: 271 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,18 @@
1111
#include "gc/Dialect/LLVMIR/XeVMDialect.h"
1212
#include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
1313
#include "mlir/Conversion/LLVMCommon/Pattern.h"
14+
#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
1415
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
1516
#include "mlir/Pass/Pass.h"
1617
#include "mlir/Support/LLVM.h"
18+
#include "llvm/Support/FormatVariadic.h"
19+
20+
#include "mlir/IR/BuiltinTypes.h"
21+
#include "mlir/IR/Types.h"
22+
23+
#include "llvm/ADT/STLExtras.h"
24+
#include "llvm/ADT/TypeSwitch.h"
25+
#include "llvm/Support/raw_ostream.h"
1726

1827
#define DEBUG_TYPE "xevm-to-llvm"
1928

@@ -26,6 +35,230 @@ using namespace mlir;
2635
using namespace xevm;
2736

2837
namespace {
38+
struct LLVMFuncAttributeOptions {
39+
bool isConvergent = false;
40+
bool isNoUnwind = false;
41+
bool isWillReturn = false;
42+
LLVM::MemoryEffectsAttr memEffectsAttr{};
43+
};
44+
static constexpr LLVMFuncAttributeOptions convergentAttrs = {
45+
true, false, false, {}};
46+
static constexpr LLVMFuncAttributeOptions noUnwindAttrs = {
47+
false, true, false, {}};
48+
static constexpr LLVMFuncAttributeOptions noUnwindWillReturnAttrs = {
49+
false, true, true, {}};
50+
static constexpr LLVMFuncAttributeOptions convergentNoUnwindWillReturnAttrs = {
51+
true, true, true, {}};
52+
53+
std::string getTypeMangling(Type ty, bool isUnsigned = false) {
54+
return TypeSwitch<Type, std::string>(ty)
55+
.Case([isUnsigned](VectorType ty) -> std::string {
56+
return "Dv" + std::to_string(ty.getNumElements()) + "_" +
57+
getTypeMangling(ty.getElementType(), isUnsigned);
58+
})
59+
.Case([](Float16Type) -> std::string { return "Dh"; })
60+
.Case([](Float32Type) -> std::string { return "f"; })
61+
.Case([](Float64Type) -> std::string { return "d"; })
62+
.Case([isUnsigned](IntegerType ty) -> std::string {
63+
switch (ty.getWidth()) {
64+
case 8:
65+
return isUnsigned ? "h" : "c";
66+
case 16:
67+
return isUnsigned ? "t" : "s";
68+
case 32:
69+
return isUnsigned ? "j" : "i";
70+
case 64:
71+
return isUnsigned ? "m" : "l";
72+
default:
73+
llvm_unreachable("unhandled integer type");
74+
}
75+
});
76+
}
77+
78+
template <typename OpType>
79+
static std::optional<ArrayAttr>
80+
getCacheControlMetadata(ConversionPatternRewriter &rewriter, OpType op,
81+
const bool isLoad) {
82+
if ((op.getL1CacheControlAttr() ==
83+
xevm::L1StoreCacheControlAttr::get(
84+
rewriter.getContext(), xevm::L1StoreCacheControl::DEFAULT) &&
85+
op.getL3CacheControlAttr() ==
86+
xevm::L3StoreCacheControlAttr::get(
87+
rewriter.getContext(), xevm::L3StoreCacheControl::DEFAULT)) ||
88+
89+
(op.getL1CacheControlAttr() ==
90+
xevm::L1LoadCacheControlAttr::get(
91+
rewriter.getContext(), xevm::L1LoadCacheControl::DEFAULT) &&
92+
op.getL3CacheControlAttr() ==
93+
xevm::L3LoadCacheControlAttr::get(
94+
rewriter.getContext(), xevm::L3LoadCacheControl::DEFAULT))) {
95+
return {};
96+
}
97+
constexpr int32_t decorationCacheControlArity{4};
98+
constexpr int32_t loadCacheControlKey{6442};
99+
constexpr int32_t storeCacheControlKey{6443};
100+
constexpr int32_t l1Level{0};
101+
constexpr int32_t l3Level{1};
102+
const int32_t controlKey{isLoad ? loadCacheControlKey : storeCacheControlKey};
103+
SmallVector<int32_t, decorationCacheControlArity> decorationsL1{
104+
controlKey, l1Level, static_cast<int32_t>(op.getL1CacheControl()), 0};
105+
SmallVector<int32_t, decorationCacheControlArity> decorationsL3{
106+
controlKey, l3Level, static_cast<int32_t>(op.getL3CacheControl()), 0};
107+
auto arrayAttrL1 = rewriter.getI32ArrayAttr(decorationsL1);
108+
auto arrayAttrL3 = rewriter.getI32ArrayAttr(decorationsL3);
109+
110+
SmallVector<Attribute, 2> combinedAttrs = {arrayAttrL1, arrayAttrL3};
111+
return rewriter.getArrayAttr(combinedAttrs);
112+
}
113+
114+
static LLVM::CallOp createDeviceFunctionCall(
115+
ConversionPatternRewriter &rewriter, StringRef funcName, Type retType,
116+
ArrayRef<Type> argTypes, ArrayRef<Value> args,
117+
mlir::ArrayRef<std::pair<unsigned, mlir::StringRef>> paramAttrs,
118+
LLVMFuncAttributeOptions funcAttributeOptions) {
119+
auto moduleOp = rewriter.getBlock()->getParent()->getParentOfType<ModuleOp>();
120+
MLIRContext *ctx = rewriter.getContext();
121+
Location loc = UnknownLoc::get(ctx);
122+
123+
LLVM::LLVMFuncOp funcOp =
124+
LLVM::lookupOrCreateFn(moduleOp, funcName, argTypes, retType);
125+
funcOp.setCConv(LLVM::cconv::CConv::SPIR_FUNC);
126+
funcOp.setConvergent(funcAttributeOptions.isConvergent);
127+
funcOp.setNoUnwind(funcAttributeOptions.isNoUnwind);
128+
funcOp.setWillReturn(funcAttributeOptions.isWillReturn);
129+
130+
if (funcAttributeOptions.memEffectsAttr)
131+
funcOp.setMemoryEffectsAttr(funcAttributeOptions.memEffectsAttr);
132+
133+
for (auto [idx, attrName] : paramAttrs)
134+
funcOp.setArgAttr(idx, attrName, rewriter.getUnitAttr());
135+
136+
// if (!passthroughAttrs.getFnAttributes().empty())
137+
// funcOp->setAttrs(passthroughAttrs.getFnAttributes().getDictionary(ctx));
138+
139+
auto callOp = rewriter.create<LLVM::CallOp>(loc, funcOp, args);
140+
callOp->setAttrs(funcOp->getAttrs());
141+
142+
return callOp;
143+
}
144+
145+
template <typename OpType>
146+
class LoadStorePrefetchNdToOCLPattern : public OpConversionPattern<OpType> {
147+
using OpConversionPattern<OpType>::OpConversionPattern;
148+
LogicalResult
149+
matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
150+
ConversionPatternRewriter &rewriter) const override {
151+
constexpr bool isLoad = std::is_same_v<OpType, xevm::BlockLoad2dOp>;
152+
constexpr bool isStore = std::is_same_v<OpType, xevm::BlockStore2dOp>;
153+
constexpr bool isPrefetch = std::is_same_v<OpType, xevm::BlockPrefetch2dOp>;
154+
auto loc = op.getLoc();
155+
VectorType vecType;
156+
if constexpr (isLoad) {
157+
vecType = op.getRes().getType();
158+
} else if constexpr (isStore) {
159+
vecType = op.getStoredVal().getType();
160+
}
161+
162+
auto i32Type = rewriter.getI32Type();
163+
bool vnni = false;
164+
bool transpose = false;
165+
if constexpr (isLoad) {
166+
vnni = op.getVnniTransform();
167+
transpose = op.getTranspose();
168+
}
169+
170+
Value byteCoord =
171+
rewriter.create<LLVM::UndefOp>(loc, VectorType::get(2, i32Type));
172+
Value zero = rewriter.create<LLVM::ConstantOp>(
173+
loc, i32Type, rewriter.getI32IntegerAttr(0));
174+
Value one = rewriter.create<LLVM::ConstantOp>(
175+
loc, i32Type, rewriter.getI32IntegerAttr(1));
176+
byteCoord = rewriter.create<LLVM::InsertElementOp>(
177+
loc, VectorType::get(2, i32Type), byteCoord, op.getX(), zero);
178+
byteCoord = rewriter.create<LLVM::InsertElementOp>(
179+
loc, VectorType::get(2, i32Type), byteCoord, op.getY(), one);
180+
SmallVector<Value> args{op.getPtr(), op.getBaseWidth(), op.getBaseHeight(),
181+
op.getBasePitch(), byteCoord};
182+
SmallVector<Type> retTypes;
183+
Value spvLoadDstPtr;
184+
std::string funcName, bitWidthId;
185+
SmallVector<std::pair<unsigned, mlir::StringRef>, 4> paramAttrs;
186+
if constexpr (isPrefetch) { // Prefetch
187+
funcName = "intel_sub_group_2d_block_prefetch";
188+
paramAttrs = {std::make_pair(0, LLVM::LLVMDialect::getNonNullAttrName())};
189+
} else {
190+
auto vecElemType = vecType.getElementType();
191+
auto vecElemBitWidth = vecElemType.getIntOrFloatBitWidth();
192+
Value numElems = rewriter.create<LLVM::ConstantOp>(
193+
loc, i32Type, vecType.getNumElements());
194+
auto dstOrSrcPtr = rewriter.create<LLVM::AllocaOp>(
195+
loc, LLVM::LLVMPointerType::get(rewriter.getContext()), vecElemType,
196+
numElems);
197+
args.push_back(dstOrSrcPtr);
198+
if constexpr (isLoad) { // Load
199+
funcName = "intel_sub_group_2d_block_read";
200+
bitWidthId = getTypeMangling(vecElemType, /*isUnsigned=*/true);
201+
if (vnni)
202+
funcName += "_transform";
203+
else if (transpose)
204+
funcName += "_transpose";
205+
spvLoadDstPtr = dstOrSrcPtr;
206+
retTypes.push_back(vecType);
207+
paramAttrs = {
208+
std::make_pair(0, LLVM::LLVMDialect::getNonNullAttrName()),
209+
std::make_pair(0, LLVM::LLVMDialect::getReadonlyAttrName()),
210+
std::make_pair(5, LLVM::LLVMDialect::getNonNullAttrName()),
211+
std::make_pair(5, LLVM::LLVMDialect::getWriteOnlyAttrName()),
212+
};
213+
} else { // Store
214+
funcName = "intel_sub_group_2d_block_write";
215+
bitWidthId = (vecElemBitWidth == 32)
216+
? "j"
217+
: ((vecElemBitWidth == 16) ? "t" : "h");
218+
rewriter.create<LLVM::StoreOp>(loc, op.getStoredVal(), dstOrSrcPtr);
219+
paramAttrs = {
220+
std::make_pair(0, LLVM::LLVMDialect::getNonNullAttrName()),
221+
std::make_pair(0, LLVM::LLVMDialect::getWriteOnlyAttrName()),
222+
std::make_pair(5, LLVM::LLVMDialect::getNonNullAttrName()),
223+
std::make_pair(5, LLVM::LLVMDialect::getReadonlyAttrName()),
224+
};
225+
}
226+
}
227+
228+
// !X = !{i32 %decoration_kind%, i32 %level%, i32 %control%, i32 %operand of
229+
// the instruction to decorate%}
230+
funcName =
231+
llvm::formatv("{0}_{1}b_{2}r{3}x{4}c", funcName, op.getElemSizeInBits(),
232+
op.getTileHeight(), op.getTileWidth(), op.getVBlocks())
233+
.str();
234+
funcName = llvm::formatv("_Z{0}{1}PU3AS1viiiDv2_i{2}{3}", funcName.size(),
235+
funcName, isPrefetch ? "" : "P", bitWidthId)
236+
.str();
237+
SmallVector<Type> argTypes;
238+
for (auto arg : args) {
239+
argTypes.push_back(arg.getType());
240+
}
241+
LLVM::CallOp call = createDeviceFunctionCall(
242+
rewriter, funcName, LLVM::LLVMVoidType::get(rewriter.getContext()),
243+
argTypes, args, paramAttrs, noUnwindWillReturnAttrs);
244+
if (std::optional<ArrayAttr> optCacheControls =
245+
getCacheControlMetadata(rewriter, op, isLoad || isPrefetch)) {
246+
call->setAttr(xevm::XeVMDialect::getCacheControlsAttrName(),
247+
*optCacheControls);
248+
}
249+
if constexpr (isLoad)
250+
rewriter.replaceOp(
251+
op, rewriter.create<LLVM::LoadOp>(loc, vecType, spvLoadDstPtr));
252+
else
253+
rewriter.eraseOp(op);
254+
return success();
255+
}
256+
};
257+
258+
//===----------------------------------------------------------------------===//
259+
// Pass Definition
260+
//===----------------------------------------------------------------------===//
261+
29262
struct ConvertXeVMToLLVMPass
30263
: public impl::ConvertXeVMToLLVMPassBase<ConvertXeVMToLLVMPass> {
31264
using Base::Base;
@@ -37,19 +270,51 @@ struct ConvertXeVMToLLVMPass
37270
void runOnOperation() override {
38271
ConversionTarget target(getContext());
39272
target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
40-
RewritePatternSet pattern(&getContext());
41-
mlir::populateXeVMToLLVMConversionPatterns(pattern);
42-
if (failed(
43-
applyPartialConversion(getOperation(), target, std::move(pattern))))
273+
target.addIllegalDialect<xevm::XeVMDialect>();
274+
RewritePatternSet patterns(&getContext());
275+
mlir::populateXeVMToLLVMConversionPatterns(patterns);
276+
if (failed(applyPartialConversion(getOperation(), target,
277+
std::move(patterns))))
44278
signalPassFailure();
45279
}
46280
};
47281
} // namespace
48282

283+
//===----------------------------------------------------------------------===//
284+
// Pattern Population
285+
//===----------------------------------------------------------------------===//
286+
49287
void mlir::populateXeVMToLLVMConversionPatterns(RewritePatternSet &patterns) {
50-
/*TODO*/
288+
patterns.add<LoadStorePrefetchNdToOCLPattern<xevm::BlockLoad2dOp>,
289+
LoadStorePrefetchNdToOCLPattern<xevm::BlockStore2dOp>,
290+
LoadStorePrefetchNdToOCLPattern<xevm::BlockPrefetch2dOp>>(
291+
patterns.getContext());
51292
}
52293

294+
//===----------------------------------------------------------------------===//
295+
// ConvertToLLVMPatternInterface implementation
296+
//===----------------------------------------------------------------------===//
297+
298+
namespace {
299+
/// Implement the interface to convert XeVM to LLVM.
300+
struct XeVMToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
301+
using ConvertToLLVMPatternInterface::ConvertToLLVMPatternInterface;
302+
void loadDependentDialects(MLIRContext *context) const final {
303+
context->loadDialect<LLVM::LLVMDialect>();
304+
}
305+
306+
/// Hook for derived dialect interface to provide conversion patterns
307+
/// and mark dialect legal for the conversion target.
308+
void populateConvertToLLVMConversionPatterns(
309+
ConversionTarget &target, LLVMTypeConverter &typeConverter,
310+
RewritePatternSet &patterns) const final {
311+
populateXeVMToLLVMConversionPatterns(patterns);
312+
}
313+
};
314+
} // namespace
315+
53316
void mlir::registerConvertXeVMToLLVMInterface(DialectRegistry &registry) {
54-
/*TODO*/
317+
registry.addExtension(+[](MLIRContext *ctx, xevm::XeVMDialect *dialect) {
318+
dialect->addInterfaces<XeVMToLLVMDialectInterface>();
319+
});
55320
}

0 commit comments

Comments
 (0)