Code

PIN to 2023_11_29#4

Code

Merged

wbmc merged 110 commits intomainintelligent-machine-learning/xla:mainfrom pin_2023_11_29intelligent-machine-learning/xla:pin_2023_11_29Copy head branch name to clipboardDec 13, 2023

+11,639-10,244

Commits

Commits on Nov 26, 2023

• [XLA:GPU] NFC: use matchers in PriorityFusionTest.

  Show description for aa20614
  

  chsigg

  authored and

  copybara-github

  committed
  aa20614

  Copy full SHA for aa20614

• [XLA:GPU] NFC: make IsReadCoalesced() slightly easier to read.

  Show description for 048e5e4
  

  chsigg

  authored and

  copybara-github

  committed
  048e5e4

  Copy full SHA for 048e5e4

Commits on Nov 27, 2023

• Make cudnn_fused_conv_rewriter_test work in OSS.

  Show description for 0154c7b
  

  akuegel

  authored and

  copybara-github

  committed
  0154c7b

  Copy full SHA for 0154c7b

• NFC: Reduce number of function arguments in reduction.cc.

  Show description for f4d1db4
  

  jreiffers

  authored and

  copybara-github

  committed
  f4d1db4

  Copy full SHA for f4d1db4

• Enable tests that require v100 in OSS.

  Show description for 3249850
  

  akuegel

  authored and

  copybara-github

  committed
  3249850

  Copy full SHA for 3249850

• Move ReductionCodegenState to reduction.cc and rename to ReductionGroupEmitter.

  Show description for 77f55b9
  

  jreiffers

  authored and

  copybara-github

  committed
  77f55b9

  Copy full SHA for 77f55b9

• [XLA:GPU] Add can_fuse cache.

  Show description for 12adbd1
  

  olegshyshkov

  authored and

  copybara-github

  committed
  12adbd1

  Copy full SHA for 12adbd1

• Integrate LLVM at llvm/llvm-project@5e5a22caf88a

  Show description for 3ccca51
  

  gribozavr

  authored and

  copybara-github

  committed
  3ccca51

  Copy full SHA for 3ccca51

• [XLA:GPU] (NFC) Remove producer_user_count_.

  Show description for 59b1d67
  

  olegshyshkov

  authored and

  copybara-github

  committed
  59b1d67

  Copy full SHA for 59b1d67

• [XLA:GPU][NFC] Clean up tile analysis tests to use TF_ASSERT_OK_AND_ASSIGN.

  Show description for bb07b21
  

  bchetioui

  authored and

  copybara-github

  committed
  bb07b21

  Copy full SHA for bb07b21

• Update TFRT dependency to use revision

  Show description for ab750b1
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  ab750b1

  Copy full SHA for ab750b1

• [XLA:GPU][NFC] Transform dimension propagation to the functional paradigm

  Show description for eb3f22f
  

  tdanyluk

  authored and

  copybara-github

  committed
  eb3f22f

  Copy full SHA for eb3f22f

• [XLA:GPU] Fix for "Transform dimension propagation to the functional paradigm

  Show description for 2e383a1
  

  tdanyluk

  authored and

  copybara-github

  committed
  2e383a1

  Copy full SHA for 2e383a1

• [XLA] Collective pipeliner goes through small reduces as they do not increase memory pressure much.

  Show description for 7506f69
  

  Marcello Maggioni

  authored and

  copybara-github

  committed
  7506f69

  Copy full SHA for 7506f69

• Adds back the solver parameter string.

  Show description for 5320016
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  5320016

  Copy full SHA for 5320016

• [stream_executor] Add Case conditional command to CommandBuffer openxla#6973

  Show description for c9e2046
  

  ezhulenev

  authored and

  copybara-github

  committed
  c9e2046

  Copy full SHA for c9e2046

• [stream_executor] Add For conditional command to CommandBuffer

  Show description for b71268c
  

  ezhulenev

  authored and

  copybara-github

  committed
  b71268c

  Copy full SHA for b71268c

• [stream_executor] NFC: Do not leak internal stream executor header

  Show description for dee190a
  

  ezhulenev

  authored and

  copybara-github

  committed
  dee190a

  Copy full SHA for dee190a

• [XLA:GPU] Triton GEMM: enable fusion of inputs concatenated along non-contracting dimensions.

  Show description for 2c5ce09
  

  Ilia Sergachev

  authored and

  copybara-github

  committed
  2c5ce09

  Copy full SHA for 2c5ce09

• Deduplicate some code by using a refactored portion of MaybeFollowInsStrategyGroup to generate strategies for elementwise ops as well.

  Show description for fea3aa5
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  fea3aa5

  Copy full SHA for fea3aa5

• Search for possible predecessors instead of only caching one hops.

  Show description for 48d695d
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  48d695d

  Copy full SHA for 48d695d

• PR openxla#7269: Fix an incorrect static_assert

  Show description for ea34cf8
  

  elfringham

  authored and

  copybara-github

  committed
  ea34cf8

  Copy full SHA for ea34cf8

• PR openxla#7277: [ROCM] rocm 6.0 fixes followup

  Show description for e02661f
  

  pemeliya

  authored and

  copybara-github

  committed
  e02661f

  Copy full SHA for e02661f

Commits on Nov 28, 2023

• Add c-api support for GetCompiledMemoryStats.

  Show description for 7c23e8c
  

  pschuh

  authored and

  copybara-github

  committed
  7c23e8c

  Copy full SHA for 7c23e8c

• Update visibility for tf.data.

  Show description for ecae151
  

  yangustc07

  authored and

  copybara-github

  committed
  ecae151

  Copy full SHA for ecae151

• [xla:gpu] NFC: move thunk->CMD conversion to runtime3 openxla#6528

  Show description for 889fa5c
  

  anlunx

  authored and

  copybara-github

  committed
  889fa5c

  Copy full SHA for 889fa5c

• [stream_executor] NFC: Do not leak internal stream executor header

  Show description for 6d09f41
  

  ezhulenev

  authored and

  copybara-github

  committed
  6d09f41

  Copy full SHA for 6d09f41

• [stream_executor] Do not use KernelArgs directly in GpuExecutor::Launch

  Show description for 6643428
  

  tyb0807

  authored and

  copybara-github

  committed
  6643428

  Copy full SHA for 6643428

• [stream_executor] Add While conditional command to CommandBuffer openxla#6973

  Show description for e795dc3
  

  ezhulenev

  authored and

  copybara-github

  committed
  e795dc3

  Copy full SHA for e795dc3

• [stream_executor] NFC: Do not leak internal stream executor header

  Show description for 537ef5e
  

  ezhulenev

  authored and

  copybara-github

  committed
  537ef5e

  Copy full SHA for 537ef5e

• Adds tests for AbsOp, DivOp, ExpOp, MaxOp, MulOp, PowOp, SubOp in xla_builder_test and shape_inference_test. The tests test shape inference for unbounded dynamism.

  Show description for 2521abf
  

  yishuangP

  authored and

  copybara-github

  committed
  2521abf

  Copy full SHA for 2521abf

• [stream_executor] NFC: Restore private visibility of stream executor private targets

  Show description for 33fc605
  

  ezhulenev

  authored and

  copybara-github

  committed
  33fc605

  Copy full SHA for 33fc605

• [XLA/GPU]Add methods to construct GpuCompilationEnvironment.

  Show description for 8e3830c
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  8e3830c

  Copy full SHA for 8e3830c

• [XLA:GPU] Remove xla_gpu_single_wave_autotuning flag

  Show description for 9ea69b0
  

  tdanyluk

  authored and

  copybara-github

  committed
  9ea69b0

  Copy full SHA for 9ea69b0

• Add a vector->SCF pass to hlo_xla_runtime_pipeline.

  Show description for 370895f
  

  pifon2a

  authored and

  copybara-github

  committed
  370895f

  Copy full SHA for 370895f

• [XLA:GPU] Add cache of runtime data for unfused instruction.

  Show description for 80fe7b1
  

  olegshyshkov

  authored and

  copybara-github

  committed
  80fe7b1

  Copy full SHA for 80fe7b1

• Retire deallocation dialect.

  Show description for d44ee9f
  

  jreiffers

  authored and

  copybara-github

  committed
  d44ee9f

  Copy full SHA for d44ee9f

• PR openxla#7287: Fix TF32 in gemm calls

  Show description for 9afe429
  

  benbarsdell

  authored and

  copybara-github

  committed
  9afe429

  Copy full SHA for 9afe429

• [XLA:GPU] Add cache for fusion runtime.

  Show description for 3ee4ca8
  

  olegshyshkov

  authored and

  copybara-github

  committed
  3ee4ca8

  Copy full SHA for 3ee4ca8

• [XLA:GPU] Fix crash in (Triton) tiling of reductions.

  Show description for a51b04c
  

  jvstokes

  authored and

  copybara-github

  committed
  a51b04c

  Copy full SHA for a51b04c

• PR openxla#7323: [ROCm] Restore private visibility of stream executor private targets

  Show description for d0aed34
  

  i-chaochen

  authored and

  copybara-github

  committed
  d0aed34

  Copy full SHA for d0aed34

• [XLA:CPU] Refactor local collectives into a separate file behind an interface.

  Show description for 95ddb7a
  

  hawkinsp

  authored and

  copybara-github

  committed
  95ddb7a

  Copy full SHA for 95ddb7a

• [XLA:GPU][NFC] Split GemmRewriterTriton into 4 parts

  Show description for bff1238
  

  tdanyluk

  authored and

  copybara-github

  committed
  bff1238

  Copy full SHA for bff1238

• Failed compatibility tests were disabled for older version plugins and we can remove the compatibility support here.

  Show description for 3032620
  

  Jieying Luo

  authored and

  copybara-github

  committed
  3032620

  Copy full SHA for 3032620

• third_party/gpus changes from PR openxla#7277 that were missed

  Show description for 0ca186f
  

  ddunl

  authored and

  copybara-github

  committed
  0ca186f

  Copy full SHA for 0ca186f

Commits on Nov 29, 2023

• [xla:runtime] NFC: Remove runner library

  Show description for 7a7eddc
  

  ezhulenev

  authored and

  copybara-github

  committed
  7a7eddc

  Copy full SHA for 7a7eddc

• PR openxla#7136: [XLA:GPU] Add Allocate command to command buffer

  Show description for 1347b18
  

  shawnwang18

  authored and

  copybara-github

  committed
  1347b18

  Copy full SHA for 1347b18

• Update XLA GPU config with NVCC compiler.

  Show description for 6046ab5
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  6046ab5

  Copy full SHA for 6046ab5

• For convolutions that can be interpreted as dots, rely on DotHandler to generate sharding strategies. For those that cannot be, we rely on pre-existing convolution handling code.

  Show description for 6b4d712
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  6b4d712

  Copy full SHA for 6b4d712

• [stream_executor] Record cond_builder before evaluating loop condition in a While command

  Show description for 85a4f70
  

  ezhulenev

  authored and

  copybara-github

  committed
  85a4f70

  Copy full SHA for 85a4f70

• [stream_executor] Add Memset command to CommandBuffer

  Show description for c58bc73
  

  ezhulenev

  authored and

  copybara-github

  committed
  c58bc73

  Copy full SHA for c58bc73

• [stream_executor] Add support for updating Memcpy command parameters

  Show description for fede5db
  

  ezhulenev

  authored and

  copybara-github

  committed
  fede5db

  Copy full SHA for fede5db

• Converts the AutoShardingSolverRequest object into a proto (which provides built-in utilities for saving & loading).

  Show description for 6c4ff70
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  6c4ff70

  Copy full SHA for 6c4ff70

• Integrate StableHLO at openxla/stablehlo@83f095e7

  Show description for 8d49e11
  

  ghpvnist

  authored and

  copybara-github

  committed
  8d49e11

  Copy full SHA for 8d49e11

• [stream_executor] NFC: Guard new features with CUDA_VERSION check

  Show description for 12a88f3
  

  ezhulenev

  authored and

  copybara-github

  committed
  12a88f3

  Copy full SHA for 12a88f3

• [XLA] Add support for while loop simplifier to remove duplicated "dynamic-update-slice" stored values.

  Show description for efcf825
  

  Marcello Maggioni

  authored and

  copybara-github

  committed
  efcf825

  Copy full SHA for efcf825

• [XLA:GPU] Add a pass merging producer fusions in to Triton Softmax fusions.

  Show description for 1317a62
  

  jvstokes

  authored and

  copybara-github

  committed
  1317a62

  Copy full SHA for 1317a62

• [XLA] Fix the MHLO TopK assemblyFormat to correctly support the optional largest attribute.

  Show description for 5cd07b6
  

  dimitar-asenov

  authored and

  copybara-github

  committed
  5cd07b6

  Copy full SHA for 5cd07b6

• [XLA:GPU][NFC] Beautify code related to Triton fusions

  Show description for e9500e3
  

  tdanyluk

  authored and

  copybara-github

  committed
  e9500e3

  Copy full SHA for e9500e3

• Fix a compile error caused by in_process_collectives

  Show description for ec0b228
  

  apaszke

  authored and

  copybara-github

  committed
  ec0b228

  Copy full SHA for ec0b228

• [TileAnalysis] Add indexing computation for reshape.

  Show description for 422b8e7
  

  pifon2a

  authored and

  copybara-github

  committed
  422b8e7

  Copy full SHA for 422b8e7

• Basic simplifier for indexing maps.

  Show description for b98258d
  

  jreiffers

  authored and

  copybara-github

  committed
  b98258d

  Copy full SHA for b98258d

• [TileAnalysis] Add indexing computation for HloBitcast.

  Show description for c9c16d6
  

  pifon2a

  authored and

  copybara-github

  committed
  c9c16d6

  Copy full SHA for c9c16d6

• [XLA:GPU] Append ".0" suffix to all instructions names.

  Show description for 757f635
  

  olegshyshkov

  authored and

  copybara-github

  committed
  757f635

  Copy full SHA for 757f635

• [xla:gpu] NFC: Migrate custom CUTLASS kernel to GemmUniversal templates

  Show description for fecd5e7
  

  ezhulenev

  authored and

  copybara-github

  committed
  fecd5e7

  Copy full SHA for fecd5e7

• [XLA:GPU][NFC] Refactor tiling propagation.

  Show description for d555040
  

  Ilia Sergachev

  authored and

  copybara-github

  committed
  d555040

  Copy full SHA for d555040

• [XLA] Lower verbosity in MSA.

  Show description for 4e6e1c5
  

  berkinilbeyi

  authored and

  copybara-github

  committed
  4e6e1c5

  Copy full SHA for 4e6e1c5

• Record free_gpu_system_memory using addressable_devices() to avoid potential OOM issues across multiple processes.

  Show description for 0e751bf
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  0e751bf

  Copy full SHA for 0e751bf

• Integrate LLVM at llvm/llvm-project@3287ae8f6520

  Show description for 7eacd1b
  

  gribozavr

  authored and

  copybara-github

  committed
  7eacd1b

  Copy full SHA for 7eacd1b

• [XLA:CPU] Add a direct implementation of AllGather, rather than lowering AllGather to AllReduce.

  Show description for c4f0a9d
  

  hawkinsp

  authored and

  copybara-github

  committed
  c4f0a9d

  Copy full SHA for c4f0a9d

• User sharding annotations can sometimes be invalid wrt to the shape of the instruction. When this happens, and a consumer instruction generates sharding startegies for itself based on the strategie…

  Show description for 8e8b121
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  8e8b121

  Copy full SHA for 8e8b121

• Some cleanup: Inline a function used once, remove some dead code, and compactify some code.

  Show description for 856d51c
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  856d51c

  Copy full SHA for 856d51c

• Reenable BitcastConvert tests in JAX CI now that the tests pass again

  Show description for c28baac
  

  ddunl

  authored and

  copybara-github

  committed
  c28baac

  Copy full SHA for c28baac

• Update TFRT dependency to use revision

  Show description for 222c77a
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  222c77a

  Copy full SHA for 222c77a

• Adds the ability to dump solver request protos.

  Show description for 89f3b13
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  89f3b13

  Copy full SHA for 89f3b13

• Fastpath for setting disable-jit.

  Show description for 57edcfa
  

  pschuh

  authored and

  copybara-github

  committed
  57edcfa

  Copy full SHA for 57edcfa

• When iterative solving is turned on, shardings annotations from a previous iteration associated with a partial mesh conflict with the current larger mesh. This CL handles this case when calling sha…

  Show description for c756093
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  c756093

  Copy full SHA for c756093

• Move large functions in auto_sharding_dot_handler.cc out of class definitions

  Show description for 17d9572
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  17d9572

  Copy full SHA for 17d9572

• [XLA]

  Show description for 9e172db
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  9e172db

  Copy full SHA for 9e172db

• Changed Version of Bazel to version 6.4.0

  Show description for 954527c
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  954527c

  Copy full SHA for 954527c

• [XLA:CPU] Add a direct implementation of ReduceScatter, instead of lowering ReduceScatter to AllReduce+DynamicSlice.

  Show description for 69f26cf
  

  hawkinsp

  authored and

  copybara-github

  committed
  69f26cf

  Copy full SHA for 69f26cf

• PR openxla#7370: [XLA:GPU] fix command_buffer_thunk_test failure

  Show description for aa1f21c
  

  shawnwang18

  authored and

  copybara-github

  committed
  aa1f21c

  Copy full SHA for aa1f21c

• [xla:ffi] Added error reporting to existing decoders

  Show description for ae7fead
  

  superbobry

  authored and

  copybara-github

  committed
  ae7fead

  Copy full SHA for ae7fead

• Reverts 69f26cf

  Show description for f69ec88
  

  hawkinsp

  authored and

  copybara-github

  committed
  f69ec88

  Copy full SHA for f69ec88

• Remove obselete TODO

  Show description for 3400945
  

  ddunl

  authored and

  copybara-github

  committed
  3400945

  Copy full SHA for 3400945

• Reverts c4f0a9d

  Show description for baf6958
  

  hawkinsp

  authored and

  copybara-github

  committed
  baf6958

  Copy full SHA for baf6958

• Remove tensorflow namespace from tsl/platform/status_matchers.h

  Show description for a680ac3
  

  ddunl

  authored and

  copybara-github

  committed
  a680ac3

  Copy full SHA for a680ac3

• [XLA:LatencyHidingScheduler] Prefer picking the instruction with less serial-resource conflicts.

  Show description for 3981d17
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  3981d17

  Copy full SHA for 3981d17

• [xla:ffi] Parameterized ffi::BaseBuffer with its rank

  Show description for b328fe8
  

  superbobry

  authored and

  copybara-github

  committed
  b328fe8

  Copy full SHA for b328fe8

• [XLA] Allow moving for HloSharding in xla::HloInstruction::set_sharding.

  Show description for 4bf4b1e
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  4bf4b1e

  Copy full SHA for 4bf4b1e

• Add EvalOrPattern to StablehloRefineShapes pass

  Show description for 6449702
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  6449702

  Copy full SHA for 6449702

• [XLA] Allow tuple_shardings move in the HloSharding ctor.

  Show description for aab9273
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  aab9273

  Copy full SHA for aab9273

Commits on Nov 30, 2023

• [xla:gpu] Move external allocation implementation details from StreamExecutor to XLA:GPU level

  Show description for 135b0ea
  

  ezhulenev

  authored and

  copybara-github

  committed
  135b0ea

  Copy full SHA for 135b0ea

• [XLA] Do not create temporary vector in HloSharding::GetSubSharding.

  Show description for 3521147
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  3521147

  Copy full SHA for 3521147

• [XLA] Remove an unnecessary HloSharding assignment.

  Show description for d6a1cf3
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  d6a1cf3

  Copy full SHA for d6a1cf3

• Use clang+NVCC compilers for all XLA GPU jobs.

  Show description for 4695163
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  4695163

  Copy full SHA for 4695163

• [XLA] Move HloSharding where possible.

  Show description for 3f03d1c
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  3f03d1c

  Copy full SHA for 3f03d1c

• [XLA] Change the prototype of ReturnImprovedShardingImpl to minmize copies.

  Show description for c5d313d
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  c5d313d

  Copy full SHA for c5d313d

• [XLA] Optimize HloValue::ComputeUses().

  Show description for ecfe268
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  ecfe268

  Copy full SHA for ecfe268

• [xla:gpu] Own command buffer allocations at a Thunk level

  Show description for da00f2b
  

  ezhulenev

  authored and

  copybara-github

  committed
  da00f2b

  Copy full SHA for da00f2b

• Integrate LLVM at llvm/llvm-project@f688e0901213

  Show description for e86ce2e
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  e86ce2e

  Copy full SHA for e86ce2e

• [stream_executor] Use new CUDA runtime API for TopK

  Show description for c6f0a6e
  

  tyb0807

  authored and

  copybara-github

  committed
  c6f0a6e

  Copy full SHA for c6f0a6e

• [xla:gpu] Add the AOT compilation pipeline for thunk runtime openxla#7360

  Show description for a36626c
  

  anlunx

  authored and

  copybara-github

  committed
  a36626c

  Copy full SHA for a36626c

• Update TFRT dependency to use revision

  Show description for 46d4355
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  46d4355

  Copy full SHA for 46d4355

• [XLA:GPU] Improve errors if callbacks are not provided to the CollectivePipeliner.

  Show description for 6eed44a
  

  hawkinsp

  authored and

  copybara-github

  committed
  6eed44a

  Copy full SHA for 6eed44a

• Relocates the evaluation output into the core solver.

  Show description for dbe129e
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  dbe129e

  Copy full SHA for dbe129e

• Lower tensor.from_elements and shape.broadcast ops in ShapeLegalizeToHLO

  Show description for 59dea45
  

  tensorflower-gardener

  authored and

  copybara-github

  committed
  59dea45

  Copy full SHA for 59dea45

• Internal infrastructure change

  Show description for 91b8ddb
  

  GleasonK

  authored and

  copybara-github

  committed
  91b8ddb

  Copy full SHA for 91b8ddb

• [stream_executor][NFC] More doc for While and For command

  Show description for 58e6b42
  

  tyb0807

  authored and

  copybara-github

  committed
  58e6b42

  Copy full SHA for 58e6b42

Commits on Dec 13, 2023

• Merge commit '58e6b428e22e40c4100a7b66790fbe86dc9d7845' into HEAD

  

  mochen.bmc

  committed
  9e419d5

  Copy full SHA for 9e419d5