[Matrix] Emit assumption that matrix indices are valid.

[fhahn]

Summary

The matrix extension requires the indices for matrix subscript
expression to be valid and it is UB otherwise.

extract/insertelement produce poison if the index is invalid, which
limits the optimizer to not be bale to scalarize load/extract pairs for
example, which causes very suboptimal code to be generated when using
matrix subscript expressions with variable indices for large matrixes.

This patch updates IRGen to emit assumes to for index expression to
convey the information that the index must be valid.

This also adjusts the order in which operations are emitted slightly, so
indices & assumes are added before the load of the matrix value.

[fhahn]

Thanks for taking a look! I added a phase ordering test and updated the pipeline tests as well.

[fhahn]

Undo unrelated changes.

[xbolva00]

llvm/include/llvm/IR/MatrixBuilder.h

242

Prefer early exit?

[xbolva00]

llvm/include/llvm/IR/MatrixBuilder.h

245

B.CreateAssumption(Cmp) may work well as well? and shorter.

[rjmccall]

Do you want to generate these even at -O0?

llvm/include/llvm/IR/MatrixBuilder.h

242

Should this do something special if the index is statically out of bounds?

[fhahn]

Thanks for the comments!

Updates:

• Only emit assumptions for optimized builds; the assumes are not helpful for unoptimized builds
• Use builder::CreateAssumption instead of manually constructing assume calls
• Add assert that condition must be true, if simplified to a constant.

[fhahn]

llvm/include/llvm/IR/MatrixBuilder.h

242

I think this cause should never happen and if it does it means earlier checks have been missed. I added an assert.

245

Updated, thanks!

[fhahn]

Fix clang-tidy warnings.

[fhahn]

Fix failing clang/test/CodeGenObjC/matrix-type-operators.m.

[fhahn]

ping :)

[erichkeane]

Just a couple of nits here, basically see how much we can put in the 'cheap to check' branch.

clang/lib/CodeGen/CGExpr.cpp

1944

I think the linter here is right, we require const auto * here.

1946

Is it worth putting the MatrixBuilder line inside of this branch too? Perhaps all the added code?

2091

Same here on both.

[fhahn]

In D102478#2817768, @erichkeane wrote:

Just a couple of nits here, basically see how much we can put in the 'cheap to check' branch.

Thanks! I moved the suggested bits inside the then branch.

[fhahn]

rebased after latest changes

[fhahn]

All required changes to make use of this have recently landed or are ready to land. So I am going to commit this momentarily.

clang/lib/CodeGen/CGExpr.cpp

 #include "clang/Basic/CodeGenOptions.h"
 #include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/MatrixBuilder.h"
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Transforms/Utils/SanitizerStats.h"
 
 #include <string>
 
if (LV.isExtVectorElt()) {
     return EmitLoadOfExtVectorElementLValue(LV);
   }
 
   // Global Register variables always invoke intrinsics
   if (LV.isGlobalReg())
     return EmitLoadOfGlobalRegLValue(LV);
 
   if (LV.isMatrixElt()) {
+    llvm::Value *Idx = LV.getMatrixIdx();
+    auto MatTy = LV.getType()->getAs<ConstantMatrixType>();
+    llvm::MatrixBuilder<CGBuilderTy> MB(Builder);
+    MB.CreateIndexAssumption(Idx, MatTy->getNumElementsFlattened());
     llvm::LoadInst *Load =
         Builder.CreateLoad(LV.getMatrixAddress(), LV.isVolatileQualified());
     return RValue::get(
         Builder.CreateExtractElement(Load, LV.getMatrixIdx(), "matrixext"));
   }
 
   assert(LV.isBitField() && "Unknown LValue type!");
   return EmitLoadOfBitfieldLValue(LV, Loc);
if (!Dst.isSimple()) {
     // appropriate.
     if (Dst.isExtVectorElt())
       return EmitStoreThroughExtVectorComponentLValue(Src, Dst);
 
     if (Dst.isGlobalReg())
       return EmitStoreThroughGlobalRegLValue(Src, Dst);
 
     if (Dst.isMatrixElt()) {
-      llvm::Value *Vec = Builder.CreateLoad(Dst.getMatrixAddress());
-      Vec = Builder.CreateInsertElement(Vec, Src.getScalarVal(),
-                                        Dst.getMatrixIdx(), "matins");
+      llvm::Value *Idx = Dst.getMatrixIdx();
+      auto MatTy = Dst.getType()->getAs<ConstantMatrixType>();
+      llvm::MatrixBuilder<CGBuilderTy> MB(Builder);
+      MB.CreateIndexAssumption(Idx, MatTy->getNumElementsFlattened());
+      llvm::Instruction *Load = Builder.CreateLoad(Dst.getMatrixAddress());
+      llvm::Value *Vec =
+          Builder.CreateInsertElement(Load, Src.getScalarVal(), Idx, "matins");
       Builder.CreateStore(Vec, Dst.getMatrixAddress(),
                           Dst.isVolatileQualified());
       return;
     }
 
     assert(Dst.isBitField() && "Unknown LValue type");
     return EmitStoreThroughBitfieldLValue(Src, Dst);
   }

clang/lib/CodeGen/CGExprScalar.cpp

 
 Value *ScalarExprEmitter::VisitMatrixSubscriptExpr(MatrixSubscriptExpr *E) {
   TestAndClearIgnoreResultAssign();
 
   // Handle the vector case.  The base must be a vector, the index must be an
   // integer value.
   Value *RowIdx = Visit(E->getRowIdx());
   Value *ColumnIdx = Visit(E->getColumnIdx());
+
+  auto *MatrixTy = E->getBase()->getType()->castAs<ConstantMatrixType>();
+  unsigned NumRows = MatrixTy->getNumRows();
+  llvm::MatrixBuilder<CGBuilderTy> MB(Builder);
+  Value *Idx = MB.CreateIndex(RowIdx, ColumnIdx, NumRows);
+  MB.CreateIndexAssumption(Idx, MatrixTy->getNumElementsFlattened());
+
   Value *Matrix = Visit(E->getBase());
 
   // TODO: Should we emit bounds checks with SanitizerKind::ArrayBounds?
-  llvm::MatrixBuilder<CGBuilderTy> MB(Builder);
-  return MB.CreateExtractElement(
-      Matrix, RowIdx, ColumnIdx,
-      E->getBase()->getType()->castAs<ConstantMatrixType>()->getNumRows());
+  return Builder.CreateExtractElement(Matrix, Idx, "matrixext");
 }
 
 static int getMaskElt(llvm::ShuffleVectorInst *SVI, unsigned Idx,
                       unsigned Off) {
   int MV = SVI->getMaskValue(Idx);
   if (MV == -1)
     return -1;
   return Off + MV;

clang/test/CodeGen/matrix-type-operators.c

void insert_float_matrix_idx_i_u_float(fx2x3_t b, float e, int j, unsigned k) {
   // CHECK-LABEL: @insert_float_matrix_idx_i_u_float(
   // CHECK:         [[E:%.*]] = load float, float* %e.addr, align 4
   // CHECK-NEXT:    [[J:%.*]] = load i32, i32* %j.addr, align 4
   // CHECK-NEXT:    [[J_EXT:%.*]] = sext i32 [[J]] to i64
   // CHECK-NEXT:    [[K:%.*]] = load i32, i32* %k.addr, align 4
   // CHECK-NEXT:    [[K_EXT:%.*]] = zext i32 [[K]] to i64
   // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[K_EXT]], 2
   // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[J_EXT]]
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 6
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[MAT:%.*]] = load <6 x float>, <6 x float>* [[MAT_ADDR:%.*]], align 4
   // CHECK-NEXT:    [[MATINS:%.*]] = insertelement <6 x float> [[MAT]], float [[E]], i64 [[IDX2]]
   // CHECK-NEXT:    store <6 x float> [[MATINS]], <6 x float>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:    ret void
 
   b[j][k] = e;
 }
 
 void insert_float_matrix_idx_s_ull_float(fx2x3_t b, float e, short j, unsigned long long k) {
   // CHECK-LABEL: @insert_float_matrix_idx_s_ull_float(
   // CHECK:         [[E:%.*]] = load float, float* %e.addr, align 4
   // CHECK-NEXT:    [[J:%.*]] = load i16, i16* %j.addr, align 2
   // CHECK-NEXT:    [[J_EXT:%.*]] = sext i16 [[J]] to i64
   // CHECK-NEXT:    [[K:%.*]] = load i64, i64* %k.addr, align 8
   // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[K]], 2
   // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[J_EXT]]
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 6
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[MAT:%.*]] = load <6 x float>, <6 x float>* [[MAT_ADDR:%.*]], align 4
   // CHECK-NEXT:    [[MATINS:%.*]] = insertelement <6 x float> [[MAT]], float [[E]], i64 [[IDX2]]
   // CHECK-NEXT:    store <6 x float> [[MATINS]], <6 x float>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:    ret void
 
   (b)[j][k] = e;
 }
 
 // Check that we can can use matrix index expressions on integer matrixes.
 typedef int ix9x3_t __attribute__((matrix_type(9, 3)));
 void insert_int_idx_expr(ix9x3_t a, int i) {
   // CHECK-LABEL: @insert_int_idx_expr(
   // CHECK:         [[I1:%.*]] = load i32, i32* %i.addr, align 4
   // CHECK-NEXT:    [[I2:%.*]] = load i32, i32* %i.addr, align 4
   // CHECK-NEXT:    [[I2_ADD:%.*]] = add nsw i32 4, [[I2]]
   // CHECK-NEXT:    [[ADD_EXT:%.*]] = sext i32 [[I2_ADD]] to i64
   // CHECK-NEXT:    [[IDX2:%.*]] = add i64 18, [[ADD_EXT]]
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 27
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[MAT:%.*]] = load <27 x i32>, <27 x i32>* [[MAT_ADDR:%.*]], align 4
   // CHECK-NEXT:    [[MATINS:%.*]] = insertelement <27 x i32> [[MAT]], i32 [[I1]], i64 [[IDX2]]
   // CHECK-NEXT:    store <27 x i32> [[MATINS]], <27 x i32>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:    ret void
 
   a[4 + i][1 + 1u] = i;
 }
 
double extract_float(fx3x3_t b) {
 
   return b[2][1];
 }
 
 int extract_int(ix9x3_t c, unsigned long j) {
   // CHECK-LABEL: @extract_int(
   // CHECK:         [[J1:%.*]] = load i64, i64* %j.addr, align 8
   // CHECK-NEXT:    [[J2:%.*]] = load i64, i64* %j.addr, align 8
-  // CHECK-NEXT:    [[MAT:%.*]] = load <27 x i32>, <27 x i32>* {{.*}}, align 4
   // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[J2]], 9
   // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[J1]]
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 27
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+  // CHECK-NEXT:    [[MAT:%.*]] = load <27 x i32>, <27 x i32>* {{.*}}, align 4
   // CHECK-NEXT:    [[MATEXT:%.*]] = extractelement <27 x i32> [[MAT]], i64 [[IDX2]]
   // CHECK-NEXT:    ret i32 [[MATEXT]]
 
   return c[j][j];
 }
 
 typedef double dx3x2_t __attribute__((matrix_type(3, 2)));
 
 double test_extract_matrix_pointer1(dx3x2_t **ptr, unsigned j) {
   // CHECK-LABEL: @test_extract_matrix_pointer1(
   // CHECK:         [[J:%.*]] = load i32, i32* %j.addr, align 4
   // CHECK-NEXT:    [[J_EXT:%.*]] = zext i32 [[J]] to i64
+  // CHECK-NEXT:    [[IDX:%.*]] = add i64 3, [[J_EXT]]
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX]], 6
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[PTR:%.*]] = load [6 x double]**, [6 x double]*** %ptr.addr, align 8
   // CHECK-NEXT:    [[PTR_IDX:%.*]] = getelementptr inbounds [6 x double]*, [6 x double]** [[PTR]], i64 1
   // CHECK-NEXT:    [[PTR2:%.*]] = load [6 x double]*, [6 x double]** [[PTR_IDX]], align 8
   // CHECK-NEXT:    [[PTR2_IDX:%.*]] = getelementptr inbounds [6 x double], [6 x double]* [[PTR2]], i64 2
   // CHECK-NEXT:    [[MAT_ADDR:%.*]] = bitcast [6 x double]* [[PTR2_IDX]] to <6 x double>*
   // CHECK-NEXT:    [[MAT:%.*]] = load <6 x double>, <6 x double>* [[MAT_ADDR]], align 8
-  // CHECK-NEXT:    [[IDX:%.*]] = add i64 3, [[J_EXT]]
   // CHECK-NEXT:    [[MATEXT:%.*]] = extractelement <6 x double> [[MAT]], i64 [[IDX]]
   // CHECK-NEXT:    ret double [[MATEXT]]
 
   return ptr[1][2][j][1];
 }
 
 double test_extract_matrix_pointer2(dx3x2_t **ptr) {
   // CHECK-LABEL: @test_extract_matrix_pointer2(
double test_extract_matrix_pointer2(dx3x2_t **ptr) {
 
   return (*(*(ptr + 4) + 6))[2][1 * 3 - 2];
 }
 
 void insert_extract(dx5x5_t a, fx3x3_t b, unsigned long j, short k) {
   // CHECK-LABEL: @insert_extract(
   // CHECK:         [[K:%.*]] = load i16, i16* %k.addr, align 2
   // CHECK-NEXT:    [[K_EXT:%.*]] = sext i16 [[K]] to i64
-  // CHECK-NEXT:    [[MAT:%.*]] = load <9 x float>, <9 x float>* [[MAT_ADDR:%.*]], align 4
   // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[K_EXT]], 3
   // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], 0
-  // CHECK-NEXT:    [[MATEXT:%.*]] = extractelement <9 x float> [[MAT]], i64 [[IDX]]
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 9
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+  // CHECK-NEXT:    [[MAT:%.*]] = load <9 x float>, <9 x float>* [[MAT_ADDR:%.*]], align 4
+  // CHECK-NEXT:    [[MATEXT:%.*]] = extractelement <9 x float> [[MAT]], i64 [[IDX2]]
   // CHECK-NEXT:    [[J:%.*]] = load i64, i64* %j.addr, align 8
   // CHECK-NEXT:    [[IDX3:%.*]] = mul i64 [[J]], 3
   // CHECK-NEXT:    [[IDX4:%.*]] = add i64 [[IDX3]], 2
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX4]], 9
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[MAT2:%.*]] = load <9 x float>, <9 x float>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:    [[MATINS:%.*]] = insertelement <9 x float> [[MAT2]], float [[MATEXT]], i64 [[IDX4]]
   // CHECK-NEXT:    store <9 x float> [[MATINS]], <9 x float>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:    ret void
 
   b[2][j] = b[0][k];
 }
 
void insert_compound_stmt_field(struct Foo *a, float f, unsigned i, unsigned j) {
   // CHECK-LABEL: define{{.*}} void @insert_compound_stmt_field(%struct.Foo* %a, float %f, i32 %i, i32 %j)
   // CHECK:         [[I:%.*]] = load i32, i32* %i.addr, align 4
   // CHECK-NEXT:    [[I_EXT:%.*]] = zext i32 [[I]] to i64
   // CHECK-NEXT:    [[J:%.*]] = load i32, i32* %j.addr, align 4
   // CHECK-NEXT:    [[J_EXT:%.*]] = zext i32 [[J]] to i64
   // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[J_EXT]], 2
   // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[I_EXT]]
   // CHECK-NEXT:    [[MAT_PTR:%.*]] = bitcast [6 x float]* %mat to <6 x float>*
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 6
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[MAT:%.*]] = load <6 x float>, <6 x float>* [[MAT_PTR]], align 4
   // CHECK-NEXT:    [[EXT:%.*]] = extractelement <6 x float> [[MAT]], i64 [[IDX2]]
   // CHECK-NEXT:    [[SUM:%.*]] = fadd float [[EXT]], {{.*}}
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 6
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[MAT2:%.*]] = load <6 x float>, <6 x float>* [[MAT_PTR]], align 4
   // CHECK-NEXT:    [[INS:%.*]] = insertelement <6 x float> [[MAT2]], float [[SUM]], i64 [[IDX2]]
   // CHECK-NEXT:    store <6 x float> [[INS]], <6 x float>* [[MAT_PTR]], align 4
   // CHECK-NEXT:    ret void
 
   a->mat[i][j] += f;
 }
 
 void matrix_as_idx(ix9x3_t a, int i, int j, dx5x5_t b) {
   // CHECK-LABEL: define{{.*}} void @matrix_as_idx(<27 x i32> %a, i32 %i, i32 %j, <25 x double> %b)
   // CHECK:       [[I1:%.*]] = load i32, i32* %i.addr, align 4
   // CHECK-NEXT:  [[I1_EXT:%.*]] = sext i32 [[I1]] to i64
   // CHECK-NEXT:  [[J1:%.*]] = load i32, i32* %j.addr, align 4
   // CHECK-NEXT:  [[J1_EXT:%.*]] = sext i32 [[J1]] to i64
-  // CHECK-NEXT:  [[A:%.*]] = load <27 x i32>, <27 x i32>* %0, align 4
   // CHECK-NEXT:  [[IDX1_1:%.*]] = mul i64 [[J1_EXT]], 9
   // CHECK-NEXT:  [[IDX1_2:%.*]] = add i64 [[IDX1_1]], [[I1_EXT]]
+  // CHECK-NEXT:  [[CMP:%.*]] = icmp ult i64 [[IDX1_2]], 27
+  // CHECK-NEXT:  call void @llvm.assume(i1 [[CMP]])
+  // CHECK-NEXT:  [[A:%.*]] = load <27 x i32>, <27 x i32>* %0, align 4
   // CHECK-NEXT:  [[MI1:%.*]] = extractelement <27 x i32> [[A]], i64 [[IDX1_2]]
   // CHECK-NEXT:  [[MI1_EXT:%.*]] = sext i32 [[MI1]] to i64
   // CHECK-NEXT:  [[J2:%.*]] = load i32, i32* %j.addr, align 4
   // CHECK-NEXT:  [[J2_EXT:%.*]] = sext i32 [[J2]] to i64
   // CHECK-NEXT:  [[I2:%.*]] = load i32, i32* %i.addr, align 4
   // CHECK-NEXT:  [[I2_EXT:%.*]] = sext i32 [[I2]] to i64
-  // CHECK-NEXT:  [[A2:%.*]] = load <27 x i32>, <27 x i32>* {{.*}}, align 4
   // CHECK-NEXT:  [[IDX2_1:%.*]] = mul i64 [[I2_EXT]], 9
   // CHECK-NEXT:  [[IDX2_2:%.*]] = add i64 [[IDX2_1]], [[J2_EXT]]
+  // CHECK-NEXT:  [[CMP:%.*]] = icmp ult i64 [[IDX2_2]], 27
+  // CHECK-NEXT:  call void @llvm.assume(i1 [[CMP]])
+  // CHECK-NEXT:  [[A2:%.*]] = load <27 x i32>, <27 x i32>* {{.*}}, align 4
   // CHECK-NEXT:  [[MI2:%.*]] = extractelement <27 x i32> [[A2]], i64 [[IDX2_2]]
   // CHECK-NEXT:  [[MI3:%.*]] = add nsw i32 [[MI2]], 2
   // CHECK-NEXT:  [[MI3_EXT:%.*]] = sext i32 [[MI3]] to i64
   // CHECK-NEXT:  [[IDX3_1:%.*]] = mul i64 [[MI3_EXT]], 5
   // CHECK-NEXT:  [[IDX3_2:%.*]] = add i64 [[IDX3_1]], [[MI1_EXT]]
+  // CHECK-NEXT:  [[CMP:%.*]] = icmp ult i64 [[IDX3_2]], 25
+  // CHECK-NEXT:  call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:  [[B:%.*]] = load <25 x double>, <25 x double>* [[B_PTR:%.*]], align 8
   // CHECK-NEXT:  [[INS:%.*]] = insertelement <25 x double> [[B]], double 1.500000e+00, i64 [[IDX3_2]]
   // CHECK-NEXT:  store <25 x double> [[INS]], <25 x double>* [[B_PTR]], align 8
   b[a[i][j]][a[j][i] + 2] = 1.5;
 }

clang/test/CodeGenCXX/matrix-type-operators.cpp

void test_insert_template1(MyMatrix<unsigned, 2, 2> &Mat, unsigned e, unsigned i, unsigned j) {
   // CHECK:         [[E:%.*]] = load i32, i32* %e.addr, align 4
   // CHECK:         [[I:%.*]] = load i32, i32* %i.addr, align 4
   // CHECK-NEXT:    [[I_EXT:%.*]] = zext i32 [[I]] to i64
   // CHECK-NEXT:    [[J:%.*]] = load i32, i32* %j.addr, align 4
   // CHECK-NEXT:    [[J_EXT:%.*]] = zext i32 [[J]] to i64
   // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[J_EXT]], 2
   // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[I_EXT]]
   // CHECK-NEXT:    [[MAT_ADDR:%.*]] = bitcast [4 x i32]* {{.*}} to <4 x i32>*
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 4
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[MAT:%.*]] = load <4 x i32>, <4 x i32>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:    [[MATINS:%.*]] = insertelement <4 x i32> [[MAT]], i32 [[E]], i64 [[IDX2]]
   // CHECK-NEXT:    store <4 x i32> [[MATINS]], <4 x i32>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:    ret void
 
   insert(Mat, e, i, j);
 }
 
 void test_insert_template2(MyMatrix<float, 3, 8> &Mat, float e) {
   // CHECK-LABEL: @_Z21test_insert_template2R8MyMatrixIfLj3ELj8EEf(
   // CHECK:         [[MAT_ADDR:%.*]] = load %struct.MyMatrix.4*, %struct.MyMatrix.4** %Mat.addr, align 8
   // CHECK-NEXT:    [[E:%.*]] = load float, float* %e.addr, align 4
   // CHECK-NEXT:    call void @_Z6insertIfLj3ELj8EEvR8MyMatrixIT_XT0_EXT1_EES1_jj(%struct.MyMatrix.4* nonnull align 4 dereferenceable(96) [[MAT_ADDR]], float [[E]], i32 2, i32 5)
   // CHECK-NEXT:    ret void
   //
   // CHECK-LABEL: define linkonce_odr void @_Z6insertIfLj3ELj8EEvR8MyMatrixIT_XT0_EXT1_EES1_jj(
   // CHECK:         [[E:%.*]] = load float, float* %e.addr, align 4
   // CHECK:         [[I:%.*]] = load i32, i32* %i.addr, align 4
   // CHECK-NEXT:    [[I_EXT:%.*]] = zext i32 [[I]] to i64
   // CHECK-NEXT:    [[J:%.*]] = load i32, i32* %j.addr, align 4
   // CHECK-NEXT:    [[J_EXT:%.*]] = zext i32 [[J]] to i64
   // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[J_EXT]], 3
   // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[I_EXT]]
   // CHECK-NEXT:    [[MAT_ADDR:%.*]] = bitcast [24 x float]* {{.*}} to <24 x float>*
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 24
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[MAT:%.*]] = load <24 x float>, <24 x float>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:    [[MATINS:%.*]] = insertelement <24 x float> [[MAT]], float [[E]], i64 [[IDX2]]
   // CHECK-NEXT:    store <24 x float> [[MATINS]], <24 x float>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:    ret void
 
   insert(Mat, e, 2, 5);
 }
 

 double extract_IntWrapper_idx(double4x4 &m, IntWrapper i, UnsignedWrapper j) {
   // CHECK-LABEL: define{{.*}} double @_Z22extract_IntWrapper_idxRu11matrix_typeILm4ELm4EdE10IntWrapper15UnsignedWrapper(
   // CHECK:         [[I:%.*]] = call i32 @_ZN10IntWrappercviEv(%struct.IntWrapper* {{[^,]*}} %i)
   // CHECK-NEXT:    [[I_ADD:%.*]] = add nsw i32 [[I]], 1
   // CHECK-NEXT:    [[I_ADD_EXT:%.*]] = sext i32 [[I_ADD]] to i64
   // CHECK-NEXT:    [[J:%.*]] = call i32 @_ZN15UnsignedWrappercvjEv(%struct.UnsignedWrapper* {{[^,]*}} %j)
   // CHECK-NEXT:    [[J_SUB:%.*]] = sub i32 [[J]], 1
   // CHECK-NEXT:    [[J_SUB_EXT:%.*]] = zext i32 [[J_SUB]] to i64
+  // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[J_SUB_EXT]], 4
+  // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[I_ADD_EXT]]
+  // CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 16
+  // CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:    [[MAT_ADDR:%.*]] = load [16 x double]*, [16 x double]** %m.addr, align 8
   // CHECK-NEXT:    [[MAT_ADDR2:%.*]] = bitcast [16 x double]* [[MAT_ADDR]] to <16 x double>*
   // CHECK-NEXT:    [[MAT:%.*]] = load <16 x double>, <16 x double>* [[MAT_ADDR2]], align 8
-  // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[J_SUB_EXT]], 4
-  // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[I_ADD_EXT]]
   // CHECK-NEXT:    [[MATEXT:%.*]]  = extractelement <16 x double> [[MAT]], i64 [[IDX2]]
   // CHECK-NEXT:    ret double [[MATEXT]]
   return m[i + 1][j - 1];
 }
 
 template <class T, unsigned R, unsigned C>
 using matrix_type = T __attribute__((matrix_type(R, C)));
 struct identmatrix_t {
void test_constexpr1(matrix_type<float, 4, 4> &m) {
   // CHECK-LABEL: for.body:                                         ; preds = %for.cond
   // CHECK-NEXT:   [[I:%.*]] = load i32, i32* %i, align 4
   // CHECK-NEXT:   [[I_EXT:%.*]] = zext i32 [[I]] to i64
   // CHECK-NEXT:   [[I2:%.*]] = load i32, i32* %i, align 4
   // CHECK-NEXT:   [[I2_EXT:%.*]] = zext i32 [[I2]] to i64
   // CHECK-NEXT:   [[IDX1:%.*]] = mul i64 [[I2_EXT]], 4
   // CHECK-NEXT:   [[IDX2:%.*]] = add i64 [[IDX1]], [[I_EXT]]
   // CHECK-NEXT:   [[MAT_ADDR:%.*]] = bitcast [16 x float]* %result to <16 x float>*
+  // CHECK-NEXT:   [[CMP:%.*]] = icmp ult i64 [[IDX2]], 16
+  // CHECK-NEXT:   call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:   [[MAT:%.*]] = load <16 x float>, <16 x float>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:   [[MATINS:%.*]] = insertelement <16 x float> [[MAT]], float 1.000000e+00, i64 [[IDX2]]
   // CHECK-NEXT:   store <16 x float> [[MATINS]], <16 x float>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:   br label %for.inc
   m = m + identmatrix;
 }
 
 void test_constexpr2(matrix_type<int, 5, 5> &m) {
void test_constexpr2(matrix_type<int, 5, 5> &m) {
   // CHECK-LABEL: for.body:                                         ; preds = %for.cond
   // CHECK-NEXT:   [[I:%.*]] = load i32, i32* %i, align 4
   // CHECK-NEXT:   [[I_EXT:%.*]] = zext i32 [[I]] to i64
   // CHECK-NEXT:   [[I2:%.*]] = load i32, i32* %i, align 4
   // CHECK-NEXT:   [[I2_EXT:%.*]] = zext i32 [[I2]] to i64
   // CHECK-NEXT:   [[IDX1:%.*]] = mul i64 [[I2_EXT]], 5
   // CHECK-NEXT:   [[IDX2:%.*]] = add i64 [[IDX1]], [[I_EXT]]
   // CHECK-NEXT:   [[MAT_ADDR:%.*]] = bitcast [25 x i32]* %result to <25 x i32>*
+  // CHECK-NEXT:   [[CMP:%.*]] = icmp ult i64 [[IDX2]], 25
+  // CHECK-NEXT:   call void @llvm.assume(i1 [[CMP]])
   // CHECK-NEXT:   [[MAT:%.*]] = load <25 x i32>, <25 x i32>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:   [[MATINS:%.*]] = insertelement <25 x i32> [[MAT]], i32 1, i64 [[IDX2]]
   // CHECK-NEXT:   store <25 x i32> [[MATINS]], <25 x i32>* [[MAT_ADDR]], align 4
   // CHECK-NEXT:   br label %for.inc
 
   m = identmatrix - m + 1;
 }

clang/test/CodeGenObjC/matrix-type-operators.m

 // CHECK-NEXT:    [[IV_PTR:%.*]] = bitcast %0* [[IV]] to i8*
 // CHECK-NEXT:    [[CALL:%.*]] = call i32 bitcast (i8* (i8*, i8*, ...)* @objc_msgSend to i32 (i8*, i8*)*)(i8* [[IV_PTR]], i8* [[SEL]])
 // CHECK-NEXT:    [[CONV:%.*]] = sext i32 [[CALL]] to i64
 // CHECK-NEXT:    [[IV2:%.*]] = load %0*, %0** [[IV_ADDR]], align 8
 // CHECK-NEXT:    [[SEL2:%.*]] = load i8*, i8** @OBJC_SELECTOR_REFERENCES_, align 8, !invariant.load !7
 // CHECK-NEXT:    [[IV2_PTR:%.*]] = bitcast %0* [[IV2]] to i8*
 // CHECK-NEXT:    [[CALL1:%.*]] = call i32 bitcast (i8* (i8*, i8*, ...)* @objc_msgSend to i32 (i8*, i8*)*)(i8* [[IV2_PTR]], i8* [[SEL2]])
 // CHECK-NEXT:    [[CONV2:%.*]] = sext i32 [[CALL1]] to i64
-// CHECK-NEXT:    [[MAT:%.*]] = load <16 x double>, <16 x double>* {{.*}} align 8
 // CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[CONV2]], 4
 // CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[CONV]]
+// CHECK-NEXT:   [[CMP:%.*]] = icmp ult i64 [[IDX2]], 16
+// CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+// CHECK-NEXT:    [[MAT:%.*]] = load <16 x double>, <16 x double>* {{.*}} align 8
 // CHECK-NEXT:    [[MATEXT:%.*]] = extractelement <16 x double> [[MAT]], i64 [[IDX2]]
 // CHECK-NEXT:    ret double [[MATEXT]]
 //
 double test_index_placeholders(double4x4 m, IntValue *iv) {
 
   return m[iv.value][iv.value];
 }
 
 __attribute__((objc_root_class))
 @interface MatrixValue
 @property double4x4 value;
 @end
 
 // CHECK-LABEL: @test_base_and_index_placeholders(
 // CHECK:         [[IV:%.*]] = load %0*, %0** [[IV_ADDR:%.*]], align 8
 // CHECK-NEXT:    [[SEL:%.*]] = load i8*, i8** @OBJC_SELECTOR_REFERENCES_, align 8, !invariant.load !7
 // CHECK-NEXT:    [[IV_PTR:%.*]] = bitcast %0* [[IV]] to i8*
 // CHECK-NEXT:    [[CALL:%.*]] = call i32 bitcast (i8* (i8*, i8*, ...)* @objc_msgSend to i32 (i8*, i8*)*)(i8* [[IV_PTR]], i8* [[SEL]])
 // CHECK-NEXT:    [[CONV:%.*]] = sext i32 [[CALL]] to i64
 // CHECK-NEXT:    [[IV2:%.*]] = load %0*, %0** [[IV_ADDR]], align 8
 // CHECK-NEXT:    [[SEL2:%.*]] = load i8*, i8** @OBJC_SELECTOR_REFERENCES_, align 8, !invariant.load !7
 // CHECK-NEXT:    [[IV2_PTR:%.*]] = bitcast %0* [[IV2]] to i8*
 // CHECK-NEXT:    [[CALL1:%.*]] = call i32 bitcast (i8* (i8*, i8*, ...)* @objc_msgSend to i32 (i8*, i8*)*)(i8* [[IV2_PTR]], i8* [[SEL2]])
 // CHECK-NEXT:    [[CONV2:%.*]] = sext i32 [[CALL1]] to i64
+// CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[CONV2]], 4
+// CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[CONV]]
+// CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX2]], 16
+// CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
 // CHECK-NEXT:    [[M:%.*]] = load %1*, %1** %m.addr, align 8
 // CHECK-NEXT:    [[SEL3:%.*]] = load i8*, i8** @OBJC_SELECTOR_REFERENCES_, align 8, !invariant.load !7
 // CHECK-NEXT:    [[M_PTR:%.*]] = bitcast %1* [[M]] to i8*
 // CHECK-NEXT:    [[MAT:%.*]] = call <16 x double> bitcast (i8* (i8*, i8*, ...)* @objc_msgSend to <16 x double> (i8*, i8*)*)(i8* [[M_PTR]], i8* [[SEL3]])
-// CHECK-NEXT:    [[IDX1:%.*]] = mul i64 [[CONV2]], 4
-// CHECK-NEXT:    [[IDX2:%.*]] = add i64 [[IDX1]], [[CONV]]
 // CHECK-NEXT:    [[MATEXT:%.*]] = extractelement <16 x double> [[MAT]], i64 [[IDX2]]
 // CHECK-NEXT:    ret double [[MATEXT]]
 //
 double test_base_and_index_placeholders(MatrixValue *m, IntValue *iv) {
 
   return m.value[iv.value][iv.value];
 }

llvm/include/llvm/IR/MatrixBuilder.h

RHS =
                             RHS, "scalar.splat");
     return cast<VectorType>(LHS->getType())
                    ->getElementType()
                    ->isFloatingPointTy()
                ? B.CreateFDiv(LHS, RHS)
                : (IsUnsigned ? B.CreateUDiv(LHS, RHS) : B.CreateSDiv(LHS, RHS));
   }
 
-  /// Extracts the element at (\p RowIdx, \p ColumnIdx) from \p Matrix.
-  Value *CreateExtractElement(Value *Matrix, Value *RowIdx, Value *ColumnIdx,
-                              unsigned NumRows, Twine const &Name = "") {
+  /// Create an assumption that \p Idx is less than \p NumElements.
+  void CreateIndexAssumption(Value *Idx, unsigned NumElements,
+                             Twine const &Name = "") {
+
+    Value *NumElts =
+        B.getIntN(Idx->getType()->getScalarSizeInBits(), NumElements);
+
+    auto *Cmp = B.CreateICmpULT(Idx, NumElts);
+    if (!isa<llvm::ConstantInt>(Cmp)) {
+      Function *TheFn =
+          Intrinsic::getDeclaration(getModule(), Intrinsic::assume);
+      B.CreateCall(TheFn->getFunctionType(), TheFn, {Cmp}, Name);
+    }
+  }
+
+  /// 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 = "") {
 
     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 *NumRowsV = B.getIntN(MaxWidth, NumRows);
-    return B.CreateExtractElement(
-        Matrix, B.CreateAdd(B.CreateMul(ColumnIdx, NumRowsV), RowIdx),
-        "matext");
+    return B.CreateAdd(B.CreateMul(ColumnIdx, NumRowsV), RowIdx);
   }
 };
 
 } // end namespace llvm
 
 #endif // LLVM_IR_MATRIXBUILDER_H