From d75a0910ec1853c4dbb962270a6c238a35cff32d Mon Sep 17 00:00:00 2001 From: Rahul Joshi Date: Wed, 11 Nov 2020 10:41:08 -0800 Subject: [PATCH] [XLA:GPU] Add test for SliceToDynamic IR emission. - Also fixed IsBlock0Thread0 to split out nested CreateICmpEQ calls out to prevent compiler dependent change in order of the instructions generated. PiperOrigin-RevId: 341860846 Change-Id: Ic65ee73d3b916e925e9745abcfb21186909feed4 --- .../xla/service/gpu/ir_emission_utils.cc | 15 ++-- .../service/gpu/tests/slice_to_dynamic.hlo | 88 +++++++++++++++++++ 2 files changed, 96 insertions(+), 7 deletions(-) create mode 100644 tensorflow/compiler/xla/service/gpu/tests/slice_to_dynamic.hlo diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc index 7743d19497d..53474dcdc66 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc @@ -486,13 +486,14 @@ string CudnnConvKindToString(CudnnConvKind kind) { } llvm::Value* IsBlock0Thread0(llvm::IRBuilder<>* b) { - return b->CreateAnd( - b->CreateICmpEQ( - b->getInt32(0), - EmitCallToTargetIntrinsic(TargetIntrinsicID::kThreadIdx, {}, {}, b)), - b->CreateICmpEQ( - b->getInt32(0), - EmitCallToTargetIntrinsic(TargetIntrinsicID::kBlockIdx, {}, {}, b))); + llvm::Value* is_thread0 = b->CreateICmpEQ( + b->getInt32(0), + EmitCallToTargetIntrinsic(TargetIntrinsicID::kThreadIdx, {}, {}, b)); + + llvm::Value* is_block0 = b->CreateICmpEQ( + b->getInt32(0), + EmitCallToTargetIntrinsic(TargetIntrinsicID::kBlockIdx, {}, {}, b)); + return b->CreateAnd(is_thread0, is_block0); } bool AreFusedReductionOutputsConsistent( diff --git a/tensorflow/compiler/xla/service/gpu/tests/slice_to_dynamic.hlo b/tensorflow/compiler/xla/service/gpu/tests/slice_to_dynamic.hlo new file mode 100644 index 00000000000..1f4f2602094 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/tests/slice_to_dynamic.hlo @@ -0,0 +1,88 @@ +// RUN: hlo_to_llvm_ir %s | FileCheck %s + +// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py + +// CHECK-LABEL: entry: +// CHECK: %[[VAL_0:.*]] = getelementptr inbounds i8, i8* %[[VAL_1:.*]], i64 0 +// CHECK: %[[VAL_2:.*]] = bitcast i8* %[[VAL_0]] to [2 x [2 x [2 x i32]]]* +// CHECK: %[[VAL_3:.*]] = getelementptr inbounds i8, i8* %[[VAL_4:.*]], i64 0 +// CHECK: %[[VAL_5:.*]] = bitcast i8* %[[VAL_3]] to [2 x [2 x [2 x i32]]]* +// CHECK: %[[VAL_6:.*]] = bitcast [2 x [2 x [2 x i32]]]* %[[VAL_2]] to i8* +// CHECK: %[[VAL_7:.*]] = load i32, i32* bitcast ([4 x i8]* @buffer_for_static to i32*), align 4 +// CHECK: %[[VAL_8:.*]] = load i32, i32* bitcast ([4 x i8]* @buffer_for_dynamic to i32*), align 4 +// CHECK: %[[VAL_9:.*]] = load i32, i32* bitcast ([4 x i8]* @buffer_for_static to i32*), align 4 +// CHECK: %[[VAL_10:.*]] = call i32 @llvm.nvvm.read.ptx.sreg.tid.x() +// CHECK: %[[VAL_11:.*]] = icmp eq i32 0, %[[VAL_10]] +// CHECK: %[[VAL_12:.*]] = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() +// CHECK: %[[VAL_13:.*]] = icmp eq i32 0, %[[VAL_12]] +// CHECK: %[[VAL_14:.*]] = and i1 %[[VAL_11]], %[[VAL_13]] +// CHECK: br i1 %[[VAL_14]], label %[[VAL_15:.*]], label %[[VAL_16:.*]] +// CHECK: is_thred_0-after: ; preds = %[[VAL_15]], %[[VAL_17:.*]] +// CHECK: %[[VAL_18:.*]] = mul i32 1, %[[VAL_7]] +// CHECK: %[[VAL_19:.*]] = mul i32 %[[VAL_18]], %[[VAL_8]] +// CHECK: %[[VAL_20:.*]] = mul i32 %[[VAL_19]], %[[VAL_9]] +// CHECK: %[[VAL_21:.*]] = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x(), !range !2 +// CHECK: %[[VAL_22:.*]] = call i32 @llvm.nvvm.read.ptx.sreg.tid.x(), !range !3 +// CHECK: %[[VAL_23:.*]] = mul nuw nsw i32 %[[VAL_21]], 8 +// CHECK: %[[VAL_24:.*]] = add nuw nsw i32 %[[VAL_23]], %[[VAL_22]] +// CHECK: %[[VAL_25:.*]] = icmp ult i32 %[[VAL_24]], 8 +// CHECK: call void @llvm.assume(i1 %[[VAL_25]]) +// CHECK: %[[VAL_26:.*]] = udiv i32 %[[VAL_24]], 1 +// CHECK: %[[VAL_27:.*]] = urem i32 %[[VAL_26]], 2 +// CHECK: %[[VAL_28:.*]] = udiv i32 %[[VAL_24]], 2 +// CHECK: %[[VAL_29:.*]] = urem i32 %[[VAL_28]], 2 +// CHECK: %[[VAL_30:.*]] = udiv i32 %[[VAL_24]], 4 +// CHECK: %[[VAL_31:.*]] = icmp ult i32 %[[VAL_24]], 8 +// CHECK: br i1 %[[VAL_31]], label %[[VAL_32:.*]], label %[[VAL_33:.*]] +// CHECK: custom-call.in_bounds-after: ; preds = %[[VAL_34:.*]], %[[VAL_16]] +// CHECK: ret void +// CHECK: is_thred_0-true: ; preds = %[[VAL_17]] +// CHECK: %[[VAL_35:.*]] = getelementptr inbounds i8, i8* %[[VAL_6]], i32 32 +// CHECK: %[[VAL_36:.*]] = bitcast i8* %[[VAL_35]] to i32* +// CHECK: store i32 %[[VAL_7]], i32* %[[VAL_36]], align 4 +// CHECK: %[[VAL_37:.*]] = getelementptr inbounds i8, i8* %[[VAL_6]], i32 36 +// CHECK: %[[VAL_38:.*]] = bitcast i8* %[[VAL_37]] to i32* +// CHECK: store i32 %[[VAL_8]], i32* %[[VAL_38]], align 4 +// CHECK: %[[VAL_39:.*]] = getelementptr inbounds i8, i8* %[[VAL_6]], i32 40 +// CHECK: %[[VAL_40:.*]] = bitcast i8* %[[VAL_39]] to i32* +// CHECK: store i32 %[[VAL_9]], i32* %[[VAL_40]], align 4 +// CHECK: br label %[[VAL_16]] +// CHECK: custom-call.in_bounds-true: ; preds = %[[VAL_16]] +// CHECK: %[[VAL_41:.*]] = mul nuw nsw i32 %[[VAL_27]], 1 +// CHECK: %[[VAL_42:.*]] = add nuw nsw i32 0, %[[VAL_41]] +// CHECK: %[[VAL_43:.*]] = mul nuw nsw i32 %[[VAL_30]], 2 +// CHECK: %[[VAL_44:.*]] = add nuw nsw i32 %[[VAL_42]], %[[VAL_43]] +// CHECK: %[[VAL_45:.*]] = mul nuw nsw i32 %[[VAL_29]], 4 +// CHECK: %[[VAL_46:.*]] = add nuw nsw i32 %[[VAL_44]], %[[VAL_45]] +// CHECK: %[[VAL_47:.*]] = icmp ult i32 %[[VAL_46]], %[[VAL_20]] +// CHECK: br i1 %[[VAL_47]], label %[[VAL_48:.*]], label %[[VAL_34]] +// CHECK: custom-call.in_dyn_bounds-after: ; preds = %[[VAL_48]], %[[VAL_32]] +// CHECK: br label %[[VAL_33]] +// CHECK: custom-call.in_dyn_bounds-true: ; preds = %[[VAL_32]] +// CHECK: %[[VAL_49:.*]] = udiv i32 %[[VAL_46]], 1 +// CHECK: %[[VAL_50:.*]] = urem i32 %[[VAL_49]], %[[VAL_9]] +// CHECK: %[[VAL_51:.*]] = mul i32 1, %[[VAL_9]] +// CHECK: %[[VAL_52:.*]] = udiv i32 %[[VAL_46]], %[[VAL_51]] +// CHECK: %[[VAL_53:.*]] = urem i32 %[[VAL_52]], %[[VAL_7]] +// CHECK: %[[VAL_54:.*]] = mul i32 %[[VAL_51]], %[[VAL_7]] +// CHECK: %[[VAL_55:.*]] = udiv i32 %[[VAL_46]], %[[VAL_54]] +// CHECK: %[[VAL_56:.*]] = getelementptr inbounds [2 x [2 x [2 x i32]]], [2 x [2 x [2 x i32]]]* %[[VAL_5]], i32 0, i32 %[[VAL_55]], i32 %[[VAL_53]], i32 %[[VAL_50]] +// CHECK: %[[VAL_57:.*]] = load i32, i32* %[[VAL_56]], align 4, !invariant.load !4 +// CHECK: %[[VAL_58:.*]] = bitcast [2 x [2 x [2 x i32]]]* %[[VAL_2]] to i32* +// CHECK: %[[VAL_59:.*]] = getelementptr inbounds i32, i32* %[[VAL_58]], i32 %[[VAL_24]] +// CHECK: store i32 %[[VAL_57]], i32* %[[VAL_59]], align 4 +// CHECK: br label %[[VAL_34]] + +HloModule SliceToDynamic + +ENTRY main { + %param = s32[2,2,2]{2,0,1} parameter(0) + %static = s32[] constant(2) + %dynamic = s32[] constant(1) + ROOT %custom-call = s32[2,<=2, 2]{2,0,1} custom-call(s32[2,2,2]{2,0,1} %param, + s32[] %static, + s32[] %dynamic, + s32[] %static), + custom_call_target="SliceToDynamic", + backend_config="" +}