54 lines
1.8 KiB
C++
54 lines
1.8 KiB
C++
|
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
|
||
|
|
#include <gtest/gtest.h>
|
||
|
|
|
||
|
|
using namespace mlir;
|
||
|
|
using mlir::triton::gpu::SharedEncodingAttr;
|
||
|
|
|
||
|
|
struct swizzleParams {
|
||
|
|
int vec;
|
||
|
|
int perPhase;
|
||
|
|
int maxPhase;
|
||
|
|
};
|
||
|
|
|
||
|
|
struct ParamT {
|
||
|
|
std::array<int64_t, 2> shape;
|
||
|
|
int opIdx;
|
||
|
|
int typeWidth;
|
||
|
|
swizzleParams refSwizzle;
|
||
|
|
};
|
||
|
|
|
||
|
|
class SwizzleDotOperandTestFixture : public ::testing::TestWithParam<ParamT> {
|
||
|
|
protected:
|
||
|
|
ParamType param;
|
||
|
|
};
|
||
|
|
|
||
|
|
TEST_P(SwizzleDotOperandTestFixture, DotOperands) {
|
||
|
|
auto params = GetParam();
|
||
|
|
// init context
|
||
|
|
MLIRContext ctx;
|
||
|
|
ctx.loadDialect<triton::gpu::TritonGPUDialect>();
|
||
|
|
// create encoding
|
||
|
|
auto parent = triton::gpu::MmaEncodingAttr::get(&ctx, 2, 0, {1, 1});
|
||
|
|
auto encoding =
|
||
|
|
triton::gpu::DotOperandEncodingAttr::get(&ctx, params.opIdx, parent);
|
||
|
|
|
||
|
|
// create element type
|
||
|
|
Type eltType = IntegerType::get(&ctx, params.typeWidth);
|
||
|
|
auto layout =
|
||
|
|
SharedEncodingAttr::get(&ctx, encoding, params.shape, {1, 0}, eltType);
|
||
|
|
|
||
|
|
ASSERT_EQ(layout.getVec(), params.refSwizzle.vec);
|
||
|
|
ASSERT_EQ(layout.getPerPhase(), params.refSwizzle.perPhase);
|
||
|
|
ASSERT_EQ(layout.getMaxPhase(), params.refSwizzle.maxPhase);
|
||
|
|
}
|
||
|
|
|
||
|
|
INSTANTIATE_TEST_SUITE_P(TestDotOperands, SwizzleDotOperandTestFixture,
|
||
|
|
::testing::Values(ParamT{{128, 64}, 0, 16, {8, 1, 8}},
|
||
|
|
ParamT{{64, 256}, 1, 16, {8, 1, 8}},
|
||
|
|
ParamT{{128, 32}, 0, 16, {8, 2, 4}},
|
||
|
|
ParamT{{32, 128}, 1, 16, {8, 1, 8}},
|
||
|
|
ParamT{{32, 32}, 0, 16, {8, 2, 4}},
|
||
|
|
ParamT{{32, 32}, 1, 16, {8, 2, 4}},
|
||
|
|
ParamT{{16, 16}, 0, 16, {8, 4, 2}},
|
||
|
|
ParamT{{16, 16}, 1, 16, {8, 4, 2}}));
|