[Large Tensor] Fixed Spatial Transformer op

[connorgoggins]

Description

The Spatial Transformer op was previously breaking on large tensor (dimension >= 2^32) data. With the following input:

run_performance_test(nd.SpatialTransformer, run_backward=True, inputs=[{'data': (2, 2**29,1,6), 'loc': nd.random_normal(shape=(2,6)), 'transform_type': 'affine', 'sampler_type': 'bilinear', 'target_shape': (2,6)}], warmup=1, runs=1)

the following error was thrown:

*** Error in `python3': double free or corruption (out): 0x00007f36bbffe010 ***

To root cause this issue, I ran the previous command in a Python script with GDB, and found that the underlying problem was in the iteration portion of the forward and backward methods of spatial_transformer.cc. Several of the variables used in the iteration used the int dtype when they should have been using index_t to properly handle long int indices. I switched these variables to index_t in the forward and backward methods, and after rebuilding, the previous input command displayed the correct output:

INFO:root:Begin Benchmark - SpatialTransformer
INFO:root:Complete Benchmark - SpatialTransformer
[{'SpatialTransformer': [{'inputs': {'data': (2, 536870912, 1, 6), 'loc': '<NDArray 2x6 @cpu(0)>', 'transform_type': 'affine', 'sampler_type': 'bilinear', 't
arget_shape': (2, 6)}, 'max_storage_mem_alloc_cpu/0': 102005472.0, 'avg_time_forward_SpatialTransformer': 551614.125, 'avg_time_backward_SpatialTransformer':
 511034.0938}]}]

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• M src/operator/tensor/spatial_transformer.cc

Comments

Tested on r5dn.24xl-ubuntu 16.04 and p2.16xl-ubuntu 16.04 with

1. Individual op run
2. Full OpPerf run

Results

The key difference between CPU and GPU tests was the instance type (r5dn.24xl for CPU, p2.16xl for GPU). All relevant build flags remain the same, and both were tested using CPU context.

Single operator test - SpatialTransformer op (GPU)
Single operator test - SpatialTransformer op (CPU)

Full OpPerf test (GPU)
Full OpPerf test (CPU)

@apeforest @access2rohit @ChaiBapchya

[connorgoggins]

@mxnet-label-bot add [pr-awaiting-review]

[access2rohit]

@connorgoggins can you tell me what are the possible types for DType here:
https://github.com/apache/incubator-mxnet/pull/17617/files#diff-64979618e95b2211ff9317b0830fbbd2R37

[connorgoggins]

@access2rohit DType here represents the type of the interior elements of the input tensors, and since SpatialTransformer only supports floating point types then DType could be float32 or float64 (also potentially float16, but this DType doesn't seem to have been implemented for GEMM on CPU yet).

[ChaiBapchya]

why is this changed back to Dtype?
It's the index position right?

[connorgoggins]

@ChaiBapchya changing DType to index_t for these specific variables makes the op generate incorrect output on standard inputs (e.g. the inputs in the CI run) - the values generated in the output NDArray are all integers instead of floats. This is due to the fact that these variables do not represent the index positions (as I also originally believed), but instead represent the underlying values at the vertices.

[ChaiBapchya]

Cool. good catch.
_v indicates the value at that particular index. Thanks!

[ChaiBapchya]

LGTM!