Skip to content

[HLSL][Matrix] Change MatrixSubscriptExpr flattened index to row major #165666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

[HLSL][Matrix] Change MatrixSubscriptExpr flattened index to row major #165666

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 12 additions & 5 deletions clang/lib/CodeGen/CGExpr.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4824,11 +4824,18 @@ LValue CodeGenFunction::EmitMatrixSubscriptExpr(const MatrixSubscriptExpr *E) {
llvm::Value *RowIdx = EmitMatrixIndexExpr(E->getRowIdx());
llvm::Value *ColIdx = EmitMatrixIndexExpr(E->getColumnIdx());

llvm::Value *NumRows = Builder.getIntN(
RowIdx->getType()->getScalarSizeInBits(),
E->getBase()->getType()->castAs<ConstantMatrixType>()->getNumRows());
llvm::Value *FinalIdx =
Builder.CreateAdd(Builder.CreateMul(ColIdx, NumRows), RowIdx);
llvm::Value *FinalIdx;
if (getLangOpts().HLSL) {
llvm::Value *NumCols = Builder.getIntN(
RowIdx->getType()->getScalarSizeInBits(),
E->getBase()->getType()->castAs<ConstantMatrixType>()->getNumColumns());
FinalIdx = Builder.CreateAdd(Builder.CreateMul(RowIdx, NumCols), ColIdx);
} else {
llvm::Value *NumRows = Builder.getIntN(
RowIdx->getType()->getScalarSizeInBits(),
E->getBase()->getType()->castAs<ConstantMatrixType>()->getNumRows());
FinalIdx = Builder.CreateAdd(Builder.CreateMul(ColIdx, NumRows), RowIdx);
}
return LValue::MakeMatrixElt(
MaybeConvertMatrixAddress(Base.getAddress(), *this), FinalIdx,
E->getBase()->getType(), Base.getBaseInfo(), TBAAAccessInfo());
Expand Down
12 changes: 10 additions & 2 deletions clang/lib/CodeGen/CGExprScalar.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2118,9 +2118,17 @@ Value *ScalarExprEmitter::VisitMatrixSubscriptExpr(MatrixSubscriptExpr *E) {
Value *ColumnIdx = CGF.EmitMatrixIndexExpr(E->getColumnIdx());

const auto *MatrixTy = E->getBase()->getType()->castAs<ConstantMatrixType>();
unsigned NumRows = MatrixTy->getNumRows();

llvm::MatrixBuilder MB(Builder);
Value *Idx = MB.CreateIndex(RowIdx, ColumnIdx, NumRows);
Value *Idx;
if (CGF.getLangOpts().HLSL) {
unsigned NumCols = MatrixTy->getNumColumns();
Idx = MB.createRowMajorIndex(RowIdx, ColumnIdx, NumCols);
} else {
unsigned NumRows = MatrixTy->getNumRows();
Idx = MB.createColumnMajorIndex(RowIdx, ColumnIdx, NumRows);
}

if (CGF.CGM.getCodeGenOpts().OptimizationLevel > 0)
MB.CreateIndexAssumption(Idx, MatrixTy->getNumElementsFlattened());

Expand Down
35 changes: 35 additions & 0 deletions clang/test/CodeGen/matrix-type-indexing.c
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
// RUN: %clang_cc1 -fenable-matrix -triple x86_64-apple-darwin %s -emit-llvm -disable-llvm-passes -o - | FileCheck %s

typedef float fx2x3_t __attribute__((matrix_type(2, 3)));
float Out[6];

void one_index_test(int index, fx2x3_t M) {
// CHECK-LABEL: one_index_test
// CHECK: %row = alloca i32, align 4
// CHECK: %col = alloca i32, align 4
// CHECK: [[row_load:%.*]] = load i32, ptr %row, align 4
// CHECK-NEXT: [[row_load_zext:%.*]] = zext i32 [[row_load]] to i64
// CHECK-NEXT: [[col_load:%.*]] = load i32, ptr %col, align 4
// CHECK-NEXT: [[col_load_zext:%.*]] = zext i32 [[col_load]] to i64
// CHECK-NEXT: [[col_offset:%.*]] = mul i64 [[col_load_zext]], 2
// CHECK-NEXT: [[col_major_index:%.*]] = add i64 [[col_offset]], [[row_load_zext]]
// CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x float>, ptr %M.addr, align 4
// CHECK-NEXT: %matrixext = extractelement <6 x float> [[matrix_as_vec]], i64 [[col_major_index]]
const unsigned int COLS = 3;
unsigned int row = index / COLS;
unsigned int col = index % COLS;
Out[index] = M[row][col];
}

float two_index_test(int row, int col, fx2x3_t M) {
// CHECK-LABEL: two_index_test
// CHECK: [[row_load:%.*]] = load i32, ptr %row.addr, align 4
// CHECK-NEXT: [[row_load_sext:%.*]] = sext i32 [[row_load]] to i64
// CHECK-NEXT: [[col_load:%.*]] = load i32, ptr %col.addr, align 4
// CHECK-NEXT: [[col_load_sext:%.*]] = sext i32 [[col_load]] to i64
// CHECK-NEXT: [[col_offset:%.*]] = mul i64 [[col_load_sext]], 2
// CHECK-NEXT: [[col_major_index:%.*]] = add i64 [[col_offset]], [[row_load_sext]]
// CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x float>, ptr %M.addr, align 4
// CHECK-NEXT: %matrixext = extractelement <6 x float> [[matrix_as_vec]], i64 [[col_major_index]]
return M[row][col];
}
35 changes: 35 additions & 0 deletions clang/test/CodeGenCXX/matrix-type-indexing.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
// RUN: %clang_cc1 -fenable-matrix -triple x86_64-apple-darwin %s -emit-llvm -disable-llvm-passes -o - -std=c++11 | FileCheck %s

typedef float fx2x3_t __attribute__((matrix_type(2, 3)));
float Out[6];

void one_index_test(int index, fx2x3_t M) {
// CHECK-LABEL: one_index_test
// CHECK: %row = alloca i32, align 4
// CHECK: %col = alloca i32, align 4
// CHECK: [[row_load:%.*]] = load i32, ptr %row, align 4
// CHECK-NEXT: [[row_load_zext:%.*]] = zext i32 [[row_load]] to i64
// CHECK-NEXT: [[col_load:%.*]] = load i32, ptr %col, align 4
// CHECK-NEXT: [[col_load_zext:%.*]] = zext i32 [[col_load]] to i64
// CHECK-NEXT: [[col_offset:%.*]] = mul i64 [[col_load_zext]], 2
// CHECK-NEXT: [[col_major_index:%.*]] = add i64 [[col_offset]], [[row_load_zext]]
// CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x float>, ptr %M.addr, align 4
// CHECK-NEXT: %matrixext = extractelement <6 x float> [[matrix_as_vec]], i64 [[col_major_index]]
const unsigned int COLS = 3;
unsigned int row = index / COLS;
unsigned int col = index % COLS;
Out[index] = M[row][col];
}

float two_index_test(int row, int col, fx2x3_t M) {
// CHECK-LABEL: two_index_test
// CHECK: [[row_load:%.*]] = load i32, ptr %row.addr, align 4
// CHECK-NEXT: [[row_load_sext:%.*]] = sext i32 [[row_load]] to i64
// CHECK-NEXT: [[col_load:%.*]] = load i32, ptr %col.addr, align 4
// CHECK-NEXT: [[col_load_sext:%.*]] = sext i32 [[col_load]] to i64
// CHECK-NEXT: [[col_offset:%.*]] = mul i64 [[col_load_sext]], 2
// CHECK-NEXT: [[col_major_index:%.*]] = add i64 [[col_offset]], [[row_load_sext]]
// CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x float>, ptr %M.addr, align 4
// CHECK-NEXT: %matrixext = extractelement <6 x float> [[matrix_as_vec]], i64 [[col_major_index]]
return M[row][col];
}
30 changes: 30 additions & 0 deletions clang/test/CodeGenHLSL/BasicFeatures/matrix-type-indexing.hlsl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
// RUN: %clang_cc1 -finclude-default-header -x hlsl -triple dxil-pc-shadermodel6.3-library %s -fnative-half-type -emit-llvm -disable-llvm-passes -o - | FileCheck %s

// NOTE: This test is only to confirm we can do codgen with the matrix alias.

RWBuffer<int> Out : register(u1);

void one_index_test(int index, half2x3 M) {
// CHECK-LABEL: one_index_test
// CHECK: %row = alloca i32, align 4
// CHECK: %col = alloca i32, align 4
// CHECK: [[row_load:%.*]] = load i32, ptr %row, align 4
// CHECK-NEXT: [[col_load:%.*]] = load i32, ptr %col, align 4
// CHECK-NEXT: [[row_offset:%.*]] = mul i32 [[row_load]], 3
// CHECK-NEXT: [[row_major_index:%.*]] = add i32 [[row_offset]], [[col_load]]
// CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x half>, ptr %M.addr, align 2
// CHECK-NEXT: %matrixext = extractelement <6 x half> [[matrix_as_vec]], i32 [[row_major_index]]
const uint COLS = 3;
uint row = index / COLS;
uint col = index % COLS;
Out[index] = M[row][col];
}

half two_index_test(int row, int col, half2x3 M) {
// CHECK-LABEL: two_index_test
// CHECK: [[row_offset:%.*]] = mul i32 [[row_load:%.*]], 3
// CHECK-NEXT: [[row_major_index:%.*]] = add i32 [[row_offset]], [[col_load:%.*]]
// CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x half>, ptr %M.addr, align 2
// CHECK-NEXT: %matrixext = extractelement <6 x half> [[matrix_as_vec]], i32 [[row_major_index]]
return M[row][col];
}
17 changes: 15 additions & 2 deletions llvm/include/llvm/IR/MatrixBuilder.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,8 +241,8 @@ class MatrixBuilder {

/// Compute the index to access the element at (\p RowIdx, \p ColumnIdx) from
/// a matrix with \p NumRows embedded in a vector.
Value *CreateIndex(Value *RowIdx, Value *ColumnIdx, unsigned NumRows,
Twine const &Name = "") {
Value *createColumnMajorIndex(Value *RowIdx, Value *ColumnIdx,
unsigned NumRows, Twine const &Name = "") {
unsigned MaxWidth = std::max(RowIdx->getType()->getScalarSizeInBits(),
ColumnIdx->getType()->getScalarSizeInBits());
Type *IntTy = IntegerType::get(RowIdx->getType()->getContext(), MaxWidth);
Expand All @@ -251,6 +251,19 @@ class MatrixBuilder {
Value *NumRowsV = B.getIntN(MaxWidth, NumRows);
return B.CreateAdd(B.CreateMul(ColumnIdx, NumRowsV), RowIdx);
}

/// Compute the index to access the element at (\p RowIdx, \p ColumnIdx) from
/// a matrix with \p NumRows embedded in a vector.
Copy link
Member Author

@farzonl farzonl Oct 30, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
/// a matrix with \p NumRows embedded in a vector.
/// a matrix with \p NumCols embedded in a vector.

Value *createRowMajorIndex(Value *RowIdx, Value *ColumnIdx, unsigned NumCols,
Twine const &Name = "") {
unsigned MaxWidth = std::max(RowIdx->getType()->getScalarSizeInBits(),
ColumnIdx->getType()->getScalarSizeInBits());
Type *IntTy = IntegerType::get(RowIdx->getType()->getContext(), MaxWidth);
RowIdx = B.CreateZExt(RowIdx, IntTy);
ColumnIdx = B.CreateZExt(ColumnIdx, IntTy);
Value *NumColsV = B.getIntN(MaxWidth, NumCols);
return B.CreateAdd(B.CreateMul(RowIdx, NumColsV), ColumnIdx);
}
};

} // end namespace llvm
Expand Down