Resultant [2,2] matrix on multiplication of [2,784] x [784,2]

2x2_matmul.ll

 ; ModuleID = '__compute_module'
source_filename = "__compute_module"
target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-unknown-linux_gnu"

define void @"cluster_2[_XlaCompiledKernel=true,_XlaNumConstantArgs=0].v5"(i8* align 16 dereferenceable(16) %retval, i8* noalias %run_options, i8** noalias %params, i8** noalias %temps, i64* noalias %prof_counters) {
entry:
  %accum_address = alloca float
  %invar_address.reduction = alloca i64
  %invar_address.rhs1 = alloca i64
  %invar_address.lhs0 = alloca i64
  %0 = getelementptr inbounds i8*, i8** %params, i64 1
  %1 = load i8*, i8** %0, !tbaa !0, !dereferenceable !3, !align !4
  %2 = bitcast i8* %1 to [2 x [784 x float]]*
  %3 = getelementptr inbounds i8*, i8** %params, i64 0
  %4 = load i8*, i8** %3, !tbaa !5, !dereferenceable !3, !align !4
  %5 = bitcast i8* %4 to [784 x [2 x float]]*
  %6 = bitcast i8* %retval to [2 x [2 x float]]*
  store i64 0, i64* %invar_address.lhs0
  br label %loop_header.lhs0

loop_header.lhs0:                                 ; preds = %loop_exit.rhs1, %entry
  %indvar.lhs0 = load i64, i64* %invar_address.lhs0
  %7 = icmp uge i64 %indvar.lhs0, 2
  br i1 %7, label %loop_exit.lhs0, label %loop_body.lhs0

loop_body.lhs0:                                   ; preds = %loop_header.lhs0
  store i64 0, i64* %invar_address.rhs1
  br label %loop_header.rhs1

loop_header.rhs1:                                 ; preds = %loop_exit.reduction, %loop_body.lhs0
  %indvar.rhs1 = load i64, i64* %invar_address.rhs1
  %8 = icmp uge i64 %indvar.rhs1, 2
  br i1 %8, label %loop_exit.rhs1, label %loop_body.rhs1

loop_body.rhs1:                                   ; preds = %loop_header.rhs1
  store i64 0, i64* %invar_address.reduction
  store float 0.000000e+00, float* %accum_address
  br label %loop_header.reduction

loop_header.reduction:                            ; preds = %loop_body.reduction, %loop_body.rhs1
  %indvar.reduction = load i64, i64* %invar_address.reduction
  %9 = icmp uge i64 %indvar.reduction, 784
  br i1 %9, label %loop_exit.reduction, label %loop_body.reduction

loop_body.reduction:                              ; preds = %loop_header.reduction
  %10 = getelementptr inbounds [2 x [784 x float]], [2 x [784 x float]]* %2, i64 0, i64 %indvar.lhs0, i64 %indvar.reduction
  %11 = load float, float* %10, !tbaa !7, !invariant.load !9, !noalias !10
  %12 = getelementptr inbounds [784 x [2 x float]], [784 x [2 x float]]* %5, i64 0, i64 %indvar.reduction, i64 %indvar.rhs1
  %13 = load float, float* %12, !tbaa !13, !invariant.load !9, !noalias !10
  %14 = fmul fast float %11, %13
  %15 = load float, float* %accum_address
  %16 = fadd fast float %15, %14
  store float %16, float* %accum_address
  %invar.inc2 = add nuw nsw i64 %indvar.reduction, 1
  store i64 %invar.inc2, i64* %invar_address.reduction
  br label %loop_header.reduction

loop_exit.reduction:                              ; preds = %loop_header.reduction
  %17 = load float, float* %accum_address
  %18 = getelementptr inbounds [2 x [2 x float]], [2 x [2 x float]]* %6, i64 0, i64 %indvar.lhs0, i64 %indvar.rhs1
  store float %17, float* %18, !tbaa !15, !alias.scope !10
  %invar.inc1 = add nuw nsw i64 %indvar.rhs1, 1
  store i64 %invar.inc1, i64* %invar_address.rhs1
  br label %loop_header.rhs1

loop_exit.rhs1:                                   ; preds = %loop_header.rhs1
  %invar.inc = add nuw nsw i64 %indvar.lhs0, 1
  store i64 %invar.inc, i64* %invar_address.lhs0
  br label %loop_header.lhs0

loop_exit.lhs0:                                   ; preds = %loop_header.lhs0
  %prof_counter_computation = getelementptr i64, i64* %prof_counters, i64 0
  ret void
}

!0 = !{!1, !1, i64 0}
!1 = !{!"pointer-to element_type: F32 dimensions: 2 dimensions: 784 layout { minor_to_major: 1 minor_to_major: 0 }", !2}
!2 = !{!"XLA TBAA"}
!3 = !{i64 6272}
!4 = !{i64 16}
!5 = !{!6, !6, i64 0}
!6 = !{!"pointer-to element_type: F32 dimensions: 784 dimensions: 2 layout { minor_to_major: 1 minor_to_major: 0 }", !2}
!7 = !{!8, !8, i64 0}
!8 = !{!"element_type: F32 dimensions: 2 dimensions: 784 layout { minor_to_major: 1 minor_to_major: 0 }", !2}
!9 = !{}
!10 = !{!11}
!11 = !{!"buffer: 2", !12}
!12 = distinct !{!12}
!13 = !{!14, !14, i64 0}
!14 = !{!"element_type: F32 dimensions: 784 dimensions: 2 layout { minor_to_major: 1 minor_to_major: 0 }", !2}
!15 = !{!16, !16, i64 0}
!16 = !{!"element_type: F32 dimensions: 2 dimensions: 2 layout { minor_to_major: 1 minor_to_major: 0 }", !2}