[HLSL] Add _m and _<numeric> based accessors to hlsl::matrix

fixes #159438

HLSL supports 0 based index and one based index accessors that
are equivalent to the column row bracket subscripting operators.

This change adds support for these accessors by hooking into
LookupMemberExpr and adding them similar to how `CheckExtVectorComponent`
and the hlsl specific scalar to vectorsplat accessors work.

Since these accessors are HLSL specific The implementation details
are kept to SemaHLSL via `tryBuildHLSLMatrixElementAccessor`.

Author: farzonl

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -7500,7 +7500,11 @@ def ext_subscript_non_lvalue : Extension<
 def err_typecheck_subscript_value : Error<
   "subscripted value is not an array, pointer, or vector">;
 def err_typecheck_subscript_not_integer : Error<
-  "array subscript is not an integer">;
+  "%select{array|matrix}0 subscript is not an integer">;
+def err_typecheck_subscript_not_in_bounds : Error<
+  "%select{row|column}0 subscript is out of bounds of %select{zero|one}1 based indexing">;
+def err_typecheck_subscript_out_of_bounds : Error<
+  "%select{row|column}0 subscripted is out of bounds">;
 def err_subscript_function_type : Error<
   "subscript of pointer to function type %0">;
 def err_subscript_incomplete_or_sizeless_type : Error<
@@ -12972,6 +12976,7 @@ def err_builtin_matrix_stride_too_small: Error<
   "stride must be greater or equal to the number of rows">;
 def err_builtin_matrix_invalid_dimension: Error<
   "%0 dimension is outside the allowed range [1, %1]">;
+def err_builtin_matrix_invalid_member: Error<"invalid matrix member">;
 
 def warn_mismatched_import : Warning<
   "import %select{module|name}0 (%1) does not match the import %select{module|name}0 (%2) of the "

clang/include/clang/Sema/SemaHLSL.h

@@ -225,6 +225,9 @@ class SemaHLSL : public SemaBase {
   bool transformInitList(const InitializedEntity &Entity, InitListExpr *Init);
   bool handleInitialization(VarDecl *VDecl, Expr *&Init);
   void deduceAddressSpace(VarDecl *Decl);
+  ExprResult tryBuildHLSLMatrixElementAccessor(Expr *Base,
+                                               SourceLocation MemberLoc,
+                                               const IdentifierInfo *MemberId);
 
 private:
   // HLSL resource type attributes need to be processed all at once.

clang/lib/Sema/SemaExprMember.cpp

@@ -19,6 +19,7 @@
 #include "clang/Sema/Overload.h"
 #include "clang/Sema/Scope.h"
 #include "clang/Sema/ScopeInfo.h"
+#include "clang/Sema/SemaHLSL.h"
 #include "clang/Sema/SemaObjC.h"
 #include "clang/Sema/SemaOpenMP.h"
 
@@ -1669,12 +1670,19 @@ static ExprResult LookupMemberExpr(Sema &S, LookupResult &R,
 
   // HLSL supports implicit conversion of scalar types to single element vector
   // rvalues in member expressions.
-  if (S.getLangOpts().HLSL && BaseType->isScalarType()) {
-    QualType VectorTy = S.Context.getExtVectorType(BaseType, 1);
-    BaseExpr = S.ImpCastExprToType(BaseExpr.get(), VectorTy, CK_VectorSplat,
-                                   BaseExpr.get()->getValueKind());
-    return LookupMemberExpr(S, R, BaseExpr, IsArrow, OpLoc, SS, ObjCImpDecl,
-                            HasTemplateArgs, TemplateKWLoc);
+  if (S.getLangOpts().HLSL) {
+    if (BaseType->isScalarType()) {
+      QualType VectorTy = S.Context.getExtVectorType(BaseType, 1);
+      BaseExpr = S.ImpCastExprToType(BaseExpr.get(), VectorTy, CK_VectorSplat,
+                                     BaseExpr.get()->getValueKind());
+      return LookupMemberExpr(S, R, BaseExpr, IsArrow, OpLoc, SS, ObjCImpDecl,
+                              HasTemplateArgs, TemplateKWLoc);
+    }
+    if (!IsArrow && BaseType->isConstantMatrixType()) {
+      if (const auto *II = MemberName.getAsIdentifierInfo())
+        return S.HLSL().tryBuildHLSLMatrixElementAccessor(BaseExpr.get(),
+                                                          MemberLoc, II);
+    }
   }
 
   S.Diag(OpLoc, diag::err_typecheck_member_reference_struct_union)

clang/lib/Sema/SemaHLSL.cpp

@@ -21,6 +21,7 @@
 #include "clang/AST/Expr.h"
 #include "clang/AST/HLSLResource.h"
 #include "clang/AST/Type.h"
+#include "clang/AST/TypeBase.h"
 #include "clang/AST/TypeLoc.h"
 #include "clang/Basic/Builtins.h"
 #include "clang/Basic/DiagnosticSema.h"
@@ -31,6 +32,7 @@
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Sema/Initialization.h"
 #include "clang/Sema/Lookup.h"
+#include "clang/Sema/Ownership.h"
 #include "clang/Sema/ParsedAttr.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Sema/Template.h"
@@ -4313,6 +4315,101 @@ bool SemaHLSL::transformInitList(const InitializedEntity &Entity,
   return true;
 }
 
+ExprResult SemaHLSL::tryBuildHLSLMatrixElementAccessor(
+    Expr *Base, SourceLocation MemberLoc, const IdentifierInfo *MemberId) {
+  if (!Base)
+    return ExprError();
+
+  QualType T = Base->getType();
+  const ConstantMatrixType *MT = T->getAs<ConstantMatrixType>();
+  if (!MT)
+    return ExprError();
+
+  auto parseHLSLMatrixAccessor =
+      [this, MemberLoc, MemberId](unsigned &Row, unsigned &Col) -> ExprResult {
+    StringRef Name = MemberId->getNameStart();
+    bool IsMaccessor = Name.consume_front("_m");
+    // consume numeric accessor second so if _m exist we don't consume _ to
+    // early.
+    bool IsNumericAccessor = Name.consume_front("_");
+    if (!IsMaccessor && !IsNumericAccessor)
+      return ExprError(
+          Diag(MemberLoc, diag::err_builtin_matrix_invalid_member));
+
+    auto isDigit = [](char c) { return c >= '0' && c <= '9'; };
+    auto isZeroBasedIndex = [](char c) { return c >= '0' && c <= '3'; };
+    auto isOneBasedIndex = [](char c) { return c >= '1' && c <= '4'; };
+    if (Name.empty() || !isDigit(Name[0]) || !isDigit(Name[1])) {
+      return ExprError(
+          Diag(MemberLoc, diag::err_typecheck_subscript_not_integer)
+          << /*matrix*/ 1);
+    }
+
+    Row = Name[0] - '0';
+    Col = Name[1] - '0';
+    bool HasIndexingError = false;
+    if (IsNumericAccessor) {
+      Row--;
+      Col--;
+      // Note: we add diagnostic errors here because otherwise we will return -1
+      if (!isOneBasedIndex(Name[0])) {
+        Diag(MemberLoc, diag::err_typecheck_subscript_not_in_bounds)
+            << /*row*/ 0 << /*one*/ 1;
+        HasIndexingError = true;
+      }
+      if (!isOneBasedIndex(Name[1])) {
+        Diag(MemberLoc, diag::err_typecheck_subscript_not_in_bounds)
+            << /*col*/ 1 << /*one*/ 1;
+        HasIndexingError = true;
+      }
+    } else {
+      if (!isZeroBasedIndex(Name[0])) {
+        Diag(MemberLoc, diag::err_typecheck_subscript_not_in_bounds)
+            << /*row*/ 0 << /*zero*/ 0;
+        HasIndexingError = true;
+      }
+      if (!isZeroBasedIndex(Name[1])) {
+        Diag(MemberLoc, diag::err_typecheck_subscript_not_in_bounds)
+            << /*col*/ 1 << /*zero*/ 0;
+        HasIndexingError = true;
+      }
+    }
+    if (HasIndexingError)
+      return ExprError();
+    return ExprEmpty();
+  };
+  unsigned Row = 0, Col = 0;
+  ExprResult ParseResult = parseHLSLMatrixAccessor(Row, Col);
+  if (ParseResult.isInvalid())
+    return ParseResult;
+
+  unsigned Rows = MT->getNumRows();
+  unsigned Cols = MT->getNumColumns();
+  bool HasBoundsError = false;
+  if (Row >= Rows) {
+    Diag(MemberLoc, diag::err_typecheck_subscript_out_of_bounds) << /*Row*/ 0;
+    HasBoundsError = true;
+  }
+  if (Col >= Cols) {
+    Diag(MemberLoc, diag::err_typecheck_subscript_out_of_bounds) << /*Col*/ 1;
+    HasBoundsError = true;
+  }
+  if (HasBoundsError)
+    return ExprError();
+
+  auto mkIdx = [&](unsigned v) -> Expr * {
+    ASTContext &Context = SemaRef.getASTContext();
+    return IntegerLiteral::Create(Context, llvm::APInt(32, v),
+                                  Context.UnsignedIntTy, MemberLoc);
+  };
+  Expr *RowIdx = mkIdx(Row);
+  Expr *ColIdx = mkIdx(Col);
+
+  // Build A[Row][Col], reusing the existing matrix subscript machinery.
+  return SemaRef.CreateBuiltinMatrixSubscriptExpr(Base, RowIdx, ColIdx,
+                                                  MemberLoc);
+}
+
 bool SemaHLSL::handleInitialization(VarDecl *VDecl, Expr *&Init) {
   const HLSLVkConstantIdAttr *ConstIdAttr =
       VDecl->getAttr<HLSLVkConstantIdAttr>();

clang/test/AST/HLSL/matrix-member-access.hlsl

@@ -0,0 +1,39 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-compute -x hlsl -ast-dump -o - %s | FileCheck %s 
+
+typedef float float3x3 __attribute__((matrix_type(3,3)));
+
+[numthreads(1,1,1)]
+void ok() {
+    float3x3 A;
+    // CHECK: BinaryOperator 0x{{[0-9a-fA-F]+}} <line:{{[0-9]+}}:{{[0-9]+}}, col:{{[0-9]+}}> 'float' lvalue matrixcomponent '='
+    // CHECK-NEXT: MatrixSubscriptExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> 'float' lvalue matrixcomponent
+    // CHECK-NEXT: DeclRefExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'float3x3':'matrix<float, 3, 3>' lvalue Var 0x{{[0-9a-fA-F]+}} 'A' 'float3x3':'matrix<float, 3, 3>'
+    // CHECK-NEXT: IntegerLiteral 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'unsigned int' 1
+    // CHECK-NEXT: IntegerLiteral 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'unsigned int' 2
+    // CHECK-NEXT: FloatingLiteral 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'float' 3.140000e+00
+    A._m12 = 3.14;
+
+    // CHECK: VarDecl 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> col:{{[0-9]+}} r 'float' cinit
+    // CHECK-NEXT: ImplicitCastExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> 'float' <LValueToRValue>
+    // CHECK-NEXT: MatrixSubscriptExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> 'float' lvalue matrixcomponent
+    // CHECK-NEXT: DeclRefExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'float3x3':'matrix<float, 3, 3>' lvalue Var 0x{{[0-9a-fA-F]+}} 'A' 'float3x3':'matrix<float, 3, 3>'
+    // CHECK-NEXT: IntegerLiteral 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'unsigned int' 0
+    // CHECK-NEXT: IntegerLiteral 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'unsigned int' 0
+    float r = A._m00;
+
+    // CHECK: VarDecl 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> col:{{[0-9]+}} good1 'float' cinit
+    // CHECK-NEXT: ImplicitCastExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> 'float' <LValueToRValue>
+    // CHECK-NEXT: MatrixSubscriptExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> 'float' lvalue matrixcomponent
+    // CHECK-NEXT: DeclRefExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'float3x3':'matrix<float, 3, 3>' lvalue Var 0x{{[0-9a-fA-F]+}} 'A' 'float3x3':'matrix<float, 3, 3>'
+    // CHECK-NEXT: IntegerLiteral 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'unsigned int' 0
+    // CHECK-NEXT: IntegerLiteral 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'unsigned int' 0
+    float good1 = A._11;
+
+    // CHECK: VarDecl 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> col:{{[0-9]+}} good2 'float' cinit
+    // CHECK-NEXT: ImplicitCastExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> 'float' <LValueToRValue>
+    // CHECK-NEXT: MatrixSubscriptExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}, col:{{[0-9]+}}> 'float' lvalue matrixcomponent
+    // CHECK-NEXT: DeclRefExpr 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'float3x3':'matrix<float, 3, 3>' lvalue Var 0x{{[0-9a-fA-F]+}} 'A' 'float3x3':'matrix<float, 3, 3>'
+    // CHECK-NEXT: IntegerLiteral 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'unsigned int' 2
+    // CHECK-NEXT: IntegerLiteral 0x{{[0-9a-fA-F]+}} <col:{{[0-9]+}}> 'unsigned int' 2    
+    float good2 = A._33;
+}

clang/test/CodeGenHLSL/matrix-member-access.hlsl

@@ -0,0 +1,25 @@
+// RUN: %clang_cc1 -std=hlsl202x -finclude-default-header -x hlsl -triple \
+// RUN:   dxil-pc-shadermodel6.3-library %s -fnative-half-type \
+// RUN:   -emit-llvm -disable-llvm-passes -o - | FileCheck %s
+
+
+half3x3 write_pi(half3x3 A) {
+    //CHECK:  [[MAT_INS:%.*]] = insertelement <9 x half> %{{[0-9]+}}, half 0xH4248, i32 7
+    //CHECK-NEXT:  store <9 x half> [[MAT_INS]], ptr %{{.*}}, align 2
+    A._m12 = 3.14;
+    return A;
+}
+
+half read_1x1(half3x3 A) {
+    //CHECK:  [[MAT_EXT:%.*]] = extractelement <9 x half> %{{[0-9]+}}, i32 0
+    return A._11;
+}
+half read_m0x0(half3x3 A) {
+    //CHECK:  [[MAT_EXT:%.*]] = extractelement <9 x half> %{{[0-9]+}}, i32 0
+    return A._m00;
+}
+
+half read_3x3(half3x3 A) {
+    //CHECK:  [[MAT_EXT:%.*]] = extractelement <9 x half> %{{[0-9]+}}, i32 8
+    return A._33;
+}

clang/test/SemaHLSL/matrix-member-access.hlsl

@@ -0,0 +1,20 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -finclude-default-header -verify %s
+
+void foo() {
+    float3x3 A;
+    float r = A._m00;      // read is ok
+    float good1 = A._11;    
+    float good2 = A._33;
+
+    float bad0 = A._m44;    // expected-error {{row subscript is out of bounds of zero based indexing}} expected-error {{column subscript is out of bounds of zero based indexing}}
+    float bad1 = A._m33;    // expected-error {{row subscripted is out of bounds}} expected-error {{column subscripted is out of bounds}}
+    float bad2 = A._mA2;    // expected-error {{matrix subscript is not an integer}}
+    float bad3 = A._m2F;    // expected-error {{matrix subscript is not an integer}}
+
+    float bad4 = A._00;    // expected-error {{row subscript is out of bounds of one based indexing}} expected-error {{column subscript is out of bounds of one based indexing}}
+    float bad5 = A._44;    // expected-error {{row subscripted is out of bounds}} expected-error {{column subscripted is out of bounds}}
+    float bad6 = A._55;    // expected-error {{row subscript is out of bounds of one based indexing}} expected-error {{column subscript is out of bounds of one based indexing}}
+
+    
+    A._m12 = 3.14;         // write is OK
+}