llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT--> @llvm/pr-subscribers-hlsl Author: Farzon Lotfi (farzonl) <details> <summary>Changes</summary> fixes #<!-- -->167617 In DXC HLSL supports different indexing modes via codegen for its equivalent of the MatrixSubscriptExpr when the /Zpr and /Zpc flags are used see: https://godbolt.org/z/bz5Y5WG36. This change modifies EmitMatrixSubscriptExpr to consider the MatrixRowMajor/MatrixColMajor Layout flags before generating an index. Similarly it introduces `createRowMajorIndex` and `createColumnMajorIndex` in `MatrixBuilder.h` for use in `VisitMatrixSubscriptExpr`. --- Full diff: https://github.com/llvm/llvm-project/pull/171564.diff 6 Files Affected: - (modified) clang/lib/CodeGen/CGExpr.cpp (+13-5) - (modified) clang/lib/CodeGen/CGExprScalar.cpp (+11-2) - (added) clang/test/CodeGen/matrix-type-indexing.c (+60) - (added) clang/test/CodeGenCXX/matrix-type-indexing.cpp (+36) - (added) clang/test/CodeGenHLSL/BasicFeatures/matrix-type-indexing.hlsl (+52) - (modified) llvm/include/llvm/IR/MatrixBuilder.h (+15-2) ``````````diff diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp index 3bde8e1fa2ac3..e47ce7050fa0f 100644 --- a/clang/lib/CodeGen/CGExpr.cpp +++ b/clang/lib/CodeGen/CGExpr.cpp @@ -4884,11 +4884,19 @@ 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().getDefaultMatrixMemoryLayout() == + LangOptions::MatrixMemoryLayout::MatrixRowMajor) { + 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()); diff --git a/clang/lib/CodeGen/CGExprScalar.cpp b/clang/lib/CodeGen/CGExprScalar.cpp index 769bc37b0e131..c5147eebb2adc 100644 --- a/clang/lib/CodeGen/CGExprScalar.cpp +++ b/clang/lib/CodeGen/CGExprScalar.cpp @@ -2118,9 +2118,18 @@ 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().getDefaultMatrixMemoryLayout() == + LangOptions::MatrixMemoryLayout::MatrixRowMajor) { + 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()); diff --git a/clang/test/CodeGen/matrix-type-indexing.c b/clang/test/CodeGen/matrix-type-indexing.c new file mode 100644 index 0000000000000..d76d14c3f67ef --- /dev/null +++ b/clang/test/CodeGen/matrix-type-indexing.c @@ -0,0 +1,60 @@ +// RUN: %clang_cc1 -fenable-matrix -fmatrix-memory-layout=row-major -triple x86_64-apple-darwin %s -emit-llvm -disable-llvm-passes -o - | FileCheck %s --check-prefixes=CHECK,ROW-CHECK +// RUN: %clang_cc1 -fenable-matrix -fmatrix-memory-layout=column-major -triple x86_64-apple-darwin %s -emit-llvm -disable-llvm-passes -o - | FileCheck %s --check-prefixes=CHECK,COL-CHECK +// RUN: %clang_cc1 -fenable-matrix -triple x86_64-apple-darwin %s -emit-llvm -disable-llvm-passes -o - | FileCheck %s --check-prefixes=CHECK,COL-CHECK + +typedef float fx2x3_t __attribute__((matrix_type(2, 3))); + float Out[6]; + + fx2x3_t gM; + +void binaryOpMatrixSubscriptExpr(int index, fx2x3_t M) { + // CHECK-LABEL: binaryOpMatrixSubscriptExpr + // 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 + // COL-CHECK-NEXT: [[col_offset:%.*]] = mul i64 [[col_load_zext]], 2 + // COL-CHECK-NEXT: [[col_major_index:%.*]] = add i64 [[col_offset]], [[row_load_zext]] + // ROW-CHECK-NEXT: [[row_offset:%.*]] = mul i64 [[row_load_zext]], 3 + // ROW-CHECK-NEXT: [[row_major_index:%.*]] = add i64 [[row_offset]], [[col_load_zext]] + // CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x float>, ptr %M.addr, align 4 + // COL-CHECK-NEXT: %matrixext = extractelement <6 x float> [[matrix_as_vec]], i64 [[col_major_index]] + // ROW-CHECK-NEXT: %matrixext = extractelement <6 x float> [[matrix_as_vec]], i64 [[row_major_index]] + const unsigned int COLS = 3; + unsigned int row = index / COLS; + unsigned int col = index % COLS; + Out[index] = M[row][col]; +} + +float returnMatrixSubscriptExpr(int row, int col, fx2x3_t M) { + // CHECK-LABEL: returnMatrixSubscriptExpr + // CHECK: [[row_load:%.*]] = load i32, ptr [[row_ptr:%.*]], align 4 + // CHECK-NEXT: [[row_load_sext:%.*]] = sext i32 [[row_load]] to i64 + // CHECK-NEXT: [[col_load:%.*]] = load i32, ptr [[col_ptr:%.*]], align 4 + // CHECK-NEXT: [[col_load_sext:%.*]] = sext i32 [[col_load]] to i64 + // COL-CHECK-NEXT: [[col_offset:%.*]] = mul i64 [[col_load_sext]], 2 + // COL-CHECK-NEXT: [[col_major_index:%.*]] = add i64 [[col_offset]], [[row_load_sext]] + // ROW-CHECK-NEXT: [[row_offset:%.*]] = mul i64 [[row_load_sext]], 3 + // ROW-CHECK-NEXT: [[row_major_index:%.*]] = add i64 [[row_offset]], [[col_load_sext]] + // CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x float>, ptr %M.addr, align 4 + // COL-CHECK-NEXT: [[matrix_after_extract:%.*]] = extractelement <6 x float> [[matrix_as_vec]], i64 [[col_major_index]] + // ROW-CHECK-NEXT: [[matrix_after_extract:%.*]] = extractelement <6 x float> [[matrix_as_vec]], i64 [[row_major_index]] + // CHECK-NEXT: ret float [[matrix_after_extract]] + return M[row][col]; +} + +void storeAtMatrixSubscriptExpr(int row, int col, float value) { + // CHECK-LABEL: storeAtMatrixSubscriptExpr + // CHECK: [[value_load:%.*]] = load float, ptr [[value_ptr:%.*]], align 4 + // ROW-CHECK: [[row_offset:%.*]] = mul i64 [[row_load:%.*]], 3 + // ROW-CHECK-NEXT: [[row_major_index:%.*]] = add i64 [[row_offset]], [[col_load:%.*]] + // COL-CHECK: [[col_offset:%.*]] = mul i64 [[col_load:%.*]], 2 + // COL-CHECK-NEXT: [[col_major_index:%.*]] = add i64 [[col_offset]], [[row_load:%.*]] + // CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x float>, ptr @gM, align 4 + // ROW-CHECK-NEXT: [[matrix_after_insert:%.*]] = insertelement <6 x float> [[matrix_as_vec]], float [[value_load]], i64 [[row_major_index]] + // COL-CHECK-NEXT: [[matrix_after_insert:%.*]] = insertelement <6 x float> [[matrix_as_vec]], float [[value_load]], i64 [[col_major_index]] + // CHECK-NEXT: store <6 x float> [[matrix_after_insert]], ptr @gM, align 4 + gM[row][col] = value; +} diff --git a/clang/test/CodeGenCXX/matrix-type-indexing.cpp b/clang/test/CodeGenCXX/matrix-type-indexing.cpp new file mode 100644 index 0000000000000..a390bcaf89119 --- /dev/null +++ b/clang/test/CodeGenCXX/matrix-type-indexing.cpp @@ -0,0 +1,36 @@ +// RUN: %clang_cc1 -fenable-matrix -fmatrix-memory-layout=column-major -triple x86_64-apple-darwin %s -emit-llvm -disable-llvm-passes -o - -std=c++11 | FileCheck %s +// 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 binaryOpMatrixSubscriptExpr(int index, fx2x3_t M) { + // CHECK-LABEL: binaryOpMatrixSubscriptExpr + // 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 returnMatrixSubscriptExpr(int row, int col, fx2x3_t M) { + // CHECK-LABEL: returnMatrixSubscriptExpr + // 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]; +} diff --git a/clang/test/CodeGenHLSL/BasicFeatures/matrix-type-indexing.hlsl b/clang/test/CodeGenHLSL/BasicFeatures/matrix-type-indexing.hlsl new file mode 100644 index 0000000000000..7a63bbb45ecf7 --- /dev/null +++ b/clang/test/CodeGenHLSL/BasicFeatures/matrix-type-indexing.hlsl @@ -0,0 +1,52 @@ +// RUN: %clang_cc1 -finclude-default-header -x hlsl -triple dxil-pc-shadermodel6.3-library %s -fnative-half-type -emit-llvm -disable-llvm-passes -fmatrix-memory-layout=row-major -o - | FileCheck %s --check-prefixes=CHECK,ROW-CHECK +// RUN: %clang_cc1 -finclude-default-header -x hlsl -triple dxil-pc-shadermodel6.3-library %s -fnative-half-type -emit-llvm -disable-llvm-passes -fmatrix-memory-layout=column-major -o - | FileCheck %s --check-prefixes=CHECK,COL-CHECK +// 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 --check-prefixes=CHECK,COL-CHECK + +RWBuffer<int> Out : register(u1); +half2x3 gM; + + +void binaryOpMatrixSubscriptExpr(int index, half2x3 M) { + // CHECK-LABEL: binaryOpMatrixSubscriptExpr + // 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 + // ROW-CHECK-NEXT: [[row_offset:%.*]] = mul i32 [[row_load]], 3 + // ROW-CHECK-NEXT: [[row_major_index:%.*]] = add i32 [[row_offset]], [[col_load]] + // COL-CHECK-NEXT: [[col_offset:%.*]] = mul i32 [[col_load]], 2 + // COL-CHECK-NEXT: [[col_major_index:%.*]] = add i32 [[col_offset]], [[row_load]] + // CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x half>, ptr %M.addr, align 2 + // ROW-CHECK-NEXT: %matrixext = extractelement <6 x half> [[matrix_as_vec]], i32 [[row_major_index]] + // COL-CHECK-NEXT: %matrixext = extractelement <6 x half> [[matrix_as_vec]], i32 [[col_major_index]] + const uint COLS = 3; + uint row = index / COLS; + uint col = index % COLS; + Out[index] = M[row][col]; +} + +half returnMatrixSubscriptExpr(int row, int col, half2x3 M) { + // CHECK-LABEL: returnMatrixSubscriptExpr + // ROW-CHECK: [[row_offset:%.*]] = mul i32 [[row_load:%.*]], 3 + // ROW-CHECK-NEXT: [[row_major_index:%.*]] = add i32 [[row_offset]], [[col_load:%.*]] + // COL-CHECK: [[col_offset:%.*]] = mul i32 [[col_load:%.*]], 2 + // COL-CHECK-NEXT: [[col_major_index:%.*]] = add i32 [[col_offset]], [[row_load:%.*]] + // CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x half>, ptr %M.addr, align 2 + // ROW-CHECK-NEXT: %matrixext = extractelement <6 x half> [[matrix_as_vec]], i32 [[row_major_index]] + // COL-CHECK-NEXT: %matrixext = extractelement <6 x half> [[matrix_as_vec]], i32 [[col_major_index]] + return M[row][col]; +} + +void storeAtMatrixSubscriptExpr(int row, int col, half value) { + // CHECK-LABEL: storeAtMatrixSubscriptExpr + // CHECK: [[value_load:%.*]] = load half, ptr [[value_ptr:%.*]], align 2 + // ROW-CHECK: [[row_offset:%.*]] = mul i32 [[row_load:%.*]], 3 + // ROW-CHECK-NEXT: [[row_major_index:%.*]] = add i32 [[row_offset]], [[col_load:%.*]] + // COL-CHECK: [[col_offset:%.*]] = mul i32 [[col_load:%.*]], 2 + // COL-CHECK-NEXT: [[col_major_index:%.*]] = add i32 [[col_offset]], [[row_load:%.*]] + // CHECK-NEXT: [[matrix_as_vec:%.*]] = load <6 x half>, ptr addrspace(2) @gM, align 2 + // ROW-CHECK-NEXT: [[matrix_after_insert:%.*]] = insertelement <6 x half> [[matrix_as_vec]], half [[value_load]], i32 [[row_major_index]] + // COL-CHECK-NEXT: [[matrix_after_insert:%.*]] = insertelement <6 x half> [[matrix_as_vec]], half [[value_load]], i32 [[col_major_index]] + // CHECK-NEXT: store <6 x half> [[matrix_after_insert]], ptr addrspace(2) @gM, align 2 + gM[row][col] = value; +} diff --git a/llvm/include/llvm/IR/MatrixBuilder.h b/llvm/include/llvm/IR/MatrixBuilder.h index 3a04ca87f2b55..0cea0f0ddafa3 100644 --- a/llvm/include/llvm/IR/MatrixBuilder.h +++ b/llvm/include/llvm/IR/MatrixBuilder.h @@ -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); @@ -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 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 `````````` </details> https://github.com/llvm/llvm-project/pull/171564 _______________________________________________ cfe-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
